"""UD evaluation metrics implementation.

Uses the official CoNLL 2018 shared task evaluation script (conll18_ud_eval.py)
for computing metrics:
- Segmentation: Tokens, Sentences, Words
- Basic: UPOS, XPOS, UFeats, AllTags, Lemmas
- Dependency: UAS, LAS
- CoNLL 2018: CLAS, MLAS, BLEX

Reference: https://universaldependencies.org/conll18/evaluation.html
"""

import io
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from ..data.conllu_reader import Sentence, Token

from . import conll18_ud_eval

# All metric names returned by conll18_ud_eval.evaluate()
ALL_METRIC_NAMES = [
    "Tokens", "Sentences", "Words",
    "UPOS", "XPOS", "UFeats", "AllTags", "Lemmas",
    "UAS", "LAS", "CLAS", "MLAS", "BLEX",
]


@dataclass
class EvaluationResult:
    """Complete evaluation results with both F1 and aligned accuracy."""

    # F1 scores for all metrics (percentage, 0-100)
    f1: dict[str, float] = field(default_factory=dict)

    # Aligned accuracy for all metrics (percentage, 0-100; None for segmentation metrics)
    aligned_accuracy: dict[str, float | None] = field(default_factory=dict)

    # Counts
    total_words: int = 0
    total_sentences: int = 0

    # Model info
    treebank: str = ""
    model_name: str = ""
    base_model: str = ""

    # Full-text evaluation: which text source was used
    # "txt_file" | "text_metadata" | "concatenated_tokens" | ""
    text_source: str = ""

    def to_dict(self) -> dict[str, float]:
        """Convert to flat dictionary with F1 percentage values (backwards compatible)."""
        result = {}
        for name in ALL_METRIC_NAMES:
            if name in self.f1:
                result[name] = self.f1[name]
        result["total_words"] = self.total_words
        result["total_sentences"] = self.total_sentences
        return result

    def to_full_dict(self) -> dict[str, any]:
        """Convert to dictionary including model info."""
        result = self.to_dict()
        result.update({
            "Treebank": self.treebank,
            "Model_name": self.model_name,
            "Base_model": self.base_model,
        })
        return result


class MetricsCalculator:
    """Calculate all UD evaluation metrics using official conll18_ud_eval.py."""

    def __init__(self):
        self._gold_cache: dict[int, Any] = {}

    def _get_gold_ud(self, gold_sentences: list[Sentence]):
        """Get parsed gold UD data, caching by identity to avoid re-conversion."""
        cache_key = id(gold_sentences)
        if cache_key not in self._gold_cache:
            print(f"  Converting {len(gold_sentences)} gold sentences to CoNLL-U (first time, will be cached)...", flush=True)
            gold_conllu = self._sentences_to_conllu(gold_sentences)
            self._gold_cache[cache_key] = conll18_ud_eval.load_conllu(
                io.StringIO(gold_conllu)
            )
            print(f"  Gold CoNLL-U cached.", flush=True)
        else:
            print(f"  Using cached gold CoNLL-U ({len(gold_sentences)} sentences).", flush=True)
        return self._gold_cache[cache_key]

    def calculate(
        self,
        gold_sentences: list[Sentence],
        pred_sentences: list[Sentence],
        treebank: str = "",
        model_name: str = "",
        base_model: str = "",
    ) -> EvaluationResult:
        """Calculate all metrics using official CoNLL 2018 evaluation script."""
        gold_ud = self._get_gold_ud(gold_sentences)

        print(f"  Converting {len(pred_sentences)} predictions to CoNLL-U...", flush=True)
        pred_conllu = self._sentences_to_conllu(pred_sentences)
        pred_ud = conll18_ud_eval.load_conllu(io.StringIO(pred_conllu))

        print(f"  Running official CoNLL-18 evaluation...", flush=True)
        evaluation = conll18_ud_eval.evaluate(gold_ud, pred_ud)
        print(f"  Evaluation complete.", flush=True)

        return self._build_result(
            evaluation, gold_sentences, treebank, model_name, base_model
        )

    def calculate_from_conllu(
        self,
        gold_conllu_path: Path,
        system_conllu: str,
        treebank: str = "",
        model_name: str = "",
        base_model: str = "",
    ) -> EvaluationResult:
        """Calculate metrics comparing system CoNLL-U output against gold file.

        Used for full-text evaluation where tokenization may differ from gold.
        """
        gold_ud = conll18_ud_eval.load_conllu_file(str(gold_conllu_path))
        system_ud = conll18_ud_eval.load_conllu(io.StringIO(system_conllu))

        evaluation = conll18_ud_eval.evaluate(gold_ud, system_ud)

        f1_scores = {}
        aligned_scores = {}
        for name in ALL_METRIC_NAMES:
            score = evaluation[name]
            f1_scores[name] = score.f1 * 100
            aligned_scores[name] = (
                score.aligned_accuracy * 100
                if score.aligned_accuracy is not None
                else None
            )

        return EvaluationResult(
            f1=f1_scores,
            aligned_accuracy=aligned_scores,
            treebank=treebank,
            model_name=model_name,
            base_model=base_model,
        )

    def _build_result(
        self,
        evaluation: dict,
        gold_sentences: list[Sentence],
        treebank: str,
        model_name: str,
        base_model: str,
    ) -> EvaluationResult:
        """Build EvaluationResult from conll18_ud_eval output."""
        total_words = sum(len(sent.words) for sent in gold_sentences)
        total_sentences = len(gold_sentences)

        f1_scores = {}
        aligned_scores = {}
        for name in ALL_METRIC_NAMES:
            score = evaluation[name]
            f1_scores[name] = score.f1 * 100
            aligned_scores[name] = (
                score.aligned_accuracy * 100
                if score.aligned_accuracy is not None
                else None
            )

        return EvaluationResult(
            f1=f1_scores,
            aligned_accuracy=aligned_scores,
            total_words=total_words,
            total_sentences=total_sentences,
            treebank=treebank,
            model_name=model_name,
            base_model=base_model,
        )

    def _sentences_to_conllu(self, sentences: list[Sentence]) -> str:
        """Convert sentences to CoNLL-U format string."""
        parts = [sent.to_conllu() for sent in sentences]
        return "\n".join(parts) + "\n" if parts else ""
