#!/usr/bin/env python3
"""Download and export prachuryyaIITG/SampurNER_Tamil_IndicBERTv2 to ONNX for the PII Redaction POC."""

from __future__ import annotations

import json
import shutil
import subprocess
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
MODELS_DIR = REPO_ROOT / "models" / "ta"
MODEL_ID = "prachuryyaIITG/SampurNER_Tamil_IndicBERTv2"
REQUIRED_PACKAGES = ("transformers", "optimum[onnxruntime]", "onnx", "torch")
SENTENCEPIECE_CANDIDATES = (
    "sentencepiece.bpe.model",
    "spiece.model",
    "tokenizer.model",
)
TOKENIZER_JSON = "tokenizer.json"
VOCAB_TXT = "vocab.txt"


def ensure_dependencies() -> None:
    try:
        import optimum.onnxruntime  # noqa: F401
        import transformers  # noqa: F401
    except ImportError:
        print("Installing Python dependencies (this may take a few minutes)...")
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", *REQUIRED_PACKAGES],
            stdout=sys.stdout,
            stderr=sys.stderr,
        )


def copy_sentencepiece_model(source_dir: Path, target_dir: Path) -> Path | None:
    for name in SENTENCEPIECE_CANDIDATES:
        candidate = source_dir / name
        if candidate.exists():
            destination = target_dir / name
            shutil.copy(candidate, destination)
            return destination

    return None


def export_wordpiece_assets(source_dir: Path, target_dir: Path) -> list[Path]:
    saved: list[Path] = []
    tokenizer_json = source_dir / TOKENIZER_JSON
    if not tokenizer_json.exists():
        return saved

    destination = target_dir / TOKENIZER_JSON
    shutil.copy(tokenizer_json, destination)
    saved.append(destination)

    with tokenizer_json.open(encoding="utf-8") as tokenizer_file:
        tokenizer_data = json.load(tokenizer_file)

    vocab = tokenizer_data.get("model", {}).get("vocab")
    if not isinstance(vocab, dict):
        return saved

    vocab_path = target_dir / VOCAB_TXT
    ordered_tokens = [token for token, _ in sorted(vocab.items(), key=lambda item: item[1])]
    vocab_path.write_text("\n".join(ordered_tokens), encoding="utf-8")
    saved.append(vocab_path)
    return saved


def copy_tokenizer_assets(source_dir: Path, target_dir: Path) -> list[Path]:
    sentencepiece_path = copy_sentencepiece_model(source_dir, target_dir)
    if sentencepiece_path is not None:
        return [sentencepiece_path]

    wordpiece_assets = export_wordpiece_assets(source_dir, target_dir)
    if wordpiece_assets:
        print(
            "WARNING: Hugging Face repo has no SentencePiece model; saved WordPiece "
            f"assets ({', '.join(path.name for path in wordpiece_assets)}). "
            "TamilOnnxNerRunner uses vocab.txt (WordPiece) when present, "
            "otherwise SentencePiece model files."
        )
        return wordpiece_assets

    raise FileNotFoundError(
        f"No tokenizer assets found in {source_dir}. "
        f"Expected one of {SENTENCEPIECE_CANDIDATES} or {TOKENIZER_JSON}."
    )


def export_model() -> None:
    from optimum.onnxruntime import ORTModelForTokenClassification
    from transformers import AutoTokenizer

    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    temp_dir = MODELS_DIR / "_export_temp"
    if temp_dir.exists():
        shutil.rmtree(temp_dir)
    temp_dir.mkdir()

    print(f"Exporting {MODEL_ID} to ONNX...")
    model = ORTModelForTokenClassification.from_pretrained(MODEL_ID, export=True)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

    model.save_pretrained(temp_dir)
    tokenizer.save_pretrained(temp_dir)

    onnx_files = sorted(temp_dir.glob("*.onnx"))
    if not onnx_files:
        raise FileNotFoundError("Export completed but no .onnx file was produced.")

    target_onnx = MODELS_DIR / "model.onnx"
    shutil.copy(onnx_files[0], target_onnx)

    config_path = temp_dir / "config.json"
    with config_path.open(encoding="utf-8") as config_file:
        config = json.load(config_file)

    id2label = config.get("id2label", {})
    labels = [id2label[str(index)] for index in range(len(id2label))]
    labels_path = MODELS_DIR / "ner-labels.txt"
    labels_path.write_text("\n".join(labels), encoding="utf-8")

    tokenizer_assets = copy_tokenizer_assets(temp_dir, MODELS_DIR)
    shutil.rmtree(temp_dir)

    print()
    print("Tamil NER model assets saved:")
    print(f"  {target_onnx}")
    for asset in tokenizer_assets:
        print(f"  {asset}")
    print(f"  {labels_path}")
    print()
    print("Run from repository root:")
    print("  dotnet test --filter Category=TamilNer")


def main() -> int:
    ensure_dependencies()
    export_model()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())