"""
Timeline / Temporal Ordering Benchmark for RNSR

Two-tier benchmark:

Tier 1 -- Synthetic Ground Truth (self-contained, zero external deps)
    Uses the test PDFs generated by ``scripts/generate_test_pdfs.py`` with
    hardcoded expected events.  Validates the full pipeline: ingestion ->
    skeleton -> KG -> extract_timeline.

Tier 2 -- LexTime (external academic dataset)
    514 instances from U.S. Federal Complaints with annotated temporal
    relation descriptions and entailment/contradiction labels.
    Dataset: https://zenodo.org/records/17157439  (CC-BY-4.0)
    Paper:  "LexTime: A Benchmark for Temporal Ordering of Legal Events"
            (EMNLP 2025 Findings)

Metrics:
    - event_recall:  fraction of expected events that were extracted
    - order_accuracy: fraction of consecutive event pairs in correct order
    - date_parse_rate: fraction of events with a successfully parsed date
    - lextime_accuracy: classification accuracy on LexTime (Tier 2 only)

Usage:
    from rnsr.benchmarks.timeline_bench import TimelineBenchLoader
    dataset = TimelineBenchLoader.load(max_samples=10)

    # Or run the full evaluation directly:
    from rnsr.benchmarks.timeline_bench import run_timeline_benchmark
    results = run_timeline_benchmark()
"""

from __future__ import annotations

import csv
import re
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

import structlog

from rnsr.benchmarks.standard_benchmarks import BenchmarkDataset, BenchmarkQuestion

logger = structlog.get_logger(__name__)

DATA_DIR = Path(__file__).resolve().parent / "data" / "lextime"
TEST_DOCS_DIR = Path(__file__).resolve().parent.parent.parent / "test-documents"


# ============================================================================
# Ground-truth events for Tier 1 (synthetic PDFs)
# ============================================================================

@dataclass
class ExpectedEvent:
    """An event we expect the timeline extractor to find."""
    date_str: str          # e.g. "15 January 2019"
    keywords: list[str]    # at least one must appear in the extracted description
    year: int
    month: int
    day: int


MERIDIAN_EVENTS: list[ExpectedEvent] = [
    # Each event has 3+ keywords with at least 2 being distinctive to that event.
    # Unique disambiguators are marked with comments where overlap exists.
    ExpectedEvent("15 January 2019", ["feasibility", "council", "approved", "study"], 2019, 1, 15),
    ExpectedEvent("3 March 2019", ["hargrove", "awarded", "2.4 million", "contract"], 2019, 3, 3),
    ExpectedEvent("28 August 2019", ["feasibility", "report", "submitted", "hargrove"], 2019, 8, 28),
    ExpectedEvent("15 October 2019", ["environmental", "lodged", "department", "eia"], 2019, 10, 15),
    ExpectedEvent("22 February 2020", ["conditional", "approval", "fauna", "eia"], 2020, 2, 22),
    ExpectedEvent("1 April 2020", ["public", "consultation", "submissions"], 2020, 4, 1),
    ExpectedEvent("14 September 2020", ["planning", "application", "state planning"], 2020, 9, 14),
    ExpectedEvent("5 March 2021", ["planning", "approval", "granted", "23 conditions"], 2021, 3, 5),
    ExpectedEvent("1 May 2021", ["detailed", "design", "bridgepoint"], 2021, 5, 1),
    ExpectedEvent("10 November 2021", ["tender", "issued", "closing"], 2021, 11, 10),
    ExpectedEvent("18 May 2022", ["pacific", "alliance", "378"], 2022, 5, 18),
    ExpectedEvent("1 August 2022", ["construction", "commenced", "earthworks", "site clearing"], 2022, 8, 1),
    ExpectedEvent("12 December 2022", ["contaminated", "soil", "dawson"], 2022, 12, 12),
    ExpectedEvent("20 June 2023", ["willow creek", "bridge", "structural"], 2023, 6, 20),
    ExpectedEvent("10 April 2024", ["fauna", "corridor", "km 4.2"], 2024, 4, 10),
]

BAXTER_EVENTS: list[ExpectedEvent] = [
    # Each event has 2+ distinctive keywords.
    ExpectedEvent("14 February 2017", ["employed", "baxter", "senior process"], 2017, 2, 14),
    ExpectedEvent("3 July 2022", ["complaint", "unsafe", "formal"], 2022, 7, 3),
    ExpectedEvent("18 July 2022", ["investigation", "commenced", "internal"], 2022, 7, 18),
    ExpectedEvent("15 September 2022", ["warning", "insubordination"], 2022, 9, 15),
    ExpectedEvent("30 November 2022", ["terminated", "termination", "employment", "dismissed", "performance"], 2022, 11, 30),
    ExpectedEvent("14 December 2022", ["unfair", "dismissal", "application", "fair work"], 2022, 12, 14),
    ExpectedEvent("2 February 2023", ["conciliation", "conference", "settlement"], 2023, 2, 2),
    ExpectedEvent("28 July 2023", ["decision", "handed down", "reinstatement"], 2023, 7, 28),
    ExpectedEvent("25 August 2023", ["appeal", "filed", "thornton"], 2023, 8, 25),
    ExpectedEvent("19 January 2024", ["appeal", "dismissed", "full bench"], 2024, 1, 19),
    ExpectedEvent("5 February 2024", ["reinstated", "baxter"], 2024, 2, 5),
]


# ============================================================================
# Tier 1: Synthetic PDF evaluation
# ============================================================================

@dataclass
class TimelineScore:
    """Scores for a single timeline extraction run."""
    doc_name: str
    total_expected: int
    events_found: int
    event_recall: float
    order_accuracy: float
    date_parse_rate: float
    details: list[dict[str, Any]] = field(default_factory=list)


def _keyword_match_count(extracted_desc: str, expected: ExpectedEvent) -> int:
    """Count how many of the expected keywords appear in the extracted text."""
    desc_lower = extracted_desc.lower()
    return sum(1 for kw in expected.keywords if kw.lower() in desc_lower)


def _event_matches(extracted_desc: str, expected: ExpectedEvent, min_keywords: int = 2) -> bool:
    """Check if an extracted event description matches an expected event.

    Requires at least ``min_keywords`` keywords to appear in the description.
    This prevents false matches from generic words like "construction" or
    "appeal" that appear in multiple events.
    """
    return _keyword_match_count(extracted_desc, expected) >= min_keywords


def _date_matches(parsed: datetime | None, expected: ExpectedEvent) -> bool:
    """Check if a parsed date matches the expected date."""
    if parsed is None:
        return False
    return (
        parsed.year == expected.year
        and parsed.month == expected.month
        and parsed.day == expected.day
    )


def _build_combined_text(evt: Any) -> str:
    """Build a searchable text string from an extracted event."""
    desc = getattr(evt, "description", "") or ""
    entities = " ".join(getattr(evt, "entities_involved", []) or [])
    return f"{desc} {entities}"


def evaluate_timeline(
    extracted_events: list[Any],  # list of TimelineEvent
    expected_events: list[ExpectedEvent],
) -> TimelineScore:
    """Score extracted timeline events against ground truth.

    Uses date-aware best-match selection with one-to-one assignment:
    1. For each expected event, collect ALL candidate matches (2+ keywords).
    2. Prefer the candidate whose parsed date matches exactly.
    3. Fall back to the candidate with the most keyword matches.
    4. Once an extracted event is matched, it cannot be reused.

    Args:
        extracted_events: Output from ``extract_timeline()``.
        expected_events: Known ground-truth events.

    Returns:
        TimelineScore with recall, ordering, and date-parse metrics.
    """
    matched: list[tuple[int, Any]] = []  # (expected_idx, extracted_event)
    unmatched_expected: list[int] = []
    details: list[dict[str, Any]] = []
    used_extracted: set[int] = set()  # indices of extracted events already matched

    for i, exp in enumerate(expected_events):
        # Collect all candidate matches (with 2+ keywords, not already used)
        candidates: list[tuple[int, int, bool]] = []  # (ext_idx, kw_count, date_ok)
        for j, evt in enumerate(extracted_events):
            if j in used_extracted:
                continue
            combined = _build_combined_text(evt)
            kw_count = _keyword_match_count(combined, exp)
            if kw_count >= 2:
                date_ok = _date_matches(getattr(evt, "date_parsed", None), exp)
                candidates.append((j, kw_count, date_ok))

        if candidates:
            # Sort: prefer date match first, then highest keyword count
            candidates.sort(key=lambda c: (c[2], c[1]), reverse=True)
            best_idx, best_kw, best_date = candidates[0]
            best_evt = extracted_events[best_idx]
            used_extracted.add(best_idx)

            matched.append((i, best_evt))
            details.append({
                "expected": exp.date_str,
                "found": True,
                "date_correct": best_date,
                "kw_matched": best_kw,
            })
        else:
            unmatched_expected.append(i)
            details.append({
                "expected": exp.date_str,
                "found": False,
                "keywords": exp.keywords,
            })

    # Event recall
    event_recall = len(matched) / len(expected_events) if expected_events else 0.0

    # Date parse rate (among matched events)
    dates_parsed = sum(
        1 for _, evt in matched
        if getattr(evt, "date_parsed", None) is not None
    )
    date_parse_rate = dates_parsed / len(matched) if matched else 0.0

    # Order accuracy: check that matched events appear in chronological order
    # by comparing their positions in the extracted list
    if len(matched) >= 2:
        extracted_positions = []
        for _, evt in matched:
            try:
                pos = extracted_events.index(evt)
            except ValueError:
                pos = -1
            extracted_positions.append(pos)

        correct_pairs = 0
        total_pairs = 0
        for a in range(len(extracted_positions)):
            for b in range(a + 1, len(extracted_positions)):
                total_pairs += 1
                if extracted_positions[a] < extracted_positions[b]:
                    correct_pairs += 1
        order_accuracy = correct_pairs / total_pairs if total_pairs else 1.0
    else:
        order_accuracy = 1.0

    return TimelineScore(
        doc_name="",
        total_expected=len(expected_events),
        events_found=len(matched),
        event_recall=event_recall,
        order_accuracy=order_accuracy,
        date_parse_rate=date_parse_rate,
        details=details,
    )


def run_tier1_benchmark(
    max_docs: int | None = None,
) -> list[TimelineScore]:
    """Run the Tier 1 synthetic-PDF timeline benchmark.

    Ingests the test PDFs, extracts timelines, and scores against ground truth.

    Returns:
        List of TimelineScore (one per document).
    """
    from rnsr import ingest_document, build_skeleton_index
    from rnsr.indexing.knowledge_graph import InMemoryKnowledgeGraph
    from rnsr.extraction.rlm_unified_extractor import extract_entities_and_relationships
    from rnsr.extraction.timeline_extractor import extract_timeline

    test_cases: list[tuple[str, list[ExpectedEvent]]] = [
        ("Timeline - Meridian Project History.pdf", MERIDIAN_EVENTS),
        ("Timeline - Baxter v Thornton.pdf", BAXTER_EVENTS),
    ]

    if max_docs:
        test_cases = test_cases[:max_docs]

    scores: list[TimelineScore] = []

    for pdf_name, expected in test_cases:
        pdf_path = TEST_DOCS_DIR / pdf_name
        if not pdf_path.exists():
            logger.warning("timeline_bench_pdf_missing", path=str(pdf_path))
            continue

        logger.info("timeline_bench_processing", doc=pdf_name)

        # Ingest
        result = ingest_document(str(pdf_path))
        skeleton, kv_store = build_skeleton_index(result.tree)

        # Build KG
        kg = InMemoryKnowledgeGraph()
        for node_id, node in skeleton.items():
            content = kv_store.get(node_id) or ""
            if len(content.strip()) < 20:
                continue
            try:
                ext_result = extract_entities_and_relationships(
                    node_id=node_id,
                    doc_id=pdf_name,
                    header=node.header,
                    content=content,
                )
                for entity in ext_result.entities:
                    kg.add_entity(entity)
                for rel in ext_result.relationships:
                    kg.add_relationship(rel)
            except Exception as exc:
                logger.debug("extraction_error", node=node_id, error=str(exc))

        # Extract timeline (KG + direct text scan for higher recall)
        events = extract_timeline(kg, skeleton=skeleton, kv_store=kv_store)

        # Score
        score = evaluate_timeline(events, expected)
        score.doc_name = pdf_name
        scores.append(score)

        logger.info(
            "timeline_bench_result",
            doc=pdf_name,
            recall=f"{score.event_recall:.0%}",
            order=f"{score.order_accuracy:.0%}",
            date_parse=f"{score.date_parse_rate:.0%}",
            found=f"{score.events_found}/{score.total_expected}",
        )

    return scores


# ============================================================================
# Tier 2: LexTime evaluation
# ============================================================================


def _load_lextime_csv() -> list[dict[str, str]] | None:
    """Load the LexTime CSV if it exists."""
    csv_path = DATA_DIR / "Entailment_Dataset.csv"
    if not csv_path.exists():
        logger.info("lextime_csv_not_found", path=str(csv_path))
        return None
    with open(csv_path, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def run_lextime_benchmark(
    llm_fn: Any | None = None,
    max_samples: int | None = None,
) -> dict[str, Any] | None:
    """Run the LexTime temporal-ordering benchmark.

    Each instance has a paragraph, a temporal relation (TR) description,
    and a label (entailment or contradiction).  We ask the LLM to classify
    whether the TR is entailed or contradicted by the paragraph.

    Args:
        llm_fn: A ``Callable[[str], str]`` for classification. If None,
            attempts to create one via ``get_llm()``.
        max_samples: Limit the number of instances evaluated.

    Returns:
        Dict with accuracy, precision, recall, f1, and per-label breakdowns.
        Returns None if the dataset is not available.
    """
    rows = _load_lextime_csv()
    if rows is None:
        return None

    if max_samples:
        rows = rows[:max_samples]

    # Get or create LLM
    if llm_fn is None:
        try:
            from rnsr.llm import get_llm
            _llm = get_llm()
            llm_fn = lambda prompt: str(_llm.complete(prompt))
        except Exception as exc:
            logger.error("lextime_llm_init_failed", error=str(exc))
            return None

    correct = 0
    total = 0
    tp = {"entailment": 0, "contradiction": 0}
    fp = {"entailment": 0, "contradiction": 0}
    fn = {"entailment": 0, "contradiction": 0}

    for row in rows:
        paragraph = row.get("paragraph", "")
        tr = row.get("TR", "")
        gold_label = row.get("label", "").strip().lower()

        if gold_label not in ("entailment", "contradiction"):
            continue

        prompt = (
            "You are a temporal reasoning expert analysing legal documents.\n\n"
            "Given the following paragraph from a legal complaint, determine "
            "whether the temporal relation statement is ENTAILED or CONTRADICTED "
            "by the paragraph.\n\n"
            f"PARAGRAPH:\n{paragraph}\n\n"
            f"TEMPORAL RELATION:\n{tr}\n\n"
            "Answer with EXACTLY one word: entailment or contradiction"
        )

        try:
            raw = llm_fn(prompt).strip().lower()
            # Extract the classification
            if "contradiction" in raw:
                predicted = "contradiction"
            elif "entailment" in raw:
                predicted = "entailment"
            else:
                predicted = "unknown"
        except Exception:
            predicted = "unknown"

        total += 1
        if predicted == gold_label:
            correct += 1
            tp[gold_label] += 1
        else:
            if predicted in fp:
                fp[predicted] += 1
            fn[gold_label] += 1

    accuracy = correct / total if total else 0.0

    results: dict[str, Any] = {
        "total": total,
        "correct": correct,
        "accuracy": accuracy,
        "per_label": {},
    }

    for label in ("entailment", "contradiction"):
        precision = tp[label] / (tp[label] + fp[label]) if (tp[label] + fp[label]) else 0.0
        recall = tp[label] / (tp[label] + fn[label]) if (tp[label] + fn[label]) else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
        results["per_label"][label] = {
            "precision": precision,
            "recall": recall,
            "f1": f1,
        }

    logger.info(
        "lextime_benchmark_result",
        accuracy=f"{accuracy:.1%}",
        total=total,
        correct=correct,
    )

    return results


# ============================================================================
# Loader (for integration with run_all_benchmarks.py)
# ============================================================================


class TimelineBenchLoader:
    """Loader for the timeline / temporal ordering benchmark.

    Returns a ``BenchmarkDataset`` of LexTime instances (Tier 2).
    Tier 1 (synthetic PDF) evaluation is run via ``run_tier1_benchmark()``.
    """

    @staticmethod
    def load(
        split: str = "test",
        max_samples: Optional[int] = None,
    ) -> BenchmarkDataset:
        """Load LexTime as a BenchmarkDataset.

        Each instance becomes a BenchmarkQuestion where:
        - question = temporal relation (TR)
        - answer = gold label (entailment / contradiction)
        - context = [paragraph]
        """
        rows = _load_lextime_csv()
        questions: list[BenchmarkQuestion] = []

        if rows is None:
            logger.warning("lextime_not_available_returning_empty_dataset")
            return BenchmarkDataset(
                name="timeline",
                description="LexTime: Temporal ordering of legal events (dataset not found)",
                questions=[],
                metrics=["accuracy"],
                source_url="https://zenodo.org/records/17157439",
            )

        for i, row in enumerate(rows):
            if max_samples and i >= max_samples:
                break
            label = row.get("label", "").strip().lower()
            if label not in ("entailment", "contradiction"):
                continue
            questions.append(
                BenchmarkQuestion(
                    id=f"lextime_{i:04d}",
                    question=row.get("TR", ""),
                    answer=label,
                    context=[row.get("paragraph", "")],
                    reasoning_type="temporal",
                    metadata={
                        "pair_type": row.get("pair_type", ""),
                        "source": "LexTime",
                    },
                )
            )

        return BenchmarkDataset(
            name="timeline",
            description=(
                "LexTime: Temporal ordering of legal events. "
                "514 instances from U.S. Federal Complaints. "
                "Task: classify temporal relation as entailment or contradiction."
            ),
            questions=questions,
            metrics=["accuracy"],
            source_url="https://zenodo.org/records/17157439",
        )


# ============================================================================
# Combined runner
# ============================================================================


def run_timeline_benchmark(
    max_samples: int | None = None,
    skip_tier1: bool = False,
    skip_tier2: bool = False,
) -> dict[str, Any]:
    """Run the full timeline benchmark (both tiers).

    Returns:
        Dict with ``tier1`` (list of TimelineScore dicts) and
        ``tier2`` (LexTime results dict or None).
    """
    results: dict[str, Any] = {"tier1": None, "tier2": None}

    if not skip_tier1:
        scores = run_tier1_benchmark(max_docs=max_samples)
        results["tier1"] = [
            {
                "doc": s.doc_name,
                "event_recall": s.event_recall,
                "order_accuracy": s.order_accuracy,
                "date_parse_rate": s.date_parse_rate,
                "found": s.events_found,
                "expected": s.total_expected,
            }
            for s in scores
        ]

    if not skip_tier2:
        results["tier2"] = run_lextime_benchmark(max_samples=max_samples)

    return results
