#!/usr/bin/env python3
import argparse
import json
import os
import shutil
import tempfile
from pathlib import Path

os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
os.environ.setdefault("TRANSFORMERS_NO_FLAX", "1")
os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1")
os.environ["USE_TF"] = "0"
os.environ["USE_FLAX"] = "0"
os.environ["USE_TORCH"] = "1"

import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline

from irish_core_decoder import repair_irish_core_spans


def load_tokenizer(model_ref: str):
    tokenizer_ref = model_ref
    tokenizer_path = Path(model_ref)
    if tokenizer_path.exists():
        tokenizer_cfg_path = tokenizer_path / "tokenizer_config.json"
        if tokenizer_cfg_path.exists():
            data = json.loads(tokenizer_cfg_path.read_text(encoding="utf-8"))
            if "fix_mistral_regex" in data:
                tmpdir = Path(tempfile.mkdtemp(prefix="openmed_tokenizer_"))
                keep = {
                    "tokenizer_config.json",
                    "tokenizer.json",
                    "special_tokens_map.json",
                    "vocab.txt",
                    "vocab.json",
                    "merges.txt",
                    "added_tokens.json",
                    "sentencepiece.bpe.model",
                    "spiece.model",
                }
                for child in tokenizer_path.iterdir():
                    if child.is_file() and child.name in keep:
                        shutil.copy2(child, tmpdir / child.name)
                data.pop("fix_mistral_regex", None)
                (tmpdir / "tokenizer_config.json").write_text(
                    json.dumps(data, ensure_ascii=False, indent=2) + "\n",
                    encoding="utf-8",
                )
                tokenizer_ref = str(tmpdir)
    try:
        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True, fix_mistral_regex=True)
    except Exception:
        pass
    try:
        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True, fix_mistral_regex=False)
    except TypeError:
        pass
    try:
        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True)
    except Exception:
        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=False)


def mask_text(text: str, spans: list[dict]) -> str:
    out = text
    for span in sorted(spans, key=lambda item: (item["start"], item["end"]), reverse=True):
        out = out[:span["start"]] + f"[{span['label']}]" + out[span["end"]:]
    return out


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=".")
    parser.add_argument("--text", required=True)
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    parser.add_argument("--ppsn-min-score", type=float, default=0.55)
    parser.add_argument("--other-min-score", type=float, default=0.50)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    tokenizer = load_tokenizer(args.model)
    model = AutoModelForTokenClassification.from_pretrained(args.model)
    if args.device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device
    model.to(device)
    model.eval()

    nlp = pipeline(
        "token-classification",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple",
        device=0 if device == "cuda" else -1,
    )
    general = nlp(args.text)
    spans = repair_irish_core_spans(
        args.text,
        model,
        tokenizer,
        general,
        other_min_score=args.other_min_score,
        ppsn_min_score=args.ppsn_min_score,
    )
    result = {
        "model": args.model,
        "masked_text": mask_text(args.text, spans),
        "spans": spans,
        "ppsn_decoder": "word_aligned",
        "general_decoder": "irish_core_label_aware",
        "ppsn_min_score": args.ppsn_min_score,
        "other_min_score": args.other_min_score,
        "backend": "transformers",
    }
    if args.json:
        print(json.dumps(result, indent=2, ensure_ascii=False))
    else:
        print(result["masked_text"])


if __name__ == "__main__":
    main()