"""
Converts a saved HuggingFace model to ONNX or TorchScript.
Called via asyncio.to_thread() — this blocks for 30-120 seconds.
"""
import logging
import tempfile
from pathlib import Path
from typing import Literal

logger = logging.getLogger(__name__)


def export_model(
    artifact_path: str,
    export_format: Literal["onnx", "torchscript"],
    opset_version: int = 14,
    optimize: bool = True,
) -> Path:
    """
    Convert the model at artifact_path to the requested format.
    Returns the path to the exported file in a temp directory.
    Caller is responsible for cleanup (shutil.rmtree on the parent dir).
    """
    import torch
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    model_path = Path(artifact_path)
    if not model_path.exists():
        raise FileNotFoundError(f"Model directory not found: {model_path}")

    logger.info("Loading model from %s for %s export", model_path, export_format)

    model = AutoModelForSequenceClassification.from_pretrained(str(model_path))
    tokenizer = AutoTokenizer.from_pretrained(str(model_path))

    # Merge PEFT/LoRA adapters into base weights before export
    if hasattr(model, "merge_and_unload"):
        logger.info("Detected PEFT model — merging LoRA adapters before export")
        model = model.merge_and_unload()

    # Dequantize only if the model is actually quantized (QLoRA / bitsandbytes).
    # transformers.PreTrainedModel always has a dequantize() method but raises
    # ValueError when called on a non-quantized model — guard with is_quantized.
    if getattr(model, "is_quantized", False) and hasattr(model, "dequantize"):
        logger.info("Dequantizing model for export")
        model = model.dequantize()

    model.eval()

    tmp_dir = Path(tempfile.mkdtemp(prefix="modelforge_export_"))

    if export_format == "torchscript":
        return _export_torchscript(model, tokenizer, tmp_dir)

    return _export_onnx(model, tokenizer, tmp_dir, opset_version, optimize)


def _make_dummy_inputs(tokenizer, max_length: int = 128):
    import torch
    dummy = tokenizer(
        "example input text",
        return_tensors="pt",
        max_length=max_length,
        padding="max_length",
        truncation=True,
    )
    return dummy


def _export_torchscript(model, tokenizer, tmp_dir: Path) -> Path:
    import torch

    dummy = _make_dummy_inputs(tokenizer)
    out_path = tmp_dir / "model.pt"

    try:
        with torch.no_grad():
            traced = torch.jit.trace(
                model,
                (dummy["input_ids"], dummy["attention_mask"]),
                strict=False,
            )
        torch.jit.save(traced, str(out_path))
        logger.info("TorchScript export successful: %s", out_path)
        return out_path
    except Exception as exc:
        logger.warning("TorchScript trace failed (%s), trying torch.jit.script()", exc)
        try:
            with torch.no_grad():
                scripted = torch.jit.script(model)
            torch.jit.save(scripted, str(out_path))
            logger.info("TorchScript script export successful: %s", out_path)
            return out_path
        except Exception as exc2:
            raise RuntimeError(
                f"TorchScript export failed (trace: {exc}; script: {exc2}). "
                "Consider using ONNX export instead."
            ) from exc2


def _export_onnx(model, tokenizer, tmp_dir: Path, opset_version: int, optimize: bool) -> Path:
    import torch

    # Try optimum first (produces optimized ONNX with graph optimizations)
    try:
        from optimum.onnxruntime import ORTModelForSequenceClassification

        out_dir = tmp_dir / "optimum_onnx"
        ort_model = ORTModelForSequenceClassification.from_pretrained(
            model.config._name_or_path,
            export=True,
        )
        ort_model.save_pretrained(str(out_dir))
        onnx_files = list(out_dir.glob("*.onnx"))
        if onnx_files:
            logger.info("Optimum ONNX export successful: %s", onnx_files[0])
            return onnx_files[0]
    except Exception as exc:
        logger.info("optimum not available or failed (%s), using torch.onnx.export fallback", exc)

    # Fallback: torch.onnx.export
    dummy = _make_dummy_inputs(tokenizer)
    out_path = tmp_dir / "model.onnx"

    try:
        with torch.no_grad():
            torch.onnx.export(
                model,
                (dummy["input_ids"], dummy["attention_mask"]),
                str(out_path),
                input_names=["input_ids", "attention_mask"],
                output_names=["logits"],
                dynamic_axes={
                    "input_ids":      {0: "batch", 1: "seq"},
                    "attention_mask": {0: "batch", 1: "seq"},
                },
                opset_version=opset_version,
                do_constant_folding=optimize,
            )
        logger.info("torch.onnx.export successful: %s", out_path)
        return out_path
    except Exception as exc:
        raise RuntimeError(f"ONNX export failed: {exc}") from exc