"""Training entry point for COMBO."""

import argparse
import random
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from dotenv import load_dotenv
from transformers import AutoTokenizer

load_dotenv()

from combo.config import Config, load_config
from combo.data import (
    CoNLLUReader,
    LabelEncoder,
    CharVocab,
    UDDataset,
    create_dataloaders,
)
from combo.data.label_encoder import MorphoEncoder
from combo.data.conllu_reader import Document
from combo.model import UDParser
from combo.training import Trainer
from combo.evaluation import Evaluator
from combo.utils import get_device, setup_for_device, WandbLogger


def set_seed(seed: int) -> None:
    """Set random seed for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def limit_documents(docs: list, max_samples: int | None) -> list:
    """Limit total sentences across all documents to max_samples."""
    if not docs or max_samples is None:
        return docs

    all_sentences = []
    for doc in docs:
        all_sentences.extend(doc.sentences)

    limited_sentences = all_sentences[:max_samples]

    return [Document(
        sentences=limited_sentences,
        treebank_id=docs[0].treebank_id,
        language=docs[0].language,
    )]


@dataclass
class SplitPaths:
    """Paths to .txt and .conllu files for a split (dev/test)."""
    text_path: Path | None = None
    conllu_path: Path | None = None


def load_language_data(config: Config) -> tuple[list, list, list, SplitPaths, SplitPaths]:
    """Load and combine all treebanks for a language.

    Returns:
        Tuple of (train_docs, dev_docs, test_docs, dev_paths, test_paths).
    """
    ud_path = Path(config.data.ud_path)
    language = config.data.language
    treebanks = config.data.treebanks

    train_docs = []
    dev_docs = []
    test_docs = []
    dev_paths = SplitPaths()
    test_paths = SplitPaths()

    print(f"\nLoading data for {language}...")

    for treebank in treebanks:
        treebank_dir = ud_path / f"UD_{language}-{treebank}"

        if not treebank_dir.exists():
            print(f"  Warning: Treebank not found: {treebank_dir}")
            continue

        train_files = list(treebank_dir.glob("*-train.conllu"))
        dev_files = list(treebank_dir.glob("*-dev.conllu"))
        test_files = list(treebank_dir.glob("*-test.conllu"))

        treebank_id = f"{language.lower()}_{treebank.lower()}"

        if train_files:
            doc = CoNLLUReader.read(train_files[0], treebank_id, language)
            train_docs.append(doc)
            print(f"  Loaded {len(doc.sentences):,} train sentences from {treebank}")

        if dev_files:
            doc = CoNLLUReader.read(dev_files[0], treebank_id, language)
            dev_docs.append(doc)
            dev_paths.conllu_path = dev_files[0]
            dev_txt = dev_files[0].with_suffix(".txt")
            if dev_txt.exists():
                dev_paths.text_path = dev_txt
            print(f"  Loaded {len(doc.sentences):,} dev sentences from {treebank}")

        if test_files:
            doc = CoNLLUReader.read(test_files[0], treebank_id, language)
            test_docs.append(doc)
            test_paths.conllu_path = test_files[0]
            test_txt = test_files[0].with_suffix(".txt")
            if test_txt.exists():
                test_paths.text_path = test_txt
            print(f"  Loaded {len(doc.sentences):,} test sentences from {treebank}")

    # Report full-text evaluation availability
    for name, paths in [("dev", dev_paths), ("test", test_paths)]:
        if paths.text_path:
            print(f"  Full-text eval available for {name}: {paths.text_path}")
        elif paths.conllu_path:
            print(f"  No .txt file for {name} — full-text eval disabled for this split")

    return train_docs, dev_docs, test_docs, dev_paths, test_paths


def build_encoders(train_docs: list) -> tuple:
    """Build label encoders from training data."""
    print("\nBuilding encoders...")

    morpho_encoder = MorphoEncoder().fit(train_docs)
    print(f"  UPOS labels: {morpho_encoder.num_upos}")
    print(f"  XPOS labels: {morpho_encoder.num_xpos}")
    print(f"  Feature categories: {morpho_encoder.num_feat_categories}")
    for cat, size in morpho_encoder.feat_vocab_sizes.items():
        print(f"    {cat}: {size}")

    deprel_encoder = LabelEncoder().fit_documents(train_docs, field="deprel")
    print(f"  Deprel labels: {len(deprel_encoder):,}")

    char_vocab = CharVocab().fit(train_docs)
    print(f"  Character vocab: {len(char_vocab):,}")

    return morpho_encoder, deprel_encoder, char_vocab


def save_encoders(
    morpho_encoder: MorphoEncoder,
    deprel_encoder: LabelEncoder,
    char_vocab: CharVocab,
    output_dir: Path,
) -> None:
    """Save encoders to disk."""
    encoders_dir = output_dir / "encoders"
    encoders_dir.mkdir(parents=True, exist_ok=True)

    morpho_encoder.save(encoders_dir / "morpho_encoder.json")
    deprel_encoder.save(encoders_dir / "deprel_encoder.json")
    char_vocab.save(encoders_dir / "char_vocab.json")
    print(f"  Saved encoders to {encoders_dir}")


def main():
    sys.stdout.reconfigure(line_buffering=True)

    parser = argparse.ArgumentParser(description="Train COMBO")
    parser.add_argument(
        "--config",
        type=Path,
        default=Path("config/base.yaml"),
        help="Path to base configuration file",
    )
    parser.add_argument(
        "--task-config",
        type=Path,
        required=True,
        help="Path to task-specific configuration override (deep-merged on top of base config)",
    )
    parser.add_argument(
        "--resume",
        type=Path,
        default=None,
        help="Path to checkpoint to resume training from",
    )
    args = parser.parse_args()

    # Load configuration
    print(f"Loading configuration from {args.config}")
    config = load_config(args.config, args.task_config)

    # Set seed
    set_seed(config.seed)
    print(f"Random seed: {config.seed}")

    # Setup device
    device = get_device(config.device)

    # Load data
    train_docs, dev_docs, test_docs, dev_paths, test_paths = load_language_data(config)

    # Apply sample limits if configured
    if config.data.max_train_samples:
        train_docs = limit_documents(train_docs, config.data.max_train_samples)
        total_train = sum(len(d.sentences) for d in train_docs)
        print(f"  Limited training to {total_train:,} sentences")
    if config.data.max_eval_samples:
        dev_docs = limit_documents(dev_docs, config.data.max_eval_samples)
        total_dev = sum(len(d.sentences) for d in dev_docs)
        print(f"  Limited dev to {total_dev:,} sentences")

    if not train_docs:
        print("Error: No training data found!")
        return

    # Build encoders
    morpho_encoder, deprel_encoder, char_vocab = build_encoders(train_docs)

    # Save encoders
    output_dir = config.language_output_dir
    save_encoders(morpho_encoder, deprel_encoder, char_vocab, output_dir)

    # Update config with encoder sizes
    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 tokenizer
    print(f"\nLoading tokenizer: {config.model.base_model}")
    tokenizer = AutoTokenizer.from_pretrained(config.model.base_model)

    # Create datasets
    print("\nCreating datasets...")
    train_dataset = UDDataset(
        train_docs,
        tokenizer,
        morpho_encoder,
        deprel_encoder,
        char_vocab,
        config.data.max_length,
    )
    print(f"  Train: {len(train_dataset):,} sentences")

    dev_dataset = UDDataset(
        dev_docs,
        tokenizer,
        morpho_encoder,
        deprel_encoder,
        char_vocab,
        config.data.max_length,
    )
    print(f"  Dev: {len(dev_dataset):,} sentences")

    test_dataset = None
    if test_docs:
        test_dataset = UDDataset(
            test_docs,
            tokenizer,
            morpho_encoder,
            deprel_encoder,
            char_vocab,
            config.data.max_length,
        )
        print(f"  Test: {len(test_dataset):,} sentences")

    # Report skipped sentences (too long for max_length)
    print(f"\nSentences exceeding max_length={config.data.max_length} (skipped):")
    for name, ds in [("Train", train_dataset), ("Dev", dev_dataset), ("Test", test_dataset)]:
        if ds is None:
            continue
        stats = ds.scan_truncation()
        if stats["truncated_count"] > 0:
            print(f"  {name}: {stats['truncated_count']}/{stats['total_sentences']} sentences "
                  f"skipped ({stats['truncated_words']} words would be dropped)")
        else:
            print(f"  {name}: none skipped")

    # Create dataloaders
    train_loader, dev_loader = create_dataloaders(
        train_dataset,
        dev_dataset,
        train_batch_size=config.data.train_batch_size,
        eval_batch_size=config.data.eval_batch_size,
        num_workers=config.data.num_workers,
        pad_token_id=tokenizer.pad_token_id,
        char_pad_id=char_vocab.pad_id,
    )

    test_loader = None
    if test_dataset is not None:
        from functools import partial
        from torch.utils.data import DataLoader
        from combo.data.dataset import collate_fn
        collate = partial(collate_fn, pad_token_id=tokenizer.pad_token_id, char_pad_id=char_vocab.pad_id)
        test_loader = DataLoader(
            test_dataset,
            batch_size=config.data.eval_batch_size,
            shuffle=False,
            num_workers=config.data.num_workers,
            collate_fn=collate,
        )

    # Create model
    print(f"\nCreating model: {config.model.base_model}")
    print(f"  LoRA enabled: {config.model.use_lora}")
    if config.model.use_lora:
        print(f"  LoRA r: {config.model.lora_r}")
        print(f"  LoRA alpha: {config.model.lora_alpha}")

    config.model.gradient_checkpointing = config.training.gradient_checkpointing
    model = UDParser(config.model)
    model = setup_for_device(model, device)

    print(f"  Total params: {model.get_total_params():,}")
    print(f"  Trainable params: {model.get_trainable_params():,}")

    # Create evaluator
    evaluator = Evaluator(
        morpho_encoder,
        deprel_encoder,
        char_vocab,
        treebank=",".join(config.data.treebanks),
        model_name=config.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,
    )

    # Create logger
    logger = WandbLogger(config.wandb, config)

    # Create trainer
    trainer = Trainer(
        model=model,
        train_dataloader=train_loader,
        dev_dataloader=dev_loader,
        config=config,
        evaluator=evaluator,
        logger=logger,
        device=device,
        test_dataloader=test_loader,
        tokenizer=tokenizer,
        dev_paths=dev_paths,
        test_paths=test_paths,
    )

    # Resume from checkpoint if specified
    if args.resume:
        print(f"\nResuming from checkpoint: {args.resume}")
        trainer.load_checkpoint(args.resume)

    # Train
    results = trainer.train()

    # Finish logging
    logger.finish()
    print("\nTraining complete!")


if __name__ == "__main__":
    main()
