"""The retail analyst assistant: a tool loop (lesson 07) that ends in a structured answer (lesson 06)."""

import json
import logging
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from pydantic import BaseModel, Field

from .db import QueryRejected, ReadOnlyDB

log = logging.getLogger("analyst")
MODEL = "gpt-5.6-luna"                     # keep in config; check the models page (lesson 04)
PRICES = {"gpt-5.6-luna": (0.20, 1.20)}    # USD per 1M tokens (input, output) when written

INSTRUCTIONS = """You are a careful retail data analyst for an Indian e-commerce shop.
Answer questions about the shop's data using the tools. Rules:
- Use describe_tables first if you are unsure of table or column names.
- Write SQLite SELECT queries. Aggregate in SQL; never ask for more rows than needed.
- Revenue = SUM(quantity * unit_price * (1 - COALESCE(discount_pct, 0) / 100.0)) from order_items,
  counting only orders with status 'COMPLETED' unless the user says otherwise.
- The data has known quality problems (duplicate customers, zero or negative quantities, guest orders).
  Mention any that affect the answer in `caveats`.
- Every number in the answer must come from a query result."""


class KeyNumber(BaseModel):
    label: str
    value: str


class Answer(BaseModel):
    summary: str = Field(description="2-4 sentences answering the question")
    key_numbers: list[KeyNumber]
    sql: list[str] = Field(description="The SQL queries whose results support the answer")
    caveats: list[str]


TOOL_SCHEMAS = [
    {"type": "function", "name": "describe_tables",
     "description": "Return the CREATE TABLE statements for every table in the shop database.",
     "parameters": {"type": "object", "properties": {}, "required": [], "additionalProperties": False},
     "strict": True},
    {"type": "function", "name": "run_sql",
     "description": "Run ONE read-only SQLite SELECT query and return at most 50 rows as a table.",
     "parameters": {"type": "object", "properties": {"query": {"type": "string"}},
                    "required": ["query"], "additionalProperties": False},
     "strict": True},
    {"type": "function", "name": "search_products",
     "description": "Find products by meaning (e.g. 'something to sit on while working'). "
                    "Returns product ids, names and categories.",
     "parameters": {"type": "object",
                    "properties": {"text": {"type": "string"}, "k": {"type": "integer"}},
                    "required": ["text", "k"], "additionalProperties": False},
     "strict": True},
]


@dataclass
class Turn:
    """What happened while answering one question: tools used, tokens, cost."""
    question: str
    tool_calls: list[tuple[str, dict]] = field(default_factory=list)
    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def cost_usd(self) -> float:
        p_in, p_out = PRICES.get(MODEL, (0.0, 0.0))
        return (self.input_tokens * p_in + self.output_tokens * p_out) / 1_000_000


class Assistant:
    def __init__(self, client, db: ReadOnlyDB | None = None, search=None, max_rounds: int = 6) -> None:
        self.client = client
        self.db = db or ReadOnlyDB()
        self.search = search                      # optional ProductSearch
        self.max_rounds = max_rounds
        self.tools: dict[str, Callable[..., Any]] = {
            "describe_tables": lambda: self.db.schema(),
            "run_sql": lambda query: self.db.query(query).as_text(),
            "search_products": self._search_products,
        }

    def _search_products(self, text: str, k: int = 5) -> str:
        if self.search is None:
            return "product search is not enabled"
        return json.dumps(self.search.search(text, k))

    def _run_tool(self, call, turn: Turn) -> dict:
        try:
            args = json.loads(call.arguments or "{}")
            turn.tool_calls.append((call.name, args))
            func = self.tools[call.name]                          # KeyError -> unknown tool
            output = func(**args)
        except KeyError:
            output = f"error: unknown tool {call.name!r}"
        except QueryRejected as e:
            output = f"error: query rejected: {e}"                # the model can fix its SQL and retry
        except Exception as e:                                    # never crash the conversation
            log.exception("tool %s failed", call.name)
            output = f"error: {type(e).__name__}: {e}"
        return {"type": "function_call_output", "call_id": call.call_id, "output": str(output)}

    def ask(self, question: str) -> tuple[Answer, Turn]:
        turn = Turn(question)
        response = self.client.responses.parse(model=MODEL, instructions=INSTRUCTIONS, input=question,
                                               tools=TOOL_SCHEMAS, text_format=Answer)
        for _ in range(self.max_rounds):
            turn.input_tokens += response.usage.input_tokens
            turn.output_tokens += response.usage.output_tokens
            calls = [item for item in response.output if item.type == "function_call"]
            if not calls:
                if response.output_parsed is None:
                    raise RuntimeError("model returned no parsable answer")
                return response.output_parsed, turn
            outputs = [self._run_tool(c, turn) for c in calls]
            response = self.client.responses.parse(model=MODEL, instructions=INSTRUCTIONS,
                                                   previous_response_id=response.id, input=outputs,
                                                   tools=TOOL_SCHEMAS, text_format=Answer)
        raise RuntimeError(f"no answer after {self.max_rounds} tool rounds")