import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import List, Dict, Any, Iterable

from zerophix.config import RedactionConfig
from zerophix.pipelines.redaction import RedactionPipeline
from zerophix.eval.metrics import Span, compute_span_metrics


ROOT = Path(__file__).resolve().parents[3]
TAB_DIR = ROOT / "data" / "benchmarks" / "tab"


@dataclass
class TabDoc:
    doc_id: str
    text: str
    gold_spans: List[Span]  # (doc_id, start, end, label)


def _load_tab_split(json_path: Path) -> List[TabDoc]:
    """
    Load a TAB JSON split (echr_train.json / echr_dev.json / echr_test.json).

    We treat any mention with identifier_type in {DIRECT, QUASI}
    as something that MUST be masked (i.e., PHI).
    """
    with json_path.open("r", encoding="utf-8") as f:
        docs_raw: List[Dict[str, Any]] = json.load(f)

    docs: List[TabDoc] = []

    for doc in docs_raw:
        doc_id = doc["doc_id"]
        text = doc["text"]
        annotations = doc["annotations"]

        spans: List[Span] = []
        
        for annot_group in annotations.values():
            # Each annotation group contains a list of entity mentions
            mentions = annot_group.get("entity_mentions", [])
            
            for ent in mentions:
                identifier_type = ent.get("identifier_type")
                
                if identifier_type not in ("DIRECT", "QUASI"):
                    continue  # not required to be masked

                start = int(ent["start_offset"])
                end = int(ent["end_offset"])
                label = ent.get("entity_type") or "PHI"

                spans.append((doc_id, start, end, label))

        docs.append(TabDoc(doc_id=doc_id, text=text, gold_spans=spans))

    return docs


def load_tab_corpus() -> Dict[str, List[TabDoc]]:
    """
    Load all splits if available. Requires `scripts/download_benchmarks.py`
    to have been run beforehand.
    """
    paths = {
        "train": TAB_DIR / "echr_train.json",
        "dev": TAB_DIR / "echr_dev.json",
        "test": TAB_DIR / "echr_test.json",
    }
    corpus: Dict[str, List[TabDoc]] = {}
    for split, p in paths.items():
        if not p.exists():
            raise FileNotFoundError(
                f"{p} not found. Run scripts/download_benchmarks.py first."
            )
        corpus[split] = _load_tab_split(p)
    return corpus


def _predict_spans(
    pipeline: RedactionPipeline, docs: Iterable[TabDoc]
) -> List[Span]:
    preds: List[Span] = []
    for d in docs:
        result = pipeline.redact(d.text)
        for ent in result.get("spans", []):
            start = int(ent.get("start", 0))
            end = int(ent.get("end", 0))
            label = ent.get("label") or ent.get("entity_type") or "PHI"
            preds.append((d.doc_id, start, end, label))
    return preds


def run_tab_benchmark(
    config: RedactionConfig | None = None,
    split: str = "test",
    ignore_label: bool = False,
):
    """
    Run ZeroPhi on the TAB benchmark and compute span-level metrics.

    Returns: (metrics, num_gold, num_pred)
    """
    corpus = load_tab_corpus()
    docs = corpus[split]

    if config is None:
        # simple, generic config
        config = RedactionConfig(
            country="EU",
            detectors=["regex", "spacy"],  # tweak as you like
            redaction_strategy="replace",
            min_confidence=0.5,
        )

    pipeline = RedactionPipeline.from_config(config)

    gold_spans: List[Span] = []
    for d in docs:
        gold_spans.extend(d.gold_spans)

    pred_spans = _predict_spans(pipeline, docs)

    metrics = compute_span_metrics(
        gold_spans=gold_spans, pred_spans=pred_spans, ignore_label=ignore_label
    )

    # --- Debug: Write mismatches to file ---
    debug_file = ROOT / "eval" / "results" / "tab_debug.txt"
    debug_file.parent.mkdir(parents=True, exist_ok=True)
    
    gold_set = set((s[0], s[1], s[2]) for s in gold_spans)
    pred_set = set((s[0], s[1], s[2]) for s in pred_spans)
    
    fp_list = list(pred_set - gold_set)[:20]
    fn_list = list(gold_set - pred_set)[:20]
    
    with debug_file.open("w", encoding="utf-8") as f:
        f.write(f"TAB Debug Info ({split})\n")
        f.write("=" * 40 + "\n")
        f.write(f"Metrics: {metrics.as_dict()}\n\n")
        
        f.write("Top 20 False Positives (Predicted but not in Gold):\n")
        for item in fp_list:
            # Find the text for this span
            doc = next((d for d in docs if d.doc_id == item[0]), None)
            text = doc.text[item[1]:item[2]] if doc else "???"
            f.write(f"  {item} -> '{text}'\n")
            
        f.write("\nTop 20 False Negatives (In Gold but not Predicted):\n")
        for item in fn_list:
            doc = next((d for d in docs if d.doc_id == item[0]), None)
            text = doc.text[item[1]:item[2]] if doc else "???"
            f.write(f"  {item} -> '{text}'\n")
            
    print(f"[TAB] Debug info written to {debug_file}")
    # ---------------------------------------

    return metrics, len(gold_spans), len(pred_spans)


if __name__ == "__main__":
    metrics, n_gold, n_pred = run_tab_benchmark()
    print(f"TAB test split: gold={n_gold}, pred={n_pred}")
    print(f"Precision: {metrics.precision:.3f}")
    print(f"Recall   : {metrics.recall:.3f}")
    print(f"F1       : {metrics.f1:.3f}")
