"""
Chart and table validation for llm-guard-kit.

ChartValidator — spec-aware judge prompt that detects wrong axis labels, unit
    mismatches, wrong chart type, and data/spec mismatches.  No new infrastructure:
    uses the same Haiku API key as AgentGuard.  Pass a chart spec dict (Vega-Lite,
    Chart.js, or plain description) and a data sample; the judge catches semantic
    errors that code-level validation misses.

Validated scope
---------------
No independent AUROC benchmark as of v0.54.0.  Judge output is heuristic — use
as a soft signal alongside unit tests and schema validation, not as a hard gate.

Usage
-----
    from llm_guard import ChartValidator

    validator = ChartValidator(api_key="sk-ant-...")
    result = validator.validate(
        question="Monthly revenue trend for 2024",
        chart_spec={
            "mark": "bar",
            "encoding": {
                "x": {"field": "month", "type": "ordinal"},
                "y": {"field": "revenue", "type": "quantitative"},
            },
        },
        data_sample=[
            {"month": "Jan", "revenue": 12000},
            {"month": "Feb", "revenue": 15400},
        ],
    )
    print(result.verdict, result.issues)
Purpose:    Detects wrong axis labels, unit mismatches, wrong chart type, and data/spec mismatches using a Haiku judge prompt. No independent AUROC benchmark; use as a soft signal alongside schema validation.
Exports:    ChartValidator, ChartValidationResult
Added:      v0.54.0
Status:     Experimental
"""

from __future__ import annotations

import json
import re
import warnings
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


# ── Judge prompt ───────────────────────────────────────────────────────────────

_CHART_JUDGE_SYSTEM = """\
You are a strict data visualization correctness evaluator.

Given:
- A natural-language question / intent
- A chart or table specification (Vega-Lite, Chart.js, or plain description)
- A data sample (first few rows)

Assess whether the visualization correctly answers the question.

Check specifically:
1. CHART TYPE — is the chart type appropriate for the data and question?
   (bar for comparison, line for trend, scatter for correlation, etc.)
2. AXES — are X and Y fields correct? Is the right field on each axis?
3. UNITS — do the axis labels, titles, and tick formats match the data units?
4. AGGREGATION — is data aggregated correctly (sum/avg/count as required)?
5. FILTERS — are any required filters applied?
6. COMPLETENESS — does the spec include all required fields from the question?

Output ONLY valid JSON:
{
  "diagnosis": "one sentence: what the chart does and why it may be wrong",
  "verdict": "VALID or INVALID or UNCERTAIN",
  "issues": ["specific", "issues", "found"],
  "confidence": 1-5
}

VALID     = chart correctly answers the question with the given data
INVALID   = clear error (wrong chart type, axes swapped, wrong unit, missing key field)
UNCERTAIN = ambiguous; spec is incomplete or cannot be evaluated from data sample alone"""

_CHART_JUDGE_USER_TMPL = """\
QUESTION / INTENT: {question}

CHART SPEC:
{spec}

DATA SAMPLE (first rows):
{data}"""

_VERDICT_RISK: Dict[str, float] = {"VALID": 0.10, "UNCERTAIN": 0.45, "INVALID": 0.85}

# ── Known chart type intent patterns ──────────────────────────────────────────
# Used by the heuristic fallback when no API key is available.
_TREND_RE   = re.compile(r"\b(trend|over time|time series|monthly|weekly|daily|year)\b", re.I)
_COMP_RE    = re.compile(r"\b(compar|rank|top|bottom|vs\.?|versus|between)\b", re.I)
_DIST_RE    = re.compile(r"\b(distribut|histogram|spread|range|outlier)\b", re.I)
_CORR_RE    = re.compile(r"\b(correlat|scatter|relationship|association)\b", re.I)

_TREND_MARKS    = {"line", "area"}
_COMP_MARKS     = {"bar", "column", "hbar"}
_DIST_MARKS     = {"histogram", "boxplot", "box"}
_CORR_MARKS     = {"point", "scatter", "circle"}


@dataclass
class ChartValidationResult:
    """Result of ChartValidator.validate()."""

    risk_score: float
    """0–1 semantic risk: 0.10=VALID, 0.45=UNCERTAIN, 0.85=INVALID."""

    verdict: str
    """'VALID' | 'INVALID' | 'UNCERTAIN'."""

    diagnosis: str
    """One-sentence judge explanation."""

    issues: List[str] = field(default_factory=list)
    """Specific issues identified (empty for VALID)."""

    judge_confidence: int = 3
    """Judge self-reported confidence 1–5."""

    judge_model: str = ""
    """Model used (or 'heuristic' for no-API fallback)."""


class ChartValidator:
    """
    Spec-aware chart and table validator via a custom Haiku judge prompt.

    Parameters
    ----------
    api_key : str, optional
        Anthropic API key. Falls back to ``ANTHROPIC_API_KEY`` env var.
    model : str
        Judge model. Default Haiku (~$0.0003/call).
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "claude-haiku-4-5-20251001",
    ) -> None:
        import os as _os
        self._api_key = api_key or _os.environ.get("ANTHROPIC_API_KEY", "")
        self._model = model
        self._client: Any = None
        if self._api_key:
            try:
                import anthropic as _ant
                self._client = _ant.Anthropic(api_key=self._api_key)
            except ImportError:
                warnings.warn(
                    "ChartValidator: anthropic package not installed. "
                    "Install with: pip install anthropic",
                    UserWarning, stacklevel=2,
                )

    def validate(
        self,
        question: str,
        chart_spec: Any,
        data_sample: Optional[List[Any]] = None,
        max_data_rows: int = 10,
    ) -> ChartValidationResult:
        """
        Validate a chart specification against the question and data sample.

        Parameters
        ----------
        question : str
            Natural-language intent / question the chart should answer.
        chart_spec : dict or str
            Chart specification.  Vega-Lite, Chart.js, or a plain text description
            are all accepted.  Dicts are serialised to JSON.
        data_sample : list, optional
            First N rows of the data the chart will render.  More rows give the
            judge more context; 5–10 is usually sufficient.
        max_data_rows : int
            Truncate data_sample to this many rows before sending to judge.

        Returns
        -------
        ChartValidationResult
        """
        spec_str = (
            json.dumps(chart_spec, indent=2)
            if isinstance(chart_spec, dict)
            else str(chart_spec)
        )
        data_str = _format_data(data_sample, max_data_rows)

        if not self._client:
            return _heuristic_chart_validate(question, chart_spec, data_sample)

        user_msg = _CHART_JUDGE_USER_TMPL.format(
            question=question,
            spec=spec_str[:2000],   # cap to avoid token explosion on large specs
            data=data_str,
        )

        try:
            resp = self._client.messages.create(
                model=self._model,
                max_tokens=512,
                system=_CHART_JUDGE_SYSTEM,
                messages=[{"role": "user", "content": user_msg}],
            )
            raw = resp.content[0].text.strip()
            parsed = _parse_json(raw)
            verdict = parsed.get("verdict", "UNCERTAIN").upper()
            if verdict not in _VERDICT_RISK:
                verdict = "UNCERTAIN"
            return ChartValidationResult(
                risk_score=_VERDICT_RISK[verdict],
                verdict=verdict,
                diagnosis=parsed.get("diagnosis", raw[:200]),
                issues=parsed.get("issues", []),
                judge_confidence=int(parsed.get("confidence", 3)),
                judge_model=self._model,
            )
        except Exception as exc:
            warnings.warn(
                f"ChartValidator: judge call failed ({exc}). "
                "Falling back to heuristic validation.",
                UserWarning, stacklevel=2,
            )
            return _heuristic_chart_validate(question, chart_spec, data_sample)

    def validate_batch(
        self,
        items: List[Dict[str, Any]],
    ) -> List[ChartValidationResult]:
        """
        Validate a batch of chart specs.

        Each item dict must have ``question`` and ``chart_spec`` keys.
        Optional: ``data_sample``.
        """
        return [
            self.validate(
                item["question"],
                item["chart_spec"],
                data_sample=item.get("data_sample"),
            )
            for item in items
        ]


# ── Helpers ────────────────────────────────────────────────────────────────────

def _format_data(data: Optional[List[Any]], max_rows: int) -> str:
    if not data:
        return "(no data sample provided)"
    rows = data[:max_rows]
    try:
        return json.dumps(rows, indent=2)[:1500]
    except (TypeError, ValueError):
        return str(rows)[:1500]


def _parse_json(raw: str) -> dict:
    cleaned = re.sub(r"```(?:json)?", "", raw).strip().rstrip("`").strip()
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        return {}


def _heuristic_chart_validate(
    question: str,
    chart_spec: Any,
    data_sample: Optional[List[Any]],
) -> ChartValidationResult:
    """
    Zero-cost heuristic fallback when no API key is available.

    Checks chart type intent alignment and obvious structural issues.
    """
    issues: List[str] = []

    # Extract mark / chart type
    mark = ""
    if isinstance(chart_spec, dict):
        mark = str(chart_spec.get("mark", chart_spec.get("type", ""))).lower()
    elif isinstance(chart_spec, str):
        for ctype in ("line", "bar", "scatter", "histogram", "area", "pie", "point"):
            if ctype in chart_spec.lower():
                mark = ctype
                break

    # Intent vs chart type alignment
    if mark:
        if _TREND_RE.search(question) and mark not in _TREND_MARKS:
            issues.append(f"Question suggests trend/time-series but chart type is '{mark}' (expected line or area)")
        if _CORR_RE.search(question) and mark not in _CORR_MARKS:
            issues.append(f"Question suggests correlation but chart type is '{mark}' (expected scatter/point)")
        if _DIST_RE.search(question) and mark not in _DIST_MARKS:
            issues.append(f"Question suggests distribution but chart type is '{mark}' (expected histogram or boxplot)")

    # Vega-Lite: check encoding fields
    if isinstance(chart_spec, dict) and "encoding" in chart_spec:
        enc = chart_spec["encoding"]
        if "x" not in enc and "y" not in enc:
            issues.append("Vega-Lite spec has encoding block but no x/y channels defined")
        for axis, info in enc.items():
            if isinstance(info, dict) and "field" not in info and "value" not in info:
                issues.append(f"Encoding channel '{axis}' has no 'field' or 'value'")

    # Data sample vs spec fields
    if data_sample and isinstance(chart_spec, dict) and "encoding" in chart_spec:
        sample_keys = set(data_sample[0].keys()) if data_sample else set()
        for axis, info in chart_spec["encoding"].items():
            if isinstance(info, dict):
                f = info.get("field", "")
                if f and sample_keys and f not in sample_keys:
                    issues.append(f"Field '{f}' in encoding not found in data sample columns {sorted(sample_keys)}")

    verdict = "INVALID" if len(issues) >= 2 else ("UNCERTAIN" if issues else "VALID")
    diagnosis = (
        "Heuristic check: " + "; ".join(issues)
        if issues
        else "No obvious structural issues detected (heuristic only — provide API key for semantic check)"
    )
    return ChartValidationResult(
        risk_score=_VERDICT_RISK[verdict],
        verdict=verdict,
        diagnosis=diagnosis,
        issues=issues,
        judge_model="heuristic",
    )
