"""Budget Studio: AI-assisted fund allocation proposals.

Builds a grounded prompt from a fund's real budget-item catalog and existing
allocations, and parses the model's JSON reply into allocation row dicts the
frontend can preview/edit before saving.
"""
import json
from decimal import Decimal, InvalidOperation

ALLOCATE_SYSTEM_PROMPT = (
    "You are a budget allocation assistant for a research organization's grant "
    "management system. Given a funding amount and a request describing how to "
    "split it, propose a breakdown across the organization's real budget items.\n\n"
    "Rules:\n"
    "- Only use budget_item_id values from the BUDGET ITEM CATALOG provided. "
    "Never invent an id, code, or item.\n"
    "- The sum of all proposed amounts should not exceed the fund's allocable "
    "amount (total fund amount minus already-allocated amount), unless the "
    "user's request explicitly asks to exceed it.\n"
    "- Each row's amount is computed as quantity * unit_cost * frequency * "
    "(time_allocated_pct / 100). Choose quantity, unit_cost, frequency, and "
    "time_allocated_pct so that this product matches the intended amount for "
    "that line — the simplest choice is quantity=1, frequency=1, "
    "time_allocated_pct=100, and unit_cost equal to the intended amount.\n"
    "- Respond with ONLY a JSON object: "
    '{"allocations": [{"budget_item_id": str, "quantity": number, '
    '"unit_cost": number, "frequency": number, "time_allocated_pct": number, '
    '"unit": str, "note": str}, ...]}\n'
    "- Do not allocate to a budget item that already has an allocation on this "
    "fund unless the user asks to adjust it — in that case set note to explain "
    "the change.\n"
    "- Keep `note` short (one sentence) explaining the reasoning for that line."
)


def build_allocation_prompt(fund, prompt):
    """Return (system, user_content) strings grounding the model on real data."""
    from apps.finance.models import BudgetItem

    allocated = sum((a.amount or Decimal("0")) for a in fund.allocations.all())
    allocable = (fund.amount or Decimal("0")) - allocated

    existing = fund.allocations.select_related("budget_item", "budget_item__category").all()
    existing_lines = [
        f"- {a.budget_item.code} {a.budget_item.name}: {fund.currency} {a.amount}"
        for a in existing
    ]

    items = BudgetItem.objects.filter(category__organization=fund.organization).select_related("category").order_by("category__name", "code")[:500]
    catalog_lines = [
        f"- id={it.id} code={it.code} name=\"{it.name}\" category=\"{it.category.name if it.category_id else ''}\""
        for it in items
    ]

    facts = [
        f"Fund source: {fund.source}",
        f"Fund total amount: {fund.currency} {fund.amount}",
        f"Already allocated: {fund.currency} {allocated}",
        f"Allocable (remaining) amount: {fund.currency} {allocable}",
    ]

    user_content = (
        "FUND:\n" + "\n".join(facts) + "\n\n"
        + "EXISTING ALLOCATIONS ON THIS FUND:\n"
        + ("\n".join(existing_lines) if existing_lines else "(none yet)") + "\n\n"
        + "BUDGET ITEM CATALOG (only use these ids):\n"
        + "\n".join(catalog_lines) + "\n\n"
        + "USER REQUEST:\n" + prompt
    )
    return ALLOCATE_SYSTEM_PROMPT, user_content


def parse_allocation_response(raw, fund):
    """Parse the model's JSON reply into a list of allocation row dicts.

    Validates budget_item_id against the real catalog and computes `amount`
    the same way ProjectFundAllocation.save() does, so the preview matches
    what will actually be stored. Raises ValueError on malformed output.
    """
    from apps.finance.models import BudgetItem

    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        raise ValueError("AI response was not valid JSON.")

    rows = data.get("allocations") if isinstance(data, dict) else None
    if not isinstance(rows, list):
        raise ValueError("AI response did not contain an allocations array.")

    item_ids = {str(it.id): it for it in BudgetItem.objects.select_related("category").filter(
        category__organization=fund.organization,
        id__in=[r.get("budget_item_id") for r in rows if isinstance(r, dict) and r.get("budget_item_id")]
    )}
    already_allocated = set(fund.allocations.values_list("budget_item_id", flat=True))

    def _num(v, default="0"):
        try:
            return Decimal(str(v))
        except (InvalidOperation, TypeError):
            return Decimal(default)

    # Keyed by budget_item id: the DB unique constraint is (fund, budget_item),
    # so both a repeated id within the AI's own response and an id that
    # already has an allocation on this fund would 400 on save. Skip both —
    # last write wins for in-response repeats.
    by_item = {}
    for r in rows:
        if not isinstance(r, dict):
            continue
        item = item_ids.get(str(r.get("budget_item_id")))
        if not item or item.id in already_allocated:
            continue  # hallucinated id, or a budget item already allocated on this fund

        quantity = _num(r.get("quantity"), "1")
        unit_cost = _num(r.get("unit_cost"), "0")
        frequency = _num(r.get("frequency"), "1")
        time_pct = _num(r.get("time_allocated_pct"), "100")
        amount = quantity * unit_cost * frequency * (time_pct / Decimal("100"))

        by_item[item.id] = {
            "budget_item": str(item.id),
            "budget_item_code": item.code,
            "budget_item_name": item.name,
            "category_name": item.category.name if item.category_id else None,
            "unit": str(r.get("unit") or "")[:50],
            "quantity": str(quantity),
            "unit_cost": str(unit_cost),
            "unit_cost_currency": fund.currency,
            "frequency": str(frequency),
            "time_allocated_pct": str(time_pct),
            "amount": str(amount),
            "note": str(r.get("note") or "")[:255],
        }

    results = list(by_item.values())
    if not results:
        raise ValueError(
            "AI did not propose any valid allocations — the proposed budget items "
            "may already be allocated on this fund."
        )
    return results