"""Evaluation entry point for COMBO.

Evaluates a trained model on the test set using the official CoNLL 2018
evaluation script (conll18_ud_eval.py).
"""

import argparse
import json
from datetime import datetime
from pathlib import Path

import torch
from transformers import AutoTokenizer

from combo.config import load_config
from combo.data import (
    CoNLLUReader,
    LabelEncoder,
    CharVocab,
    UDDataset,
    create_dataloaders,
)
from combo.data.label_encoder import MorphoEncoder
from combo.data.conllu_writer import CoNLLUWriter
from combo.evaluation import Evaluator
from combo.evaluation import conll18_ud_eval
from combo.model import UDParser
from combo.utils import get_device


def find_gold_test_file(config) -> Path:
    """Find the gold test CoNLL-U file for the configured language/treebank."""
    ud_path = Path(config.data.ud_path)
    for treebank in config.data.treebanks:
        treebank_dir = ud_path / f"UD_{config.data.language}-{treebank}"
        test_files = list(treebank_dir.glob("*-test.conllu"))
        if test_files:
            return test_files[0]
    raise FileNotFoundError(
        f"No test file found for {config.data.language} in {config.data.treebanks}"
    )


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate COMBO using official CoNLL 2018 evaluation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Evaluate exported model on a test file (local path or HuggingFace ID)
  combo-nlp-evaluate --model clarin-pl/combo-nlp-polish --test-file test.conllu
  combo-nlp-evaluate --model-dir /path/to/exported/model --test-file test.conllu

  # Evaluate using task config (auto-detects test file from UD treebank)
  combo-nlp-evaluate --task-config config/tasks/polish.yaml

  # Full-text evaluation (segmentation + parsing, end-to-end)
  combo-nlp-evaluate --model clarin-pl/combo-nlp-polish --test-file test.conllu --full-text --language Polish

  # Compare two CoNLL-U files directly (no model needed)
  combo-nlp-evaluate --gold-file gold.conllu --predictions-file predicted.conllu
""",
    )

    # File comparison mode (no model needed)
    parser.add_argument(
        "--gold-file", type=Path, default=None,
        help="Gold CoNLL-U file for direct file comparison (no model needed)",
    )
    parser.add_argument(
        "--predictions-file", type=Path, default=None,
        help="Predicted CoNLL-U file for direct file comparison",
    )

    # Model source
    parser.add_argument(
        "--model", type=str, default=None,
        help="Model: local path or HuggingFace ID (e.g. clarin-pl/combo-nlp-polish). "
             "Downloads from HF Hub if not a local directory.",
    )
    parser.add_argument(
        "--model-dir", type=Path, default=None,
        help="Path to exported model directory (contains config.json, pytorch_model.bin, encoders/). "
             "Alias for --model with a local path.",
    )
    parser.add_argument(
        "--task-config", type=Path, default=None,
        help="Path to task-specific configuration override",
    )
    parser.add_argument(
        "--config", type=Path, default=Path("config/base.yaml"),
        help="Path to base configuration file (used with --task-config)",
    )
    parser.add_argument(
        "--checkpoint", type=Path, default=None,
        help="Path to checkpoint (default: best checkpoint in training output dir)",
    )
    parser.add_argument(
        "--test-file", type=Path, default=None,
        help="Path to test file: .conllu or .txt. For --full-text, a .txt auto-resolves the matching .conllu "
             "for gold scoring (and vice versa). Required with --model/--model-dir, optional with --task-config.",
    )
    parser.add_argument(
        "--output-dir", type=Path, default=None,
        help="Directory to save results (default: <model-dir>/results/ or training output dir)",
    )
    parser.add_argument(
        "--save-predictions", action="store_true",
        help="Keep the predictions CoNLL-U file in the results directory",
    )
    parser.add_argument(
        "--full-text", action="store_true",
        help="Run full-text evaluation: segment raw text with LAMBO, then predict and compare against gold. "
             "Measures end-to-end performance including tokenization/segmentation.",
    )
    parser.add_argument(
        "--language", type=str, default=None,
        help="Language name for LAMBO segmenter (required with --full-text and --model-dir, "
             "auto-detected from task config). E.g. Polish, English, German.",
    )
    args = parser.parse_args()

    # === Mode 1: Direct file comparison ===
    if args.gold_file or args.predictions_file:
        if not (args.gold_file and args.predictions_file):
            parser.error("--gold-file and --predictions-file must be used together")

    if args.gold_file and args.predictions_file:
        print(f"Comparing files:")
        print(f"  Gold:   {args.gold_file}")
        print(f"  Predictions: {args.predictions_file}")
        evaluation = conll18_ud_eval.evaluate(
            conll18_ud_eval.load_conllu_file(str(args.gold_file)),
            conll18_ud_eval.load_conllu_file(str(args.predictions_file)),
        )
        print_results(evaluation)
        return

    # === Mode 2: Model evaluation ===
    # Resolve --model / --model-dir into a local directory
    model_dir = args.model_dir
    if args.model:
        local_path = Path(args.model)
        if local_path.exists():
            model_dir = local_path
        else:
            # Treat as HuggingFace Hub ID — download
            from combo.hub import download_model
            print(f"Downloading model from HuggingFace Hub: {args.model}")
            model_dir = download_model(args.model)

    if model_dir:
        # Load from exported model directory (local or downloaded from HF)
        if not args.test_file:
            parser.error("--test-file is required when using --model or --model-dir")

        print(f"Loading model from: {model_dir}")

        config = load_config(model_dir / "config.json")
        device = get_device(config.device)
        max_length = config.max_length

        # Load encoders
        encoders_dir = model_dir / "encoders"
        morpho_encoder = MorphoEncoder.load(encoders_dir / "morpho_encoder.json")
        deprel_encoder = LabelEncoder.load(encoders_dir / "deprel_encoder.json")
        char_vocab = CharVocab.load(encoders_dir / "char_vocab.json")

        config.model.num_upos = morpho_encoder.num_upos
        config.model.num_xpos = morpho_encoder.num_xpos
        config.model.feat_vocab_sizes = morpho_encoder.feat_vocab_sizes
        config.model.feat_categories = morpho_encoder.feat_categories
        config.model.num_deprels = len(deprel_encoder)
        config.model.char_vocab_size = len(char_vocab)

        # Load model
        checkpoint_path = model_dir / "pytorch_model.bin"
        print(f"Loading model from {checkpoint_path}")
        model = UDParser(config.model)
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        model.load_state_dict(checkpoint["model_state_dict"], strict=False)
        model = model.to(device)
        model.eval()

        tokenizer = AutoTokenizer.from_pretrained(config.model.base_model)
        gold_test_file = args.test_file
        results_dir = args.output_dir or Path("results")
        model_name = model_dir.name
        treebank_name = ""
        eval_batch_size = 32
        language = args.language

    elif args.task_config:
        # Load from task config
        config = load_config(args.config, args.task_config)
        device = get_device(config.device)
        max_length = config.data.max_length
        output_dir = config.language_output_dir

        # Load encoders
        encoders_dir = output_dir / "encoders"
        print(f"Loading encoders from {encoders_dir}")
        morpho_encoder = MorphoEncoder.load(encoders_dir / "morpho_encoder.json")
        deprel_encoder = LabelEncoder.load(encoders_dir / "deprel_encoder.json")
        char_vocab = CharVocab.load(encoders_dir / "char_vocab.json")

        config.model.num_upos = morpho_encoder.num_upos
        config.model.num_xpos = morpho_encoder.num_xpos
        config.model.feat_vocab_sizes = morpho_encoder.feat_vocab_sizes
        config.model.feat_categories = morpho_encoder.feat_categories
        config.model.num_deprels = len(deprel_encoder)
        config.model.char_vocab_size = len(char_vocab)

        # Load model from checkpoint
        checkpoint_path = args.checkpoint or (output_dir / "checkpoints" / "checkpoint_best.pt")
        print(f"Loading model from {checkpoint_path}")
        model = UDParser(config.model)
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        model.load_state_dict(checkpoint["model_state_dict"], strict=False)
        model = model.to(device)
        model.eval()

        epoch = checkpoint.get("epoch", "?")
        best_metric = checkpoint.get("best_metric", "?")
        print(f"  Checkpoint epoch: {epoch + 1 if isinstance(epoch, int) else epoch}")
        print(f"  Dev best LAS: {best_metric}")

        tokenizer = AutoTokenizer.from_pretrained(config.model.base_model)
        gold_test_file = args.test_file or find_gold_test_file(config)
        results_dir = args.output_dir or (output_dir / "results")
        model_name = config.model_name
        treebank_name = ",".join(config.data.treebanks)
        eval_batch_size = config.data.eval_batch_size
        language = args.language or config.data.language

    else:
        parser.error("Either --model, --model-dir, or --task-config is required for model evaluation")

    print(f"\nGold test file: {gold_test_file}")

    # Build evaluator (shared by both modes)
    evaluator = Evaluator(
        morpho_encoder, deprel_encoder, char_vocab,
        treebank=treebank_name,
        model_name=model_name,
        base_model=config.model.base_model,
        use_lemma_beam_search=config.model.use_lemma_beam_search,
        lemma_beam_width=config.model.lemma_beam_width,
        use_morpho_constraints=config.model.use_morpho_constraints,
    )

    results_dir = Path(results_dir)
    results_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    if args.full_text:
        # === Full-text evaluation: segmentation + parsing ===
        if not language:
            parser.error("--language is required for --full-text evaluation with --model or --model-dir")

        # Resolve .txt / .conllu pair: the .txt is raw input for the segmenter,
        # the gold .conllu is needed for scoring.
        if gold_test_file.suffix == ".txt":
            text_file = gold_test_file
            gold_conllu = gold_test_file.with_suffix(".conllu")
            if not gold_conllu.exists():
                parser.error(f"Gold .conllu file not found next to .txt: {gold_conllu}")
        else:
            gold_conllu = gold_test_file
            text_file = gold_test_file.with_suffix(".txt")
            if not text_file.exists():
                text_file = None
                print("  Warning: no .txt file found next to gold .conllu, "
                      "falling back to # text = metadata from CoNLL-U")

        print(f"\nFull-text evaluation (language: {language})")
        print(f"  Raw text:  {text_file or '(from # text = metadata)'}")
        print(f"  Gold file: {gold_conllu}")
        print("  Segmenter: LAMBO")

        result = evaluator.evaluate_full_text(
            model=model,
            text_path=text_file,
            gold_conllu_path=gold_conllu,
            tokenizer=tokenizer,
            max_length=max_length,
            batch_size=eval_batch_size,
            device=device,
            language=language,
            combo_model_id=args.model if args.model and "/" in args.model else None,
            show_progress=True,
        )

        # Print results from EvaluationResult
        print("\n" + "=" * 60)
        print("Full-text Evaluation Results (official CoNLL 2018 metrics)")
        print("=" * 60)

        fulltext_metrics = [
            "Tokens", "Sentences", "Words",
            "UPOS", "XPOS", "UFeats", "AllTags", "Lemmas",
            "UAS", "LAS", "CLAS", "MLAS", "BLEX",
        ]
        print("\n{:11}| {:>10} | {:>10}".format("Metric", "F1 Score", "Aligned"))
        print("-" * 42)
        for metric_name in fulltext_metrics:
            f1 = result.f1.get(metric_name)
            aligned = result.aligned_accuracy.get(metric_name)
            if f1 is not None:
                aligned_str = f"{aligned:10.2f}" if aligned is not None else "       n/a"
                print("{:11}| {:10.2f} | {}".format(metric_name, f1, aligned_str))

        # Save results as JSON
        metrics = {}
        for name in fulltext_metrics:
            f1 = result.f1.get(name)
            aligned = result.aligned_accuracy.get(name)
            if f1 is not None:
                metrics[name] = {"f1": f1, "aligned": aligned}
        metrics["mode"] = "full_text"
        metrics["language"] = language
        metrics["checkpoint"] = str(checkpoint_path)
        metrics["gold_file"] = str(gold_test_file)

        results_file = results_dir / f"eval_fulltext_{timestamp}.json"
        with open(results_file, "w") as f:
            json.dump(metrics, f, indent=2)
        print(f"\nResults saved to {results_file}")

    else:
        # === Aligned evaluation: gold tokenization ===
        test_doc = CoNLLUReader.read(
            gold_test_file,
            treebank_id="eval",
            language="unknown",
        )
        print(f"  {len(test_doc.sentences):,} sentences")

        test_dataset = UDDataset(
            [test_doc], tokenizer, morpho_encoder, deprel_encoder,
            char_vocab, max_length,
        )
        _, test_loader = create_dataloaders(
            test_dataset, test_dataset,
            train_batch_size=eval_batch_size,
            eval_batch_size=eval_batch_size,
            num_workers=0,
            pad_token_id=tokenizer.pad_token_id,
            char_pad_id=char_vocab.pad_id,
        )

        stats = test_dataset.scan_truncation()
        if stats["truncated_count"] > 0:
            print(f"  Skipped: {stats['truncated_count']}/{stats['total_sentences']} sentences "
                  f"exceed max_length={max_length} (removed from evaluation)")

        print("\nRunning predictions...")
        _, predictions = evaluator.evaluate_and_predict(
            model, test_loader, device, show_progress=True,
        )

        if test_dataset.truncated_count > 0:
            print(f"\n  Note: {test_dataset.truncated_count} sentences were skipped "
                  f"(exceeded max_length={max_length})")

        # Write predictions to CoNLL-U file
        pred_file = results_dir / f"predictions_{timestamp}.conllu"
        CoNLLUWriter.write(predictions, pred_file)
        print(f"Predictions written to {pred_file}")

        # Evaluate using official conll18_ud_eval
        print("\nEvaluating with official CoNLL 2018 script...")
        gold_ud = conll18_ud_eval.load_conllu_file(str(gold_test_file))
        system_ud = conll18_ud_eval.load_conllu_file(str(pred_file))
        evaluation = conll18_ud_eval.evaluate(gold_ud, system_ud)

        # Print results
        print_results(evaluation)

        # Save results as JSON
        aligned_metrics = [
            "UPOS", "XPOS", "UFeats", "AllTags", "Lemmas",
            "UAS", "LAS", "CLAS", "MLAS", "BLEX",
        ]
        metrics = {}
        for metric in aligned_metrics:
            s = evaluation[metric]
            metrics[metric] = {
                "f1": s.f1 * 100,
                "aligned": s.aligned_accuracy * 100 if s.aligned_accuracy is not None else None,
            }
        metrics["mode"] = "aligned"
        metrics["checkpoint"] = str(checkpoint_path)
        metrics["gold_file"] = str(gold_test_file)
        metrics["total_sentences"] = len(test_doc.sentences)
        metrics["total_words"] = sum(len(s.words) for s in test_doc.sentences)

        results_file = results_dir / f"eval_results_{timestamp}.json"
        with open(results_file, "w") as f:
            json.dump(metrics, f, indent=2)
        print(f"\nResults saved to {results_file}")

        # Clean up predictions file unless user wants to keep it
        if not args.save_predictions:
            pred_file.unlink()
            print(f"Predictions file removed (use --save-predictions to keep)")


def print_results(evaluation: dict) -> None:
    """Print evaluation results in a readable format."""
    print("\n" + "=" * 72)
    print("Evaluation Results (official CoNLL 2018 metrics)")
    print("=" * 72)

    print("\n{:11}| {:>10} | {:>10} | {:>10} | {:>10}".format(
        "Metric", "Precision", "Recall", "F1 Score", "Aligned"))
    print("-" * 72)

    for metric in ["UPOS", "XPOS", "UFeats", "AllTags", "Lemmas",
                    "UAS", "LAS", "CLAS", "MLAS", "BLEX"]:
        s = evaluation[metric]
        aligned = s.aligned_accuracy
        aligned_str = f"{aligned * 100:10.2f}" if aligned is not None else "       n/a"
        print("{:11}| {:10.2f} | {:10.2f} | {:10.2f} | {}".format(
            metric, s.precision * 100, s.recall * 100, s.f1 * 100, aligned_str))


if __name__ == "__main__":
    main()
