"""
RubricJudge — evaluates spans using domain-specific rubrics.

Replaces LLMJudge's freeform evaluation with structured criteria,
producing StructuredDiagnosis for precise remediation.
"""

import json
import logging
import time
from typing import Any

from .models import EnrichedContext, JudgeRubric, RootCause, StructuredDiagnosis

logger = logging.getLogger("aigie.judge.rubric_judge")


class RubricJudge:
    """Evaluates a span using a domain-specific rubric."""

    def __init__(self, rubric: JudgeRubric, llm_client: Any):
        self._rubric = rubric
        self._llm = llm_client

    async def evaluate(
        self,
        span_id: str,
        input_messages: list[dict],
        output_content: str,
        enriched_context: EnrichedContext,
    ) -> StructuredDiagnosis:
        """Evaluate using rubric criteria. Returns structured diagnosis."""
        start = time.perf_counter()
        prompt = self._build_prompt(input_messages, output_content, enriched_context)
        try:
            raw_response = await self._llm.complete(prompt)
            raw_json = self._extract_json(raw_response)
            eval_time = (time.perf_counter() - start) * 1000
            return self._parse_response(raw_json, span_id, eval_time)
        except Exception as e:
            eval_time = (time.perf_counter() - start) * 1000
            logger.warning(f"RubricJudge evaluation failed: {e}")
            return StructuredDiagnosis(
                issue_domain=self._rubric.domain,
                issue_type="",
                severity="info",
                span_id=span_id,
                root_cause=RootCause.UNKNOWN,
                confidence=0.1,
                reasoning=f"Evaluation failed: {e}",
                judge_id=f"{self._rubric.domain}_judge",
                judge_model="",
                evaluation_time_ms=eval_time,
            )

    def _build_prompt(
        self,
        input_messages: list[dict],
        output_content: str,
        enriched_context: EnrichedContext,
    ) -> str:
        """Build focused evaluation prompt from rubric."""
        criteria_section = ""
        for c in self._rubric.criteria:
            criteria_section += f"  {c.name}: {c.question}\n"
            criteria_section += f"    Scoring: {c.scoring} (weight: {c.weight})\n"
            criteria_section += f"    If fails, root cause: {c.failure_root_cause}\n\n"

        input_preview = ""
        for msg in input_messages[-3:]:
            role = msg.get("role", "unknown")
            content = str(msg.get("content", ""))[:500]
            input_preview += f"[{role}]: {content}\n"

        prev_context = ""
        if enriched_context.previous_step_had_issue:
            prev_context = (
                f"\nPREVIOUS STEP CONTEXT:\n"
                f"  Previous issue: {enriched_context.previous_step_issue_type}\n"
                f"  Consecutive issues: {enriched_context.consecutive_issues}\n"
                f"  This may indicate a cascade — consider if this step's output compounds the problem.\n"
            )

        schema_json = json.dumps(self._rubric.output_schema, indent=2)

        return f"""You are a {self._rubric.domain} evaluation specialist.

CRITERIA TO EVALUATE:
{criteria_section}
INPUT CONTEXT:
{input_preview}

OUTPUT TO EVALUATE:
{output_content[:2000]}
{prev_context}
Return JSON matching this schema:
{schema_json}

IMPORTANT:
- For root_cause, use ONLY one of: {self._rubric.root_cause_options}
- Score each criterion 0.0-1.0 (1.0 = no issues)
- Set has_issues=false and omit primary_issue if output is acceptable
- Be specific in evidence — quote the problematic text"""

    def _extract_json(self, raw_response: str) -> dict:
        """Extract JSON from LLM response, handling markdown code blocks."""
        text = raw_response.strip()
        if text.startswith("```"):
            lines = text.split("\n")
            json_lines = [line for line in lines if not line.startswith("```")]
            text = "\n".join(json_lines)
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            start = text.find("{")
            end = text.rfind("}") + 1
            if start >= 0 and end > start:
                return json.loads(text[start:end])
            raise

    def _parse_response(self, raw: dict, span_id: str, eval_time_ms: float) -> StructuredDiagnosis:
        """Parse LLM JSON output into StructuredDiagnosis. Graceful on malformed input."""
        has_issues = raw.get("has_issues", False)
        primary = raw.get("primary_issue", {}) or {}
        fix = raw.get("recommended_fix", {}) or {}

        if not has_issues:
            return StructuredDiagnosis(
                issue_domain=self._rubric.domain,
                issue_type="",
                severity="info",
                span_id=span_id,
                root_cause=RootCause.UNKNOWN,
                evidence=[],
                confidence=raw.get("confidence", 0.5),
                reasoning=raw.get("reasoning", ""),
                judge_id=f"{self._rubric.domain}_judge",
                judge_model="",
                evaluation_time_ms=eval_time_ms,
            )

        root_cause = RootCause.parse(primary.get("root_cause", "unknown"))
        severity = primary.get("severity", "medium")
        evidence = primary.get("evidence", [])
        if isinstance(evidence, str):
            evidence = [evidence]

        strategy = fix.get(
            "strategy",
            self._rubric.root_cause_to_strategy.get(root_cause.value, ""),
        )

        return StructuredDiagnosis(
            issue_domain=self._rubric.domain,
            issue_type=f"{self._rubric.domain}_{root_cause.value}",
            severity=severity,
            span_id=span_id,
            affected_region=primary.get("affected_region"),
            root_cause=root_cause,
            evidence=evidence,
            recommended_strategy=strategy,
            fix_instruction=fix.get("instruction", ""),
            fix_context=fix.get("context", {}),
            confidence=raw.get("confidence", 0.5),
            reasoning=raw.get("reasoning", ""),
            judge_id=f"{self._rubric.domain}_judge",
            judge_model="",
            evaluation_time_ms=eval_time_ms,
        )
