"""Observer Audit Layer — LLM-based response quality check.

See design spec: docs/superpowers/specs/2026-04-19-observer-audit-layer-design.md

Runs after the Executor and after the regex-based ResponseValidator. Checks
the final response against four dimensions — Hallucination, Sycophancy,
Laziness, Tool-Ignorance — with per-dimension retry strategies.

The class is additive: it never replaces existing validators and fails open
(returns a pass result) on any internal failure so the core agent is never
blocked by a broken observer.
"""

from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass, field
from typing import Literal

from cognithor.utils.logging import get_logger

log = get_logger(__name__)


@dataclass(frozen=True)
class DimensionResult:
    """Per-dimension audit outcome."""

    name: Literal["hallucination", "sycophancy", "laziness", "tool_ignorance"]
    passed: bool
    reason: str
    evidence: str
    fix_suggestion: str


@dataclass(frozen=True)
class AuditResult:
    """Aggregate audit outcome for one observer call."""

    overall_passed: bool
    dimensions: dict[str, DimensionResult]
    retry_count: int
    final_action: Literal["pass", "rejected_with_retry", "delivered_with_warning"]
    retry_strategy: Literal["response_regen", "pge_reloop", "deliver", "deliver_with_warning"]
    model: str
    duration_ms: int
    degraded_mode: bool
    # None on success. On fail-open: "timeout", "parse_failed", "circuit_open",
    # "disabled". Task 13 will add "degraded" when degraded-mode falls back.
    error_type: str | None


@dataclass(frozen=True)
class PGEReloopDirective:
    """Observer signal requesting a full PGE re-loop (not just response regen)."""

    reason: Literal["tool_ignorance"]
    missing_data: str
    suggested_tools: list[str] = field(default_factory=list)


@dataclass(frozen=True)
class ResponseEnvelope:
    """Return type of Planner.formulate_response().

    A plain content payload plus an optional directive. Directive=None means
    'deliver content to user as-is'. Otherwise the Gateway catches the
    directive and re-enters the PGE loop.
    """

    content: str
    directive: PGEReloopDirective | None


from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from cognithor.config import CognithorConfig
    from cognithor.core.observer_store import AuditStore
    from cognithor.models import ToolResult


def _render_tool_line(r: Any) -> str:
    """Render a single tool result line, handling None error messages."""
    if r.success:
        return f"- {r.tool_name}: {r.content}"
    error_msg = r.error_message or "(no error message)"
    return f"- {r.tool_name}: ERROR: {error_msg}"


_SYSTEM_PROMPT = (
    "You are a quality auditor for LLM responses. Given a user message, the assistant's"
    " draft response, and any tool call results, judge the response against FOUR"
    " dimensions. Output ONLY valid JSON matching the schema below.\n"
    "\nDIMENSIONS:\n"
    "\n1. hallucination — Does the response contain claims not supported by the tool"
    " results? If no tools were called, the response must not contain specific factual"
    " claims that require external data (dates, numbers, named entities).\n"
    "\n2. sycophancy — Is the response flattering or agreement-seeking beyond what the"
    " user asked? Does it praise the user's question or avoid correcting obvious errors"
    " in the user's premise?\n"
    "\n3. laziness — Is the response vague, placeholder-heavy, or describes what the"
    ' assistant "would do" instead of actually answering?\n'
    "\n4. tool_ignorance — Was the user's question researchable/verifiable with the"
    " tools available, but no tool was actually called? If tools WERE called and used"
    " correctly, this passes.\n"
    "\nFor each dimension, output:\n"
    "  - passed: true | false\n"
    "  - reason: one-sentence explanation\n"
    "  - evidence: exact quote from the response (or empty string if passed)\n"
    "  - fix_suggestion: one-sentence change suggestion (or empty string if passed)\n"
    "\nOUTPUT SCHEMA (valid JSON, no additional text):\n"
    "{\n"
    '  "hallucination": {\n'
    '    "passed": true, "reason": "...", "evidence": "...", "fix_suggestion": "..."\n'
    "  },\n"
    '  "sycophancy": {\n'
    '    "passed": true, "reason": "...", "evidence": "...", "fix_suggestion": "..."\n'
    "  },\n"
    '  "laziness": {\n'
    '    "passed": true, "reason": "...", "evidence": "...", "fix_suggestion": "..."\n'
    "  },\n"
    '  "tool_ignorance": {\n'
    '    "passed": true, "reason": "...", "evidence": "...", "fix_suggestion": "..."\n'
    "  }\n"
    "}"
)


class ObserverAudit:
    """Run an LLM-based audit on a draft response. Fail-open by design."""

    def __init__(
        self,
        *,
        config: CognithorConfig,
        ollama_client: Any,
        audit_store: AuditStore,
    ) -> None:
        self._config = config
        self._ollama = ollama_client
        self._store = audit_store
        self._consecutive_failures = 0
        self._circuit_open = False
        self._degraded_warned = False

    def _build_prompt(
        self,
        *,
        user_message: str,
        response: str,
        tool_results: list[ToolResult],
    ) -> list[dict[str, str]]:
        """Compose system + user messages for the audit LLM call."""
        tool_section = (
            "\n".join(_render_tool_line(r) for r in tool_results) or "(no tool calls were made)"
        )
        user_payload = (
            f"USER MESSAGE:\n{user_message}\n\n"
            f"DRAFT RESPONSE:\n{response}\n\n"
            f"TOOL RESULTS:\n{tool_section}\n"
        )
        return [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {"role": "user", "content": user_payload},
        ]

    async def _resolve_model(self) -> tuple[str, bool]:
        """Return (model_name, degraded_mode).

        Falls back to planner model if observer model is missing.
        """
        observer_model = self._config.models.observer.name
        planner_model = self._config.models.planner.name
        try:
            available = await self._ollama.list_models()
        except Exception:
            # Can't list — assume observer model exists, proceed.
            return observer_model, False
        if observer_model in available:
            return observer_model, False
        if planner_model in available:
            if not self._degraded_warned:
                log.warning(
                    "observer_degraded_mode",
                    actual_model=planner_model,
                    intended_model=observer_model,
                )
                self._degraded_warned = True
            return planner_model, True
        # Both missing: observer cannot run
        if not self._degraded_warned:
            log.warning(
                "observer_disabled_runtime",
                observer_model=observer_model,
                planner_model=planner_model,
            )
            self._degraded_warned = True
        return "", True

    async def _call_llm_audit(
        self,
        *,
        messages: list[dict[str, str]],
        model_override: str | None = None,
    ) -> str | None:
        """Call the Observer LLM with JSON format + timeout. Returns None on any failure.

        Expects the ollama client's chat() to return a dict shaped
        ``{"message": {"content": "..."}}``. If a ``ChatResponse`` dataclass is
        wired in later (Task 18), adapt accordingly.
        """
        model_name = model_override or self._config.models.observer.name
        if not model_name:
            return None
        timeout = self._config.observer.timeout_seconds
        try:
            response = await asyncio.wait_for(
                self._ollama.chat(
                    model=model_name,
                    messages=messages,
                    options={"temperature": 0.1},
                    format_json=True,
                ),
                timeout=timeout,
            )
        except TimeoutError:
            log.warning(
                "observer_timeout",
                model=model_name,
                timeout_seconds=timeout,
            )
            return None
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            log.warning(
                "observer_connection_failed",
                model=model_name,
                error=str(exc),
            )
            return None

        content = response.get("message", {}).get("content", "")
        if not content:
            log.warning("observer_empty_response", model=model_name)
            return None
        return content

    def _parse_response(self, raw_text: str) -> dict[str, DimensionResult] | None:
        """Parse LLM JSON output. Returns dict of DimensionResults, or None on total failure.

        If only some dimensions are present, missing ones are filled with a
        'skipped' DimensionResult that counts as passed (so partial responses
        still allow the audit to proceed).
        """
        try:
            payload = json.loads(raw_text)
        except json.JSONDecodeError as exc:
            log.warning("observer_json_parse_failed", error=str(exc), raw_head=raw_text[:200])
            return None

        if not isinstance(payload, dict):
            log.warning("observer_schema_validation_failed", reason="top-level not object")
            return None

        all_dims = ("hallucination", "sycophancy", "laziness", "tool_ignorance")
        present = [
            d for d in all_dims if isinstance(payload.get(d), dict) and "passed" in payload[d]
        ]
        if not present:
            log.warning("observer_schema_validation_failed", reason="no dimensions present")
            return None

        dims: dict[str, DimensionResult] = {}
        for name in all_dims:
            entry = payload.get(name)
            if isinstance(entry, dict) and "passed" in entry:
                dims[name] = DimensionResult(
                    name=name,  # type: ignore[arg-type]
                    passed=bool(entry["passed"]),
                    reason=str(entry.get("reason", "")),
                    evidence=str(entry.get("evidence", "")),
                    fix_suggestion=str(entry.get("fix_suggestion", "")),
                )
            else:
                dims[name] = DimensionResult(
                    name=name,  # type: ignore[arg-type]
                    passed=True,  # skipped = pass
                    reason="skipped (missing from LLM response)",
                    evidence="",
                    fix_suggestion="",
                )
        return dims

    def _decide_retry_strategy(
        self,
        dimensions: dict[str, DimensionResult],
        retry_count: int,
    ) -> tuple[bool, Literal["response_regen", "pge_reloop", "deliver", "deliver_with_warning"]]:
        """Determine overall pass/fail and retry strategy.

        Priority when both blocking dimensions fail: tool_ignorance wins
        because gathering new data is more fundamental than rewording.
        """
        blocking = self._config.observer.blocking_dimensions
        blocking_failed = [
            name for name in blocking if name in dimensions and not dimensions[name].passed
        ]
        overall_passed = not blocking_failed

        if overall_passed:
            return True, "deliver"

        if retry_count >= self._config.observer.max_retries:
            return False, "deliver_with_warning"

        # Priority: tool_ignorance > hallucination (more fundamental fix).
        if "tool_ignorance" in blocking_failed:
            return False, "pge_reloop"
        if "hallucination" in blocking_failed:
            return False, "response_regen"

        # Any other blocking dimension currently collapses to response_regen.
        return False, "response_regen"

    async def audit(
        self,
        *,
        user_message: str,
        response: str,
        tool_results: list[ToolResult],
        session_id: str,
        retry_count: int = 0,
    ) -> AuditResult:
        """Run the four-dimension audit. Always returns an AuditResult — never raises."""
        start = time.monotonic()

        if self._circuit_open or not self._config.observer.enabled:
            # Fail-open placeholder: treat as pass so Core is never blocked.
            return self._fail_open_result(
                model=self._config.models.observer.name,
                reason="circuit_open" if self._circuit_open else "disabled",
                duration_ms=int((time.monotonic() - start) * 1000),
                retry_count=retry_count,
            )

        model, degraded = await self._resolve_model()
        if not model:
            # Both observer and planner models unavailable.
            return self._fail_open_result(
                model=self._config.models.observer.name,
                reason="model_unavailable",
                duration_ms=int((time.monotonic() - start) * 1000),
                retry_count=retry_count,
            )

        messages = self._build_prompt(
            user_message=user_message,
            response=response,
            tool_results=tool_results,
        )
        raw = await self._call_llm_audit(messages=messages, model_override=model)

        if raw is None:
            self._record_failure_for_circuit_breaker()
            result = self._fail_open_result(
                model=model,
                reason="timeout",
                duration_ms=int((time.monotonic() - start) * 1000),
                retry_count=retry_count,
            )
            self._persist(
                session_id=session_id,
                user_message=user_message,
                response=response,
                result=result,
            )
            return result

        dims = self._parse_response(raw)
        if dims is None:
            self._record_failure_for_circuit_breaker()
            result = self._fail_open_result(
                model=model,
                reason="parse_failed",
                duration_ms=int((time.monotonic() - start) * 1000),
                retry_count=retry_count,
            )
            self._persist(
                session_id=session_id,
                user_message=user_message,
                response=response,
                result=result,
            )
            return result

        self._consecutive_failures = 0  # successful call resets breaker

        overall_passed, strategy = self._decide_retry_strategy(dims, retry_count=retry_count)
        final_action: Literal["pass", "rejected_with_retry", "delivered_with_warning"]
        if overall_passed:
            final_action = "pass"
        elif strategy == "deliver_with_warning":
            final_action = "delivered_with_warning"
        else:
            final_action = "rejected_with_retry"

        result = AuditResult(
            overall_passed=overall_passed,
            dimensions=dims,
            retry_count=retry_count,
            final_action=final_action,
            retry_strategy=strategy,
            model=model,
            duration_ms=int((time.monotonic() - start) * 1000),
            degraded_mode=degraded,
            error_type=None,
        )
        self._persist(
            session_id=session_id,
            user_message=user_message,
            response=response,
            result=result,
        )
        return result

    def _fail_open_result(
        self,
        *,
        model: str,
        reason: str,
        duration_ms: int,
        retry_count: int = 0,
    ) -> AuditResult:
        """Construct a pass result used when the observer itself couldn't run."""

        def _skipped(name: str) -> DimensionResult:
            return DimensionResult(
                name=name,  # type: ignore[arg-type]
                passed=True,
                reason=f"fail_open: {reason}",
                evidence="",
                fix_suggestion="",
            )

        return AuditResult(
            overall_passed=True,
            dimensions={
                "hallucination": _skipped("hallucination"),
                "sycophancy": _skipped("sycophancy"),
                "laziness": _skipped("laziness"),
                "tool_ignorance": _skipped("tool_ignorance"),
            },
            retry_count=retry_count,
            final_action="pass",
            retry_strategy="deliver",
            model=model,
            duration_ms=duration_ms,
            degraded_mode=False,
            error_type=reason,
        )

    def _record_failure_for_circuit_breaker(self) -> None:
        self._consecutive_failures += 1
        if self._consecutive_failures >= self._config.observer.circuit_breaker_threshold:
            self._circuit_open = True
            log.info(
                "observer_circuit_open",
                consecutive_failures=self._consecutive_failures,
                threshold=self._config.observer.circuit_breaker_threshold,
            )

    def build_retry_feedback(self, result: AuditResult) -> dict[str, str]:
        """Produce a system-message payload for response-regen retries."""
        failed = [name for name, dim in result.dimensions.items() if not dim.passed]
        payload = {
            "observer_rejection": {
                "retry_count": result.retry_count,
                "max_retries": self._config.observer.max_retries,
                "dimensions_failed": failed,
                "reasons": [result.dimensions[n].reason for n in failed],
                "fix_suggestions": [result.dimensions[n].fix_suggestion for n in failed],
            }
        }
        return {"role": "system", "content": json.dumps(payload, ensure_ascii=False)}

    def build_pge_directive(self, result: AuditResult) -> PGEReloopDirective | None:
        """Extract a PGE re-loop directive from a tool_ignorance failure.

        Returns None if tool_ignorance passed (no re-loop needed). The directive
        contains the missing-data description and the Observer's suggested tools
        parsed out of the fix_suggestion.
        """
        ti = result.dimensions.get("tool_ignorance")
        if ti is None or ti.passed:
            return None
        # Extract tool suggestions by scanning the fix_suggestion for known
        # tool name patterns. Conservative: if none match, leave empty.
        known_tools = (
            "web_search",
            "web_fetch",
            "search_memory",
            "search_and_read",
            "read_file",
            "list_directory",
            "api_call",
            "exec_command",
        )
        suggested = [t for t in known_tools if t in ti.fix_suggestion]
        return PGEReloopDirective(
            reason="tool_ignorance",
            missing_data=ti.reason or ti.evidence,
            suggested_tools=suggested,
        )

    def _persist(
        self,
        *,
        session_id: str,
        user_message: str,
        response: str,
        result: AuditResult,
    ) -> None:
        """Forward an audit result to the store. Single call site for all 3 paths."""
        self._store.record(
            session_id=session_id,
            user_message=user_message,
            response=response,
            result=result,
        )
