"""The copilot: retrieve → prompt → model → verify citations, with a trace per question (lessons 09, 17, 20).

Two interchangeable "models":
  * OpenAIModel      the real thing (Responses API); used when OPENAI_API_KEY is set
  * ExtractiveModel  an offline stand-in with no LLM: it quotes the source sentences that best match
                     the question and cites them. It lets the whole system run, be tested and be traced
                     without a key. It is not a language model and is labelled as such everywhere.
"""

from __future__ import annotations

import hashlib
import html
import re
import time
from dataclasses import dataclass, field

from opentelemetry.trace import SpanKind, Status, StatusCode

from .guardrails import injection_signals, redact
from .index import HybridIndex, tokenize
from .ingest import Chunk
from .telemetry import Telemetry

PROMPT_VERSION = "answer-v3"
INSTRUCTIONS = """You answer questions about the learn-with-project course lessons in the <source> tags.
- Use only the sources. Cite every claim inline as [1], [2].
- If the sources do not answer the question, reply exactly: I don't know.
- At most 4 sentences, for a junior engineer."""
PRICES = {"gpt-5.6-luna": (0.20, 1.20), "extractive-v1": (0.0, 0.0)}      # USD per 1M tokens (in, out)
FALLBACK = "I don't know."


@dataclass
class ModelReply:
    text: str
    input_tokens: int
    output_tokens: int
    model: str
    response_id: str
    finish_reason: str = "stop"


class OpenAIModel:
    provider = "openai"

    def __init__(self, client, model: str = "gpt-5.6-luna", max_output_tokens: int = 400) -> None:
        self.client, self.model, self.max_output_tokens = client, model, max_output_tokens

    def generate(self, instructions: str, input: str) -> ModelReply:
        r = self.client.responses.create(model=self.model, instructions=instructions, input=input,
                                         max_output_tokens=self.max_output_tokens)
        finish = "stop" if r.status == "completed" else "length"
        return ModelReply(r.output_text, r.usage.input_tokens, r.usage.output_tokens, r.model, r.id, finish)


class ExtractiveModel:
    """Offline stand-in (NOT an LLM): quotes the best-matching source sentences, with citations."""
    provider = "offline"
    model = "extractive-v1"
    SOURCE = re.compile(r'<source id="(\d+)"[^>]*>\n(.*?)\n</source>', re.S)

    GENERIC = frozenset("use using used should which will would there their does work works".split())

    def __init__(self, min_overlap: float = 0.4, max_sentences: int = 2) -> None:
        self.min_overlap, self.max_sentences = min_overlap, max_sentences

    @classmethod
    def _terms(cls, text: str) -> set[str]:
        """Content words, crudely stemmed so that page/pages/paging and cache/caching match."""
        out = set()
        for t in tokenize(text):
            if t in cls.GENERIC:
                continue
            for suffix in ("ing", "es", "ed", "s", "e"):
                if len(t) > 4 and t.endswith(suffix):
                    t = t[: -len(suffix)]
                    break
            out.add(t)
        return out

    def generate(self, instructions: str, input: str) -> ModelReply:
        question = html.unescape(re.search(r"<question>(.*?)</question>", input, re.S).group(1))
        terms = self._terms(question)
        candidates = []
        for sid, body in self.SOURCE.findall(input):
            for line in html.unescape(body).splitlines():
                for sentence in re.split(r"(?<=[.!?])\s+", line):
                    words = sentence.split()
                    if len(words) < 6 or not sentence[:1].isupper() or not sentence.rstrip().endswith((".", "!")):
                        continue                                       # skip code, headings, questions
                    overlap = len(terms & self._terms(sentence)) / max(1, len(terms))
                    candidates.append((overlap, -int(sid), sentence.strip(), int(sid)))
        # prefer higher-ranked sources: a small penalty per rank position
        ranked = sorted(candidates, key=lambda c: c[0] - 0.08 * (-c[1] - 1), reverse=True)
        chosen = [c for c in ranked if c[0] >= self.min_overlap][:self.max_sentences]
        text = " ".join(f"{s} [{sid}]" for _, _, s, sid in chosen) or FALLBACK
        digest = hashlib.sha1(input.encode()).hexdigest()[:12]
        return ModelReply(text, max(1, len(input) // 4), max(1, len(text) // 4), self.model, f"offline_{digest}")


@dataclass
class Answer:
    question: str
    text: str
    sources: list[Chunk] = field(default_factory=list)
    grounded: bool = False
    input_tokens: int = 0
    output_tokens: int = 0
    cost_usd: float = 0.0
    latency_s: float = 0.0
    trace_id: str = ""
    guardrail_signals: list[str] = field(default_factory=list)


class Copilot:
    def __init__(self, index: HybridIndex, model, telemetry: Telemetry, budget_words: int = 700) -> None:
        self.index, self.model, self.tel, self.budget_words = index, model, telemetry, budget_words

    # ------------------------------------------------------------------ retrieval
    def retrieve(self, question: str, course: str | None = None) -> list[Chunk]:
        """Top hits, at most one chunk per lesson, within a word budget (lessons 07, 09)."""
        chosen, lessons, used = [], set(), 0
        for hit in self.index.search(question, k=20, course=course):
            lesson = hit.chunk.url.split("#")[0]
            words = len(hit.chunk.text.split())
            if lesson in lessons or (chosen and used + words > self.budget_words):
                continue
            chosen.append(hit.chunk)
            lessons.add(lesson)
            used += words
        return chosen

    @staticmethod
    def build_input(question: str, chunks: list[Chunk]) -> str:
        blocks = [f'<source id="{i}" title="{html.escape(c.header())}">\n{html.escape(c.text, quote=False)}\n</source>'
                  for i, c in enumerate(chunks, start=1)]
        return "\n".join(blocks) + f"\n<question>{html.escape(question, quote=False)}</question>"

    # ------------------------------------------------------------------ the request
    def ask(self, question: str, *, user_id: str = "anonymous", session_id: str | None = None,
            course: str | None = None) -> Answer:
        tel, started = self.tel, time.perf_counter()
        with tel.tracer.start_as_current_span("copilot.ask") as root:
            root.set_attributes({"app.prompt.version": PROMPT_VERSION, "app.index.chunks": len(self.index),
                                 "app.index.embedder": self.index.embedder.name,
                                 "app.user.hash": hashlib.sha256(user_id.encode()).hexdigest()[:12]})
            if session_id:
                root.set_attribute("gen_ai.conversation.id", session_id)

            signals = injection_signals(question)                     # tripwire: record, alert, continue
            if signals:
                root.set_attribute("app.guardrail.injection", signals)
                tel.guardrail_hits.add(1, {"type": "injection"})
            clean_question, pii = redact(question)                     # nothing personal leaves the process
            if pii:
                root.set_attribute("app.guardrail.pii_redacted", len(pii))
                tel.guardrail_hits.add(1, {"type": "pii"})

            with tel.tracer.start_as_current_span("retrieval site-lessons", attributes={
                    "gen_ai.operation.name": "retrieval", "gen_ai.data_source.id": "site-lessons"}) as span:
                chunks = self.retrieve(clean_question, course)
                span.set_attribute("app.retrieval.urls", [c.url for c in chunks])

            answer = Answer(question, FALLBACK, guardrail_signals=signals)
            if chunks:
                reply = self._generate(clean_question, chunks)
                answer.input_tokens, answer.output_tokens = reply.input_tokens, reply.output_tokens
                price_in, price_out = PRICES.get(self.model.model, (0.0, 0.0))
                answer.cost_usd = (reply.input_tokens * price_in + reply.output_tokens * price_out) / 1e6
                with tel.tracer.start_as_current_span("verify_citations") as span:
                    cited = sorted({int(n) for n in re.findall(r"\[(\d+)\]", reply.text)})
                    valid = [n for n in cited if 1 <= n <= len(chunks)]
                    answer.grounded = bool(valid) and len(valid) == len(cited)
                    span.set_attributes({"app.citations": cited, "app.grounded": answer.grounded})
                    if not answer.grounded and reply.text != FALLBACK:
                        span.set_status(Status(StatusCode.ERROR, "answer cites no valid source"))
                if answer.grounded:
                    answer.text, answer.sources = reply.text, [chunks[n - 1] for n in valid]

            answer.latency_s = time.perf_counter() - started
            answer.trace_id = format(root.get_span_context().trace_id, "032x")
            root.set_attributes({"app.grounded": answer.grounded, "app.cost_usd": answer.cost_usd})
            tel.answers.add(1, {"grounded": str(answer.grounded).lower()})
            tel.cost.add(answer.cost_usd, {"gen_ai.request.model": self.model.model})
            tel.request_duration.record(answer.latency_s, {"route": "ask"})
            return answer

    def _generate(self, question: str, chunks: list[Chunk]) -> ModelReply:
        tel = self.tel
        attrs = {"gen_ai.operation.name": "chat", "gen_ai.provider.name": self.model.provider,
                 "gen_ai.request.model": self.model.model}
        started = time.perf_counter()
        with tel.tracer.start_as_current_span(f"chat {self.model.model}", kind=SpanKind.CLIENT,
                                              attributes=attrs) as span:
            try:
                reply = self.model.generate(INSTRUCTIONS, self.build_input(question, chunks))
            except Exception as e:                                     # recorded on the span automatically
                span.set_attribute("error.type", type(e).__name__)
                tel.operation_duration.record(time.perf_counter() - started, {**attrs, "error.type": type(e).__name__})
                raise
            span.set_attributes({"gen_ai.response.model": reply.model, "gen_ai.response.id": reply.response_id,
                                 "gen_ai.usage.input_tokens": reply.input_tokens,
                                 "gen_ai.usage.output_tokens": reply.output_tokens,
                                 "gen_ai.response.finish_reasons": [reply.finish_reason]})
            labels = {**attrs, "gen_ai.response.model": reply.model}
            tel.token_usage.record(reply.input_tokens, {**labels, "gen_ai.token.type": "input"})
            tel.token_usage.record(reply.output_tokens, {**labels, "gen_ai.token.type": "output"})
            tel.operation_duration.record(time.perf_counter() - started, labels)
            return reply