# SPDX-FileCopyrightText: 2026 Surendra Singh <66422685+bolnet@users.noreply.github.com>
# SPDX-License-Identifier: MIT
"""Round-level fact extraction (Phase 3.3, roadmap §A.2).

Two functions:

  extract_user_facts(turn, recent_context, model, client) -> list[ExtractedFact]
      Speaker-locked to the USER. Returns durable facts the user stated
      (preferences, plans, identity attributes...).

  extract_agent_facts(turn, recent_context, model, client) -> list[ExtractedFact]
      Speaker-locked to the ASSISTANT. Returns recommendations,
      decisions, commitments — anything an auditor must replay.

Both share:
  - JSON-only output, parsed with markdown-fence tolerance
  - Strict per-fact validation (drops malformed facts rather than crash)
  - source_span citations clamped to the turn's content range
  - Confidence clamped to [0.0, 1.0]

LLM client is injectable so tests can use a stub. By default, the client
is resolved via ``attestor.llm_trace.get_client_for_model`` against the
YAML-configured ``stack.llm.providers`` map — YAML is the sole source of
truth for provider ``base_url`` / ``api_key_env``.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import Any

from attestor.conversation.turns import ConversationTurn
from attestor.extraction.prompts import (
    AGENT_FACT_PROMPT_NAME,
    USER_FACT_PROMPT_NAME,
    format_agent_fact_prompt,
    format_user_fact_prompt,
)
from attestor.extraction.prompt_loader import prompt_version

logger = logging.getLogger("attestor.extraction.round")

# Prompt-template version each extractor stamps onto every ExtractedFact
# it returns. Bumping the version is a content edit + suffix bump in
# attestor/extraction/prompts/ — no signature change here.
USER_FACT_PROMPT_VERSION = prompt_version(USER_FACT_PROMPT_NAME)
AGENT_FACT_PROMPT_VERSION = prompt_version(AGENT_FACT_PROMPT_NAME)

def _default_model() -> str:
    """Resolve the extraction model from ``configs/attestor.yaml``."""
    from attestor.config import get_stack
    return get_stack().models.extraction


DEFAULT_MODEL = _default_model()
DEFAULT_MAX_TOKENS = 2048

USER_FACT_CATEGORIES = {
    "preference", "career", "project", "technical",
    "personal", "location", "relationship", "event", "financial",
}
AGENT_FACT_CATEGORIES = {
    "recommendation", "decision", "commitment",
    "constraint", "calculation", "refusal",
}


@dataclass(frozen=True)
class ExtractedFact:
    """One atomic fact emitted by the extractor.

    Bound to a single source turn; the (start, end) span lets the audit
    pipeline highlight which portion of the turn produced this fact.

    ``prompt_template`` / ``prompt_version`` (optional) record which
    externalized ``.md`` prompt produced this fact, so audit replay can
    correlate per-fact behavior changes back to a specific prompt edit.
    Defaults are ``None`` so legacy callers / tests that construct facts
    by hand continue to work without modification.
    """

    text: str
    category: str
    entity: str | None
    confidence: float
    source_span: list[int]
    speaker: str  # "user" | "assistant" | "<agent_id>"
    prompt_template: str | None = None
    prompt_version: str | None = None

    @property
    def source_start(self) -> int:
        return self.source_span[0]

    @property
    def source_end(self) -> int:
        return self.source_span[1]


# ──────────────────────────────────────────────────────────────────────────
# LLM client
# ──────────────────────────────────────────────────────────────────────────


def default_llm_client(model: str = DEFAULT_MODEL) -> Any:
    """Resolve the YAML-configured LLM client for ``model``.

    Routes through ``attestor.llm_trace.get_client_for_model`` so the
    provider URL/key come from ``stack.llm.providers`` in
    ``configs/attestor.yaml`` — never from a hardcoded constant or env
    fallback. Tests inject a stub via the ``client`` parameter on the
    ``extract_*`` functions, so this is only invoked when a real call
    is needed.
    """
    from attestor.llm_trace import get_client_for_model
    client, _clean_model = get_client_for_model(model)
    return client


def _call_llm(
    prompt: str,
    *,
    client: Any,
    model: str,
    max_tokens: int,
) -> str:
    """Invoke the LLM and return the raw response text.

    ``model`` may carry a ``provider/`` prefix; we strip it for the
    request payload because the prefix is just a routing hint for the
    pool. ``traced_create`` itself also tolerates a prefix, but stripping
    here keeps trace events labelled with the actual served model id.
    """
    from attestor.llm_trace import traced_create
    _, sep, tail = model.partition("/")
    clean_model = tail if sep else model
    response = traced_create(
        client,
        role="round_extractor",
        model=clean_model,
        max_tokens=max_tokens,
        timeout=30,
        max_retries=0,
        messages=[{"role": "user", "content": prompt}],
    )
    if not response.choices or not response.choices[0].message:
        logger.warning("LLM returned empty response; skipping extraction")
        return ""
    return response.choices[0].message.content or ""


# ──────────────────────────────────────────────────────────────────────────
# JSON parsing
# ──────────────────────────────────────────────────────────────────────────


from attestor.extraction.utils import strip_markdown_fences as _strip_markdown_fences  # noqa: E402,F401


def _parse_facts_payload(raw: str) -> list[dict]:
    """Parse the LLM response into a list of fact dicts.

    Tolerates: markdown fences, pure-array shape, the documented
    {"facts": [...]} envelope. Returns [] on any parse failure rather
    than raising — extractor failures must never break ingest.
    """
    text = _strip_markdown_fences(raw)
    if not text:
        return []
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError as e:
        logger.warning("extractor JSON parse failed: %s; raw=%r", e, raw[:200])
        return []
    if isinstance(parsed, list):
        return [f for f in parsed if isinstance(f, dict)]
    if isinstance(parsed, dict):
        facts = parsed.get("facts", [])
        if isinstance(facts, list):
            return [f for f in facts if isinstance(f, dict)]
    return []


def _validate_fact(
    raw: dict,
    *,
    speaker: str,
    content_len: int,
    allowed_categories: set,
    prompt_template: str | None = None,
    prompt_version_str: str | None = None,
) -> ExtractedFact | None:
    """Coerce a raw fact dict into an ExtractedFact, or return None.

    Validates:
      - text is non-empty
      - category is in the allowed set (else "general")
      - confidence is a float in [0, 1] (else 0.5)
      - source_span is [start, end] both ints, clamped to [0, content_len]
        and start <= end (else [0, 0])

    When ``prompt_template`` / ``prompt_version_str`` are supplied, they
    are stamped onto the resulting ``ExtractedFact`` so audit replay can
    correlate the fact back to the externalized prompt that produced it.
    """
    text = raw.get("text")
    if not isinstance(text, str) or not text.strip():
        return None

    category = raw.get("category")
    if not isinstance(category, str) or category not in allowed_categories:
        category = "general"

    entity_raw = raw.get("entity")
    entity = entity_raw if isinstance(entity_raw, str) and entity_raw else None

    conf_raw = raw.get("confidence", 0.5)
    try:
        confidence = float(conf_raw)
    except (TypeError, ValueError):
        confidence = 0.5
    confidence = max(0.0, min(1.0, confidence))

    span_raw = raw.get("source_span", [0, content_len])
    if (
        isinstance(span_raw, list) and len(span_raw) == 2
        and all(isinstance(x, int) for x in span_raw)
    ):
        start = max(0, min(content_len, span_raw[0]))
        end = max(start, min(content_len, span_raw[1]))
    else:
        start, end = 0, content_len

    return ExtractedFact(
        text=text.strip(),
        category=category,
        entity=entity,
        confidence=confidence,
        source_span=[start, end],
        speaker=speaker,
        prompt_template=prompt_template,
        prompt_version=prompt_version_str,
    )


# ──────────────────────────────────────────────────────────────────────────
# Public API
# ──────────────────────────────────────────────────────────────────────────


def extract_user_facts(
    turn: ConversationTurn,
    recent_context: str = "(none)",
    *,
    model: str = DEFAULT_MODEL,
    max_tokens: int = DEFAULT_MAX_TOKENS,
    client: Any | None = None,
) -> list[ExtractedFact]:
    """Extract durable facts from a USER turn.

    Speaker-locked: refuses to extract from anything but the user message
    via the IMPORTANT line in the prompt. Returns [] if the turn is not
    a user turn, the LLM returns empty, or parsing fails — never raises
    on extraction failure.
    """
    if not turn.is_user:
        logger.warning(
            "extract_user_facts called on a non-user turn (role=%r); "
            "returning empty", turn.role,
        )
        return []
    prompt = format_user_fact_prompt(
        ts=turn.ts.isoformat(),
        user_message=turn.content,
        recent_context_summary=recent_context or "(none)",
    )
    cli = client or default_llm_client(model)
    raw_text = _call_llm(prompt, client=cli, model=model, max_tokens=max_tokens)
    raw_facts = _parse_facts_payload(raw_text)

    facts: list[ExtractedFact] = []
    for rf in raw_facts:
        f = _validate_fact(
            rf, speaker=turn.speaker,
            content_len=len(turn.content),
            allowed_categories=USER_FACT_CATEGORIES,
            prompt_template=USER_FACT_PROMPT_NAME,
            prompt_version_str=USER_FACT_PROMPT_VERSION,
        )
        if f is not None:
            facts.append(f)
    return facts


def extract_agent_facts(
    turn: ConversationTurn,
    recent_context: str = "(none)",
    *,
    model: str = DEFAULT_MODEL,
    max_tokens: int = DEFAULT_MAX_TOKENS,
    client: Any | None = None,
) -> list[ExtractedFact]:
    """Extract durable statements from an ASSISTANT turn.

    Speaker-locked to the assistant message. Categories are
    recommendation / decision / commitment / constraint / calculation
    / refusal — the audit categories an auditor cares about.
    """
    if not turn.is_assistant:
        logger.warning(
            "extract_agent_facts called on a non-assistant turn (role=%r); "
            "returning empty", turn.role,
        )
        return []
    prompt = format_agent_fact_prompt(
        ts=turn.ts.isoformat(),
        assistant_message=turn.content,
        recent_context_summary=recent_context or "(none)",
    )
    cli = client or default_llm_client(model)
    raw_text = _call_llm(prompt, client=cli, model=model, max_tokens=max_tokens)
    raw_facts = _parse_facts_payload(raw_text)

    facts: list[ExtractedFact] = []
    for rf in raw_facts:
        f = _validate_fact(
            rf, speaker=turn.speaker,
            content_len=len(turn.content),
            allowed_categories=AGENT_FACT_CATEGORIES,
            prompt_template=AGENT_FACT_PROMPT_NAME,
            prompt_version_str=AGENT_FACT_PROMPT_VERSION,
        )
        if f is not None:
            facts.append(f)
    return facts
