"""Evaluation orchestrator for COMBO."""

from functools import partial
from pathlib import Path
from typing import Any, Optional

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import PreTrainedTokenizer

from ..data.conllu_reader import Document, Sentence, Token
from ..data.dataset import UDDataset, collate_fn
from ..data.label_encoder import MorphoEncoder, CharVocab, LabelEncoder
from ..inference.mst import decode_mst
from ..model import UDParser
from ..segmenters import LamboSegmenter
from .metrics import EvaluationResult, MetricsCalculator


class Evaluator:
    """Evaluation orchestrator for COMBO."""

    def __init__(
        self,
        morpho_encoder: MorphoEncoder,
        deprel_encoder: LabelEncoder,
        char_vocab: CharVocab,
        treebank: str = "",
        model_name: str = "",
        base_model: str = "",
        use_lemma_beam_search: bool = False,
        lemma_beam_width: int = 5,
        use_morpho_constraints: bool = False,
    ):
        self.morpho_encoder = morpho_encoder
        self.deprel_encoder = deprel_encoder
        self.char_vocab = char_vocab
        self.metrics_calculator = MetricsCalculator()
        self.treebank = treebank
        self.model_name = model_name
        self.base_model = base_model
        self.use_lemma_beam_search = use_lemma_beam_search
        self.lemma_beam_width = lemma_beam_width
        self.use_morpho_constraints = use_morpho_constraints
        self._gold_cache: dict[int, list[Sentence]] = {}
        self._model_ref: Optional[UDParser] = None

    def evaluate(
        self,
        model: UDParser,
        dataloader: DataLoader,
        device: torch.device,
        max_samples: Optional[int] = None,
        show_progress: bool = False,
    ) -> EvaluationResult:
        """Run evaluation on dataset.

        Returns:
            EvaluationResult with F1 and aligned accuracy for all metrics.
        """
        model.eval()
        self._model_ref = model

        # Cache gold sentences per dataloader (they never change)
        dl_key = id(dataloader)
        cache_gold = dl_key in self._gold_cache and max_samples is None
        all_gold_sentences = self._gold_cache.get(dl_key)
        collect_gold = all_gold_sentences is None
        if collect_gold:
            all_gold_sentences = []

        all_pred_sentences: list[Sentence] = []
        samples_seen = 0

        iterator = dataloader
        if show_progress:
            total = None
            if max_samples:
                total = (max_samples + dataloader.batch_size - 1) // dataloader.batch_size
            iterator = tqdm(dataloader, desc="Evaluating", total=total)

        with torch.no_grad():
            for batch in iterator:
                if max_samples and samples_seen >= max_samples:
                    break

                # Move batch to device (exclude non-tensor items)
                input_batch = self._prepare_batch(batch, device)

                # Forward pass
                outputs = model(**input_batch)

                # Decode predictions
                predictions = self._decode_batch(outputs, batch, device)

                if collect_gold:
                    all_gold_sentences.extend(batch["sentences"])
                all_pred_sentences.extend(predictions)
                samples_seen += len(batch["sentences"])

        # Cache gold sentences (skip when using max_samples as subset varies)
        if collect_gold and max_samples is None:
            self._gold_cache[dl_key] = all_gold_sentences

        # Calculate metrics
        result = self.metrics_calculator.calculate(
            all_gold_sentences,
            all_pred_sentences,
            treebank=self.treebank,
            model_name=self.model_name,
            base_model=self.base_model,
        )

        return result

    def evaluate_and_predict(
        self,
        model: UDParser,
        dataloader: DataLoader,
        device: torch.device,
        show_progress: bool = True,
    ) -> tuple[EvaluationResult, list[Sentence]]:
        """Run evaluation and return both metrics and predictions."""
        model.eval()
        self._model_ref = model

        all_gold_sentences: list[Sentence] = []
        all_pred_sentences: list[Sentence] = []

        iterator = dataloader
        if show_progress:
            iterator = tqdm(dataloader, desc="Evaluating")

        with torch.no_grad():
            for batch in iterator:
                gold_sentences = batch["sentences"]
                input_batch = self._prepare_batch(batch, device)
                outputs = model(**input_batch)
                predictions = self._decode_batch(outputs, batch, device)

                all_gold_sentences.extend(gold_sentences)
                all_pred_sentences.extend(predictions)

        result = self.metrics_calculator.calculate(
            all_gold_sentences,
            all_pred_sentences,
            treebank=self.treebank,
            model_name=self.model_name,
            base_model=self.base_model,
        )

        return result, all_pred_sentences

    @staticmethod
    def _extract_text_from_conllu(conllu_path: Path) -> str:
        """Extract raw text from gold CoNLL-U file using # text = metadata.

        This guarantees the text matches the gold annotations character-for-character,
        which is required by conll18_ud_eval.py.
        """
        texts = []
        with open(conllu_path, encoding="utf-8") as f:
            for line in f:
                if line.startswith("# text = "):
                    texts.append(line[len("# text = "):].rstrip("\n"))
        return "\n".join(texts)

    @staticmethod
    def _concatenate_tokens_from_conllu(conllu_path: Path) -> str:
        """Build raw text by joining token forms from gold CoNLL-U with spaces.

        Last-resort fallback when # text = metadata is missing.  The resulting
        text is character-compatible with gold tokens because conll18_ud_eval
        uses MWT forms (e.g. "1-2") and skips their sub-tokens, and so do we.
        """
        sentences = []
        current_tokens: list[str] = []
        mwt_end = 0  # end index of current multiword token range
        with open(conllu_path, encoding="utf-8") as f:
            for line in f:
                line = line.rstrip("\n")
                if not line:
                    if current_tokens:
                        sentences.append(" ".join(current_tokens))
                        current_tokens = []
                    mwt_end = 0
                    continue
                if line.startswith("#"):
                    continue
                fields = line.split("\t")
                if len(fields) < 2:
                    continue
                token_id = fields[0]
                # Skip empty word nodes (e.g. "1.1")
                if "." in token_id:
                    continue
                # Multiword token (e.g. "1-2"): use its form, skip sub-tokens
                if "-" in token_id:
                    start, end = token_id.split("-")
                    mwt_end = int(end)
                    current_tokens.append(fields[1])
                    continue
                # Skip sub-tokens covered by a MWT
                if int(token_id) <= mwt_end:
                    continue
                current_tokens.append(fields[1])
        if current_tokens:
            sentences.append(" ".join(current_tokens))
        return "\n".join(sentences)

    def evaluate_full_text(
        self,
        model: UDParser,
        text_path: Path,
        gold_conllu_path: Path,
        tokenizer: PreTrainedTokenizer,
        max_length: int,
        batch_size: int,
        device: torch.device,
        language: str = "Polish",
        segmenter=None,
        combo_model_id: str | None = None,
        show_progress: bool = True,
        split_level: str | None = None,
        split_multiwords: bool | None = None,
    ) -> EvaluationResult:
        """Run full-text evaluation: segment raw text, predict, compare against gold.

        This evaluates both segmentation quality (Tokens, Sentences, Words)
        and parsing quality (UPOS, LAS, etc.) on text that was segmented
        by the system rather than using gold tokenization.

        Args:
            model: UDParser model.
            text_path: Path to raw text file (.txt) for segmenter input.
            gold_conllu_path: Path to gold CoNLL-U file for evaluation.
            tokenizer: HuggingFace tokenizer for creating dataset.
            max_length: Maximum sequence length.
            batch_size: Batch size for evaluation.
            device: Device for inference.
            language: Language name for LamboSegmenter (default: "Polish").
            segmenter: Segmenter to use (default: LamboSegmenter).
            show_progress: Show progress bar.

        Returns:
            EvaluationResult with segmentation + parsing metrics.
        """
        model.eval()
        self._model_ref = model

        if segmenter is None:
            lambo_model = None
            if combo_model_id:
                from ..hub import resolve_lambo_model
                lambo_model = resolve_lambo_model(combo_model_id)
            segmenter = LamboSegmenter(language, model_name=lambo_model)

        # 1. Resolve raw text from gold CoNLL-U (guarantees character alignment).
        #    # text = metadata → concatenated tokens
        raw_text = self._extract_text_from_conllu(gold_conllu_path)
        if raw_text.strip():
            text_source = "text_metadata"
            print(f"  Full-text source: # text = metadata from {gold_conllu_path.name}")
        else:
            raw_text = self._concatenate_tokens_from_conllu(gold_conllu_path)
            if raw_text.strip():
                text_source = "concatenated_tokens"
                print(f"  Full-text source: concatenated tokens from {gold_conllu_path.name} "
                      f"(no # text = metadata)")
            else:
                print(f"  Warning: no text source available for {gold_conllu_path.name}, "
                      f"skipping full-text evaluation")
                return EvaluationResult()
        doc = Document(text=raw_text, sentences=[])
        seg_kwargs = {}
        if split_level is not None:
            seg_kwargs["split_level"] = split_level
        if split_multiwords is not None:
            seg_kwargs["split_multiwords"] = split_multiwords
        doc = segmenter(doc, **seg_kwargs)

        # Clean up segmenter output: remove whitespace-only tokens (e.g. \n)
        # and strip newlines from sentence text (they break CoNLL-U format)
        for sent in doc.sentences:
            sent.tokens = [t for t in sent.tokens if t.form.strip()]
            for i, t in enumerate(sent.tokens):
                t.id = str(i + 1)  # re-number after filtering
            if sent.text:
                sent.text = sent.text.replace("\n", " ").strip()
        doc.sentences = [s for s in doc.sentences if s.tokens]

        if not doc.sentences:
            print(f"  Warning: No sentences after segmentation of {text_path}")
            return EvaluationResult()

        # 2. Create dataset and dataloader from segmented text
        dataset = UDDataset(
            [doc], tokenizer, self.morpho_encoder, self.deprel_encoder,
            self.char_vocab, max_length,
        )
        collate = partial(
            collate_fn,
            pad_token_id=tokenizer.pad_token_id,
            char_pad_id=self.char_vocab.pad_id,
        )
        dataloader = DataLoader(
            dataset, batch_size=batch_size, shuffle=False,
            num_workers=0, collate_fn=collate,
        )

        # 3. Run predictions
        all_pred_sentences: list[Sentence] = []
        iterator = tqdm(dataloader, desc="Full-text eval") if show_progress else dataloader

        with torch.no_grad():
            for batch in iterator:
                input_batch = self._prepare_batch(batch, device)
                outputs = model(**input_batch)
                predictions = self._decode_batch(outputs, batch, device)
                all_pred_sentences.extend(predictions)

        # Report skipped sentence stats
        if dataset.truncated_count > 0:
            print(f"\n  Note: {dataset.truncated_count} sentences were skipped during dataset loading "
                  f"(exceeded max_length={max_length})")

        # 4. Validate predicted trees — fix any that don't have exactly one root
        for sent in all_pred_sentences:
            roots = [t for t in sent.tokens if t.head == 0]
            if len(roots) > 1:
                best_root = roots[0]
                for t in roots[1:]:
                    t.head = int(best_root.id)
            elif len(roots) == 0 and sent.tokens:
                sent.tokens[0].head = 0

        pred_conllu = self.metrics_calculator._sentences_to_conllu(all_pred_sentences)
        try:
            result = self.metrics_calculator.calculate_from_conllu(
                gold_conllu_path, pred_conllu,
                treebank=self.treebank,
                model_name=self.model_name,
                base_model=self.base_model,
            )
        except Exception as e:
            print(f"\n  Warning: full-text evaluation failed: {e}")
            if dataset.truncated_count > 0:
                print(f"  ({dataset.truncated_count} sentences were skipped due to "
                      f"max_length={max_length})")
            return EvaluationResult()
        result.text_source = text_source

        return result

    def _prepare_batch(
        self, batch: dict[str, Any], device: torch.device
    ) -> dict[str, Tensor]:
        """Prepare batch for model input."""
        keys_to_include = [
            "input_ids",
            "attention_mask",
            "word_start_mask",
            "word_mask",
            "form_chars",
        ]
        return {
            k: batch[k].to(device) for k in keys_to_include if k in batch
        }

    def _decode_batch(
        self,
        outputs: dict[str, Tensor],
        batch: dict[str, Any],
        device: torch.device,
    ) -> list[Sentence]:
        """Decode model outputs to Sentence objects."""
        predictions = []

        upos_preds = outputs["upos_logits"].argmax(dim=-1)  # [batch, seq]
        xpos_preds = outputs["xpos_logits"].argmax(dim=-1)  # [batch, seq]
        feats_logits = outputs["feats_logits"]  # dict[str, [batch, seq, vocab]]
        feats_preds = {
            cat: logits.argmax(dim=-1)
            for cat, logits in feats_logits.items()
        }

        arc_logits = outputs["arc_logits"]  # [batch, seq, 1+seq]
        rel_logits = outputs["rel_logits"]  # [batch, seq, 1+seq, n_rels]

        # Decode lemmas
        lemma_preds = None
        if outputs["lemma_logits"] is not None:
            if self.use_lemma_beam_search and self._model_ref is not None:
                word_repr_for_lemma = outputs["word_repr_for_lemma"]
                form_chars = batch["form_chars"].to(device)
                lemma_preds = self._model_ref.lemma_head.beam_decode(
                    word_repr_for_lemma, form_chars,
                    beam_width=self.lemma_beam_width,
                )
            else:
                lemma_preds = outputs["lemma_logits"].argmax(dim=-1)  # [batch, seq, max_len]

        batch_size = upos_preds.size(0)

        for b in range(batch_size):
            sentence = batch["sentences"][b]
            words = sentence.tokens
            n_words = len(words)

            # Decode heads with MST (Chu-Liu-Edmonds) to guarantee a valid tree
            arc_scores = arc_logits[b, :n_words, :n_words+1].cpu().numpy()
            energy = np.full((n_words + 1, n_words + 1), -1e9)
            energy[:, 1:n_words+1] = arc_scores.T
            heads, _ = decode_mst(energy, n_words + 1, has_labels=False)
            arc_preds_list = heads[1:n_words+1].tolist()

            # Enforce valid tree: exactly one root, no self-loops
            # Fix self-loops first (word pointing to itself)
            for i in range(n_words):
                if arc_preds_list[i] == i + 1:  # 1-indexed: word i+1 points to head i+1
                    arc_preds_list[i] = 0  # temporarily make it root, will be fixed below

            # Enforce single root
            root_indices = [i for i, h in enumerate(arc_preds_list) if h == 0]
            if len(root_indices) == 0:
                # No root — pick word with highest root score
                best_root = max(range(n_words), key=lambda i: arc_scores[i, 0])
                arc_preds_list[best_root] = 0
            elif len(root_indices) > 1:
                best_root = max(root_indices, key=lambda i: arc_scores[i, 0])
                best_root_1indexed = best_root + 1
                for i in root_indices:
                    if i != best_root:
                        arc_preds_list[i] = best_root_1indexed

            # Build predicted word tokens
            pred_words = []
            for i, word in enumerate(words):
                # Decode morphological labels
                upos = self.morpho_encoder.decode_upos(upos_preds[b, i].item())
                xpos = self.morpho_encoder.decode_xpos(xpos_preds[b, i].item())
                feat_ids = [feats_preds[cat][b, i].item() for cat in self.morpho_encoder.feat_categories]

                # Apply morphological constraint filtering
                if self.use_morpho_constraints:
                    feat_ids = self.morpho_encoder.constrain_feats(upos_preds[b, i].item(), feat_ids)

                feats = self.morpho_encoder.decode_feats(feat_ids)

                # Decode dependency
                head = arc_preds_list[i]
                rel_id = rel_logits[b, i, head].argmax().item()
                deprel = self.deprel_encoder.decode(rel_id)

                # Decode lemma
                if lemma_preds is not None:
                    lemma_ids = lemma_preds[b, i].tolist()
                    lemma = self.char_vocab.decode(lemma_ids)
                    if not lemma:
                        lemma = word.form
                else:
                    lemma = word.form

                pred_words.append(
                    Token(
                        id=word.id,
                        form=word.form,
                        lemma=lemma,
                        upos=upos,
                        xpos=xpos,
                        feats=feats,
                        head=head,
                        deprel=deprel,
                        multiword=word.multiword,
                    )
                )

            # Re-insert multi-word tokens (MWTs) from gold sentence so that
            # conll18_ud_eval character sequences match between gold and system.
            # MWTs are evaluation-neutral (only their sub-words carry labels).
            # If predicted tokens already carry multiword info (e.g. from
            # LAMBO segmenter), Sentence.to_conllu() will generate MWT lines
            # automatically.  For gold-tokenized evaluation we still need to
            # re-insert the gold MWT lines explicitly.
            tokens = []
            has_pred_mwt = any(pw.multiword is not None for pw in pred_words)
            if has_pred_mwt:
                # Tokens carry their own multiword info — no need to
                # re-insert from gold.  to_conllu() will emit MWT lines.
                tokens = pred_words
            else:
                gold_mwts = [t for t in sentence.tokens if "-" in t.id]
                mwt_idx = 0
                for pw in pred_words:
                    word_id = int(pw.id)
                    # Insert any MWT whose range starts at this word
                    while mwt_idx < len(gold_mwts):
                        mwt = gold_mwts[mwt_idx]
                        mwt_start = int(mwt.id.split("-")[0])
                        if mwt_start == word_id:
                            tokens.append(Token(
                                id=mwt.id, form=mwt.form,
                                lemma="_", upos="_", xpos="_",
                            ))
                            mwt_idx += 1
                        else:
                            break
                    tokens.append(pw)

            predictions.append(
                Sentence(
                    tokens=tokens,
                    sent_id=sentence.sent_id,
                    text=sentence.text,
                )
            )

        return predictions
