"""Export a trained COMBO model for distribution and optionally upload to HuggingFace Hub.

Takes a training output directory (with checkpoints/ and encoders/) and packages
the model into the format expected by COMBO.from_pretrained():

    <output>/
    ├── config.json          # ModelConfig
    ├── pytorch_model.bin    # Weights only (no optimizer state)
    ├── README.md            # Model card (if eval results found)
    └── encoders/
        ├── morpho_encoder.json
        ├── deprel_encoder.json
        └── char_vocab.json
"""

import argparse
import json
import shutil
from dataclasses import asdict
from pathlib import Path

import torch
import yaml
from dotenv import load_dotenv
from jinja2 import Environment, FileSystemLoader

load_dotenv()

from combo.config import ModelConfig, load_config
from combo.data.label_encoder import MorphoEncoder, LabelEncoder, CharVocab


def export_model(
    training_dir: Path,
    checkpoint_name: str,
    base_model: str,
    output_dir: Path,
    config_path: Path | None = None,
    max_length: int = 512,
) -> Path:
    """Export a trained model to a distributable directory.

    Args:
        training_dir: Training output directory containing checkpoints/ and encoders/.
        checkpoint_name: Name of checkpoint file inside checkpoints/ (e.g. "checkpoint_best.pt").
        base_model: HuggingFace base model name (e.g. "xlm-roberta-base").
        output_dir: Directory to write the exported model to.
        config_path: Optional path to a config YAML to use as base. If not provided,
            a config is reconstructed from the checkpoint and encoders.
        max_length: Maximum sequence length for inference (saved in config.json).

    Returns:
        Path to the output directory.
    """
    checkpoint_path = training_dir / "checkpoints" / checkpoint_name
    encoders_dir = training_dir / "encoders"

    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
    if not encoders_dir.exists():
        raise FileNotFoundError(f"Encoders directory not found: {encoders_dir}")

    output_dir.mkdir(parents=True, exist_ok=True)

    # 1. Load checkpoint and strip optimizer/scheduler state
    print(f"Loading checkpoint: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

    slim_checkpoint = {
        "model_state_dict": checkpoint["model_state_dict"],
        "epoch": checkpoint.get("epoch"),
        "metrics": checkpoint.get("metrics"),
    }

    out_checkpoint = output_dir / "pytorch_model.bin"
    torch.save(slim_checkpoint, out_checkpoint)
    orig_size = checkpoint_path.stat().st_size / (1024 * 1024)
    new_size = out_checkpoint.stat().st_size / (1024 * 1024)
    print(f"Saved checkpoint: {orig_size:.1f} MB -> {new_size:.1f} MB (stripped optimizer state)")

    # 2. Copy encoders
    out_encoders = output_dir / "encoders"
    if out_encoders.exists():
        shutil.rmtree(out_encoders)
    shutil.copytree(encoders_dir, out_encoders)
    print(f"Copied encoders to {out_encoders}")

    # 3. Build config.json
    # Load encoders to get dynamic sizes
    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")

    if config_path and config_path.exists():
        with open(config_path) as f:
            raw = yaml.safe_load(f)
        model_dict = raw.get("model", {})
    else:
        # Reconstruct from defaults — infer architecture from state dict keys
        model_dict = asdict(ModelConfig())

    # Set the base model and dynamic encoder sizes
    model_dict["base_model"] = base_model
    model_dict["num_upos"] = morpho_encoder.num_upos
    model_dict["num_xpos"] = morpho_encoder.num_xpos
    model_dict["feat_vocab_sizes"] = morpho_encoder.feat_vocab_sizes
    model_dict["feat_categories"] = morpho_encoder.feat_categories
    model_dict["num_deprels"] = len(deprel_encoder)
    model_dict["char_vocab_size"] = len(char_vocab)

    # Infer LoRA config from state dict keys
    state_keys = list(checkpoint["model_state_dict"].keys())
    has_lora = any("lora_" in k for k in state_keys)
    model_dict["use_lora"] = has_lora

    if has_lora:
        # Infer target modules from key names
        lora_modules = set()
        for k in state_keys:
            if "lora_A" in k:
                parts = k.split(".")
                lora_idx = parts.index("lora_A")
                module_name = parts[lora_idx - 1]
                lora_modules.add(module_name)
        if lora_modules:
            model_dict["lora_target_modules"] = sorted(lora_modules)

        # Infer lora_r from weight shapes
        for k in state_keys:
            if "lora_A" in k:
                shape = checkpoint["model_state_dict"][k].shape
                model_dict["lora_r"] = shape[0]
                break

    config_out = output_dir / "config.json"
    with open(config_out, "w") as f:
        json.dump({"model": model_dict, "max_length": max_length}, f, indent=2)
    print(f"Wrote config: {config_out}")

    # 4. Print summary
    metrics = checkpoint.get("metrics", {})
    if metrics:
        print("\nModel metrics:")
        for k, v in metrics.items():
            if isinstance(v, float):
                print(f"  {k}: {v:.2f}")
            else:
                print(f"  {k}: {v}")

    print(f"\nExported model ready at: {output_dir}")
    return output_dir


from combo.languages import (
    Language,
    resolve_language,
    EXAMPLE_SENTENCES,
)


def generate_model_card(
    output_dir: Path,
    results_json: Path,
    language: str,
    base_model: str,
    repo_id: str,
    ud_treebank_names: list[str] | None = None,
    license: str = "cc-by-sa-4.0",
) -> Path:
    """Generate a HuggingFace model card README.md from eval results."""
    if ud_treebank_names is None:
        ud_treebank_names = []
    with open(results_json) as f:
        results = json.load(f)

    best_split = results.get("best_model_split", "dev")

    fulltext_key = f"{best_split}_fulltext_metrics"
    fulltext_data = results.get(fulltext_key, {})
    fulltext_f1 = fulltext_data.get("f1", {}) if fulltext_data else {}
    aligned_accuracy = fulltext_data.get("aligned_accuracy", {}) if fulltext_data else {}

    # Find template relative to the combo package
    import combo
    template_dir = Path(combo.__file__).parent
    env = Environment(
        loader=FileSystemLoader(str(template_dir)),
        keep_trailing_newline=True,
    )
    template = env.get_template("model_card_template.md.j2")

    model_name = repo_id.split("/")[-1].replace("-", " ").title() if repo_id else ""

    readme_content = template.render(
        language=language,
        language_code=resolve_language(language).code,
        base_model=base_model,
        repo_id=repo_id,
        model_name=model_name,
        ud_treebank_names=ud_treebank_names,
        license=license,
        aligned_accuracy=aligned_accuracy or None,
        fulltext_f1=fulltext_f1 or None,
        total_sentences=results.get("total_sentences", ""),
        total_words=results.get("total_words", ""),
        example_sentence=EXAMPLE_SENTENCES.get(
            resolve_language(language), "Example sentence."
        ),
    )

    readme_path = output_dir / "README.md"
    with open(readme_path, "w") as f:
        f.write(readme_content)
    print(f"Generated model card: {readme_path}")
    return readme_path


def push_to_hub(output_dir: Path, repo_id: str) -> str:
    """Upload exported model directory to HuggingFace Hub."""
    from huggingface_hub import HfApi

    api = HfApi()
    api.create_repo(repo_id, exist_ok=True, repo_type="model", private=True)

    print(f"Uploading to https://huggingface.co/{repo_id} ...")
    api.upload_folder(
        folder_path=str(output_dir),
        repo_id=repo_id,
        repo_type="model",
    )

    url = f"https://huggingface.co/{repo_id}"
    print(f"Uploaded: {url}")
    return url


def find_best_results(training_dir: Path) -> Path | None:
    """Find results_best.json saved during training."""
    results_best = training_dir / "results" / "results_best.json"
    return results_best if results_best.exists() else None


def main():
    parser = argparse.ArgumentParser(
        description="Export a trained COMBO model for distribution.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
Examples:
    combo-nlp-export --task-config config/tasks/english.yaml
    combo-nlp-export --task-config config/tasks/english.yaml --push-to-hub
""",
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=Path("config/base.yaml"),
        help="Path to base config YAML (default: config/base.yaml).",
    )
    parser.add_argument(
        "--task-config",
        type=Path,
        required=True,
        help="Path to task-specific config override (deep-merged on top of base config).",
    )
    parser.add_argument(
        "--push-to-hub",
        action="store_true",
        default=False,
        help="Upload the exported model to HuggingFace Hub (repo_id derived from config).",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="checkpoint_best.pt",
        help="Checkpoint filename inside checkpoints/ (default: checkpoint_best.pt).",
    )

    args = parser.parse_args()

    # Load config — all settings derived from it
    config = load_config(args.config, args.task_config)

    training_dir = config.language_output_dir
    base_model = config.model.base_model
    output_dir = Path("data/export") / config.naming.project / config.model_name
    repo_id = config.repo_id
    language = config.data.language
    treebanks = config.data.treebanks
    ud_treebank_names = [f"UD_{language}-{t}" for t in treebanks]

    print(f"Config:       {args.config}")
    print(f"Training dir: {training_dir}")
    print(f"Output dir:   {output_dir}")
    print(f"Model name:   {config.model_name}")
    print(f"Repo ID:      {repo_id}")
    print(f"Language:      {language}")
    print(f"Treebanks:    {', '.join(ud_treebank_names)}")
    print(f"Base model:   {base_model}")
    print()

    # Export model files
    export_model(
        training_dir=training_dir,
        checkpoint_name=args.checkpoint,
        base_model=base_model,
        output_dir=output_dir,
        config_path=args.config,
        max_length=config.data.max_length,
    )

    # Generate model card if best results exist
    results_json = find_best_results(training_dir)
    if results_json:
        print(f"\nFound best results: {results_json}")
        generate_model_card(
            output_dir=output_dir,
            results_json=results_json,
            language=language,
            base_model=base_model,
            repo_id=repo_id,
            ud_treebank_names=ud_treebank_names,
        )
    else:
        print("\nNo results_best.json found — skipping model card generation.")
        print(f"  (Expected at {training_dir / 'results' / 'results_best.json'})")

    # Push to HF Hub only if explicitly requested
    if args.push_to_hub:
        push_to_hub(output_dir, repo_id)


if __name__ == "__main__":
    main()
