"""
cortex/api/routes/extraction_test.py — Diagnostic extraction test endpoint.

POST /api/v1/extraction/test — runs the full extraction pipeline as a dry-run
on provided conversation input and returns a detailed trace of every step,
without writing anything to storage.

Issue: R-242
"""

from __future__ import annotations

import dataclasses
import logging
from typing import Any, Dict, List, Optional

from fastapi import APIRouter, Depends, Request, status
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from cortex.api.deps import PREFIX, get_plugin
from cortex.extraction.prompts_v2 import ALLOWED_PREDICATES

logger = logging.getLogger(__name__)

router = APIRouter()


# ---------------------------------------------------------------------------
# Request / Response models
# ---------------------------------------------------------------------------

class ConversationMessage(BaseModel):
    role: str = "user"
    content: str = Field(default="", max_length=4000)


class ExtractionTestRequest(BaseModel):
    conversation: List[ConversationMessage] = Field(default_factory=list, max_length=20)
    # NOTE: owner_id removed — this endpoint is a dry-run that never writes to
    # storage, so accepting an owner_id field was misleading (it was silently
    # ignored).  See 100yenadmin/electric-sheep#1546.
    session_date: Optional[str] = None
    speaker_a: Optional[str] = None
    speaker_b: Optional[str] = None
    prompt_version: str = "v2"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _claim_to_dict(claim: Any) -> Dict[str, Any]:
    """Convert a ClaimExtraction dataclass to a plain dict."""
    if dataclasses.is_dataclass(claim) and not isinstance(claim, type):
        d = dataclasses.asdict(claim)
        return d
    if isinstance(claim, dict):
        return claim
    return {k: v for k, v in claim.__dict__.items() if not k.startswith("_")}


def _validate_claim(claim_dict: Dict[str, Any]) -> Dict[str, Any]:
    """Run validation checks on a single claim dict."""
    from cortex.extraction.models import VALID_CATEGORY

    warnings: List[str] = []
    empty_fields: List[str] = []

    has_subject = bool(claim_dict.get("subject"))
    has_predicate = bool(claim_dict.get("predicate"))
    has_object_value = bool(claim_dict.get("object_value"))
    has_category = bool(claim_dict.get("category"))
    has_confidence = claim_dict.get("confidence") is not None

    predicate = str(claim_dict.get("predicate", "")).lower().strip()
    predicate_in_allowed_set = predicate in ALLOWED_PREDICATES

    category = str(claim_dict.get("category", ""))
    category_valid = category in VALID_CATEGORY

    if has_confidence:
        confidence = claim_dict.get("confidence")
        try:
            conf_float = float(confidence)
            confidence_in_range = 0.0 <= conf_float <= 1.0
        except (TypeError, ValueError):
            confidence_in_range = False
    else:
        confidence_in_range = None

    object_value = str(claim_dict.get("object_value", ""))
    object_value_length = len(object_value)

    # Check for empty optional fields
    optional_fields = [
        "cause", "effect", "state_before", "state_after",
        "event_date", "temporal_marker", "human_family",
    ]
    for f in optional_fields:
        val = claim_dict.get(f)
        if val is None or (isinstance(val, str) and not val.strip()):
            empty_fields.append(f)

    # Generate warnings
    if not has_subject:
        warnings.append("Missing subject field")
    if not has_predicate:
        warnings.append("Missing predicate field")
    if not has_object_value:
        warnings.append("Missing object_value field")
    if not has_category:
        warnings.append("Empty category field")
    if not predicate_in_allowed_set and has_predicate:
        warnings.append(f"Predicate '{predicate}' not in allowed set")
    if not category_valid and has_category:
        warnings.append(f"Category '{category}' not in VALID_CATEGORY")
    if confidence_in_range is None:
        warnings.append("Missing confidence field")
    elif not confidence_in_range:
        warnings.append(f"Confidence {claim_dict.get('confidence')} out of [0, 1] range")
    if object_value_length > 150:
        warnings.append(f"object_value length ({object_value_length}) exceeds 150 char recommendation")

    return {
        "has_subject": has_subject,
        "has_predicate": has_predicate,
        "has_object_value": has_object_value,
        "has_category": has_category,
        "has_confidence": has_confidence,
        "predicate_in_allowed_set": predicate_in_allowed_set,
        "category_valid": category_valid,
        "confidence_in_range": confidence_in_range,
        "object_value_length": object_value_length,
        "empty_fields": empty_fields,
        "warnings": warnings,
    }


def _run_filter_tracking(
    claims: list,
    known_subjects: set,
    exclude_categories: set,
    enable_stale_plan_filter: bool,
) -> Dict[str, Any]:
    """Run filters both independently (for diagnostics) and sequentially (for production match).

    Independent per-filter tracking shows what each filter would remove from the
    full claim set.  Sequential application mirrors pipeline.py production order
    to produce the authoritative ``claims_after_filters`` list.

    Returns per-claim filter results, aggregate dropped_by_filter counts (independent),
    and the sequentially-filtered claim IDs (``sequential_passed_ids``).
    """
    from cortex.extraction.predicate_filter import (
        filter_claims as predicate_filter_fn,
        filter_unknown_subjects as unknown_subjects_fn,
        cap_engineering_claims as cap_engineering_fn,
    )
    from cortex.extraction.tautology_filter import filter_tautologies as tautology_filter_fn
    from cortex.extraction.stale_plan_filter import filter_stale_plans as stale_plan_filter_fn

    # --- Independent filter runs (for per-claim diagnostics) ---
    after_predicate = predicate_filter_fn(claims)
    after_tautology = tautology_filter_fn(claims, agent_name="eva")
    after_stale_plan = stale_plan_filter_fn(claims) if enable_stale_plan_filter else claims

    if exclude_categories:
        after_category = [
            c for c in claims
            if str(getattr(c, "category", None) or "misc").lower()
            not in exclude_categories
        ]
    else:
        after_category = claims

    if known_subjects:
        after_unknown = unknown_subjects_fn(claims, known_subjects)
    else:
        after_unknown = claims

    after_engineering = cap_engineering_fn(claims)

    predicate_passed_ids = {id(c) for c in after_predicate}
    tautology_passed_ids = {id(c) for c in after_tautology}
    stale_plan_passed_ids = {id(c) for c in after_stale_plan}
    category_passed_ids = {id(c) for c in after_category}
    unknown_passed_ids = {id(c) for c in after_unknown}
    engineering_passed_ids = {id(c) for c in after_engineering}

    # --- Per-claim filter results (independent) ---
    per_claim_filters: List[Dict[str, Any]] = []
    dropped_by_filter = {
        "predicate_filter": 0,
        "tautology_filter": 0,
        "stale_plan_filter": 0,
        "category_exclusion": 0,
        "unknown_subject": 0,
        "engineering_cap": 0,
    }

    for claim in claims:
        cid = id(claim)
        pf = cid in predicate_passed_ids
        tf = cid in tautology_passed_ids
        sf = cid in stale_plan_passed_ids
        cf = cid in category_passed_ids
        uf = cid in unknown_passed_ids
        ef = cid in engineering_passed_ids

        if not pf:
            dropped_by_filter["predicate_filter"] += 1
        if not tf:
            dropped_by_filter["tautology_filter"] += 1
        if not sf:
            dropped_by_filter["stale_plan_filter"] += 1
        if not cf:
            dropped_by_filter["category_exclusion"] += 1
        if not uf:
            dropped_by_filter["unknown_subject"] += 1
        if not ef:
            dropped_by_filter["engineering_cap"] += 1

        per_claim_filters.append({
            "predicate_filter": {"passed": pf, "reason": None if pf else "blocked predicate or subject"},
            "tautology_filter": {"passed": tf, "reason": None if tf else "tautological claim"},
            "stale_plan_filter": {"passed": sf, "reason": None if sf else "stale planning claim"},
            "category_exclusion": {"passed": cf, "reason": None if cf else "excluded category"},
            "unknown_subject": {"passed": uf, "reason": None if uf else "unknown subject"},
            "engineering_cap": {"passed": ef, "reason": None if ef else "engineering claim cap exceeded"},
        })

    # --- Sequential filter application (matches pipeline.py production order) ---
    # Order: predicate → tautology → stale_plan → category_exclusion → unknown_subject → engineering_cap
    seq = predicate_filter_fn(claims)
    seq = tautology_filter_fn(seq, agent_name="eva")
    if enable_stale_plan_filter:
        seq = stale_plan_filter_fn(seq)
    if exclude_categories:
        seq = [
            c for c in seq
            if str(getattr(c, "category", None) or "misc").lower()
            not in exclude_categories
        ]
    if known_subjects:
        seq = unknown_subjects_fn(seq, known_subjects)
    seq = cap_engineering_fn(seq)
    sequential_passed_ids = {id(c) for c in seq}

    return {
        "per_claim_filters": per_claim_filters,
        "dropped_by_filter": dropped_by_filter,
        "sequential_passed_ids": sequential_passed_ids,
    }


# ---------------------------------------------------------------------------
# Diagnostic pipeline runner
# ---------------------------------------------------------------------------

async def _run_diagnostic_extraction(
    plugin: Any,
    body: ExtractionTestRequest,
) -> Dict[str, Any]:
    """Run extraction pipeline in diagnostic mode and return full trace."""
    from cortex.extraction.pipeline import ExtractionPipeline
    from cortex.extraction.utils import _split_batches

    config = plugin.config

    # Build messages list from conversation
    messages = [{"role": m.role, "content": m.content} for m in body.conversation]

    # Determine prompt version
    prompt_version = body.prompt_version or getattr(config, "extraction_prompt_version", "v2")

    # Instantiate pipeline the same way production does
    pipeline = ExtractionPipeline(
        llm=plugin.llm,
        model=getattr(config, "extraction_model", None),
        batch_size=getattr(config, "extraction_batch_size", 5),
        custom_instructions=getattr(config, "extraction_custom_instructions", ""),
        max_claims_per_extraction=getattr(config, "extraction_max_claims", 25),
        max_entities_per_extraction=getattr(config, "extraction_max_entities", 100),
        prompt_version=prompt_version,
        known_subjects=list(getattr(config, "extraction_known_subjects", []) or []),
        enable_stale_plan_filter=getattr(config, "extraction_enable_stale_plan_filter", True),
        exclude_categories=getattr(config, "extraction_exclude_categories", None),
    )
    pipeline.extract_assistant_turns = getattr(config, "extract_assistant_turns", False)

    # --- Step 1: Noise filter ---
    raw_count = len(messages)
    noise_filter = pipeline.noise_filter
    filtered = noise_filter.filter_conversation(messages)
    filtered_ids = {id(m) for m in filtered}
    noise_filtered_messages = [m for m in messages if id(m) not in filtered_ids]
    noise_removed_count = len(noise_filtered_messages)

    # Gate assistant turns (tracked separately from noise filter)
    assistant_extracted = pipeline.extract_assistant_turns
    assistant_turns_removed_count = 0
    if not assistant_extracted:
        before_assistant_gate = len(filtered)
        filtered = [m for m in filtered if m.get("role", "user") != "assistant"]
        assistant_turns_removed_count = before_assistant_gate - len(filtered)

    noise_info = {
        "input_count": raw_count,
        "noise_removed_count": noise_removed_count,
        "assistant_turns_removed_count": assistant_turns_removed_count,
        "output_count": len(filtered),
        "filtered_messages": [
            {"role": m.get("role", "user"), "content": m.get("content", "")[:200]}
            for m in noise_filtered_messages
        ],
        "assistant_turns_extracted": assistant_extracted,
    }

    if not filtered:
        return {
            "input": {
                "message_count": raw_count,
                "prompt_version": prompt_version,
                "system_prompt_length": len(pipeline._system_prompt),
            },
            "noise_filter": noise_info,
            "batches": [],
            "claims_raw": [],
            "claims_after_filters": [],
            "summary": {
                "total_extracted": 0,
                "total_after_filters": 0,
                "dropped_by_filter": {},
                "category_distribution": {},
                "confidence_distribution": {"0.9+": 0, "0.7-0.89": 0, "<0.7": 0},
                "warnings": [],
            },
        }

    # --- Step 2: Batch splitting ---
    batches = _split_batches(filtered, pipeline.batch_size)

    # --- Step 3: Process each batch (LLM calls) ---
    if prompt_version == "v2":
        from cortex.extraction.prompts_v2 import format_conversation as format_conv
    else:
        from cortex.extraction.prompts import format_conversation as format_conv

    batch_traces: List[Dict[str, Any]] = []
    all_raw_claims: list = []

    from cortex.llm.router import resolve_routed_model

    for batch_idx, batch in enumerate(batches):
        if prompt_version == "v2":
            conversation_text = format_conv(
                batch,
                session_date=body.session_date,
                speaker_a=body.speaker_a,
                speaker_b=body.speaker_b,
            )
        else:
            conversation_text = format_conv(batch)

        if not conversation_text.strip():
            batch_traces.append({
                "batch_index": batch_idx,
                "message_count": len(batch),
                "raw_llm_response": "",
                "parse_result": {
                    "success": False,
                    "claims_parsed": 0,
                    "parse_errors": ["Empty conversation text after formatting"],
                },
            })
            continue

        # Call LLM
        kwargs: Dict[str, Any] = {"response_format": "json", "max_tokens": 8192}
        _routed_model = resolve_routed_model(pipeline.llm, "extraction", pipeline.model)
        if _routed_model:
            kwargs["model"] = _routed_model

        raw_json = ""
        parse_errors: List[str] = []
        batch_claims: list = []

        try:
            _resp = await pipeline.llm.complete(
                system=pipeline._system_prompt,
                user=conversation_text,
                **kwargs,
            )
            raw_json = _resp.text if hasattr(_resp, "text") else str(_resp)
        except Exception as exc:
            parse_errors.append(f"LLM call failed: {exc}")
            batch_traces.append({
                "batch_index": batch_idx,
                "message_count": len(batch),
                "raw_llm_response": "",
                "parse_result": {
                    "success": False,
                    "claims_parsed": 0,
                    "parse_errors": parse_errors,
                },
            })
            continue

        # Parse claims using the pipeline's parser
        try:
            batch_claims = pipeline._parse_claims(raw_json, f"diag-{batch_idx}")
        except Exception as exc:
            parse_errors.append(f"Parse failed: {exc}")

        all_raw_claims.extend(batch_claims)

        raw_json_truncated = raw_json[:50000] if raw_json and len(raw_json) > 50000 else raw_json
        trace_entry: Dict[str, Any] = {
            "batch_index": batch_idx,
            "message_count": len(batch),
            "raw_llm_response": raw_json_truncated,
            "parse_result": {
                "success": len(batch_claims) > 0 or not parse_errors,
                "claims_parsed": len(batch_claims),
                "parse_errors": parse_errors,
            },
        }
        if raw_json and len(raw_json) > 50000:
            trace_entry["raw_llm_response_truncated"] = True
        batch_traces.append(trace_entry)

    input_info = {
        "message_count": raw_count,
        "prompt_version": prompt_version,
        "system_prompt_length": len(pipeline._system_prompt),
    }

    # --- Step 5: Filter tracking ---
    filter_results = _run_filter_tracking(
        all_raw_claims,
        known_subjects=pipeline.known_subjects,
        exclude_categories=pipeline.exclude_categories,
        enable_stale_plan_filter=pipeline.enable_stale_plan_filter,
    )

    # --- Step 6: Build per-claim trace with validation ---
    claims_raw_trace: List[Dict[str, Any]] = []
    for i, claim in enumerate(all_raw_claims):
        claim_dict = _claim_to_dict(claim)
        validation = _validate_claim(claim_dict)
        filters = filter_results["per_claim_filters"][i] if i < len(filter_results["per_claim_filters"]) else {}

        claims_raw_trace.append({
            "claim": claim_dict,
            "filters": filters,
            "validation": validation,
        })

    # --- Step 7: Build claims_after_filters (sequential, matches production) ---
    sequential_passed_ids = filter_results["sequential_passed_ids"]
    claims_after_filters = [
        _claim_to_dict(c) for c in all_raw_claims if id(c) in sequential_passed_ids
    ]

    # --- Step 8: Summary ---
    total_extracted = len(all_raw_claims)
    total_after_filters = len(claims_after_filters)

    # Category distribution (from filtered claims)
    cat_dist: Dict[str, int] = {}
    for c in claims_after_filters:
        cat = c.get("category", "misc")
        cat_dist[cat] = cat_dist.get(cat, 0) + 1

    # Confidence distribution (from filtered claims)
    conf_dist = {"0.9+": 0, "0.7-0.89": 0, "<0.7": 0}
    for c in claims_after_filters:
        try:
            conf = float(c.get("confidence", 0))
        except (TypeError, ValueError):
            conf = 0.0
        if conf >= 0.9:
            conf_dist["0.9+"] += 1
        elif conf >= 0.7:
            conf_dist["0.7-0.89"] += 1
        else:
            conf_dist["<0.7"] += 1

    # Collect global warnings
    global_warnings: List[str] = []
    empty_cat_count = sum(
        1 for c in all_raw_claims
        if not (getattr(c, "category", None) if hasattr(c, "category") else "")
    )
    if empty_cat_count:
        global_warnings.append(f"{empty_cat_count} claim(s) have empty category field")

    return {
        "input": input_info,
        "noise_filter": noise_info,
        "batches": batch_traces,
        "claims_raw": claims_raw_trace,
        "claims_after_filters": claims_after_filters,
        "summary": {
            "total_extracted": total_extracted,
            "total_after_filters": total_after_filters,
            "dropped_by_filter": filter_results["dropped_by_filter"],
            "category_distribution": cat_dist,
            "confidence_distribution": conf_dist,
            "warnings": global_warnings,
        },
    }


# ---------------------------------------------------------------------------
# Route registration
# ---------------------------------------------------------------------------

def _make_routes(verify_api_key):
    """Register the extraction test endpoint."""

    @router.post(
        f"{PREFIX}/extraction/test",
        status_code=status.HTTP_200_OK,
        summary="Diagnostic extraction dry-run",
        tags=["extraction"],
        operation_id="extraction_test",
    )
    async def extraction_test(
        request: Request,
        body: ExtractionTestRequest,
        _: None = Depends(verify_api_key),
        plugin: Any = Depends(get_plugin),
    ) -> Dict[str, Any]:
        """Run the extraction pipeline as a dry-run and return a detailed trace.

        No data is written to storage. This endpoint is for diagnostic use only.
        """
        try:
            result = await _run_diagnostic_extraction(plugin, body)
            return result
        except Exception as exc:
            logger.error("Extraction test failed: %s", exc, exc_info=True)
            return JSONResponse(
                status_code=500,
                content={"error": "Extraction pipeline failed. Check server logs for details."},
            )
