from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
from pathlib import Path
from typing import Any, Dict, List

import mlx.core as mx
import numpy as np
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from modeling_eurobert_mlx import load_model  # noqa: E402


def _to_numpy(x: mx.array) -> np.ndarray:
    return np.array(x.astype(mx.float32))


def _timed(fn, repeats: int = 1) -> tuple[Any, float]:
    times = []
    out = None
    for _ in range(repeats):
        start = time.perf_counter()
        out = fn()
        times.append((time.perf_counter() - start) * 1000.0)
    return out, statistics.median(times)


def _load_tokenizer(model_dir: Path):
    # Keep the same tokenizer behavior used by pulpie.Extractor and by the
    # source model card. Transformers may warn about the regex; changing it
    # changes token IDs relative to the published checkpoint.
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True)
    if "<|sep|>" not in tokenizer.get_vocab():
        tokenizer.add_special_tokens({"additional_special_tokens": ["<|sep|>"]})
    return tokenizer


def run_load_checks(model_dir: Path, variants: List[str]) -> Dict[str, Any]:
    results = {}
    for variant in variants:
        model = load_model(model_dir, variant)
        input_ids = mx.array([[128000, 32, 128001]])
        attention_mask = mx.array([[1, 1, 1]])
        logits = model(input_ids, attention_mask)
        mx.eval(logits)
        results[variant] = {
            "loaded": True,
            "logits_shape": list(logits.shape),
            "logits_dtype": str(logits.dtype),
        }
    return results


def run_numerical_checks(
    source_dir: Path, model_dir: Path, variants: List[str]
) -> Dict[str, Any]:
    tokenizer = _load_tokenizer(model_dir)
    texts = ["A", "B", "C"]
    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=8)

    torch.set_num_threads(4)
    torch_model = AutoModelForTokenClassification.from_pretrained(
        str(source_dir),
        trust_remote_code=True,
        attn_implementation="eager",
        dtype=torch.float32,
    ).eval()

    with torch.inference_mode():
        torch_model(**inputs)
        torch_logits, torch_ms = _timed(
            lambda: torch_model(**inputs).logits.detach().cpu().numpy()
        )

    mx_inputs = {
        key: mx.array(value.numpy())
        for key, value in inputs.items()
        if key in {"input_ids", "attention_mask"}
    }
    results = {
        "test_inputs": texts,
        "token_shape": list(inputs["input_ids"].shape),
        "torch_reference": {
            "dtype": "float32",
            "attention": "eager",
            "latency_ms": torch_ms,
        },
        "variants": {},
    }

    for variant in variants:
        model = load_model(model_dir, variant)
        logits = model(mx_inputs["input_ids"], mx_inputs["attention_mask"])
        mx.eval(logits)
        logits, mlx_ms = _timed(
            lambda: _eval_logits(model, mx_inputs["input_ids"], mx_inputs["attention_mask"])
        )
        diff = np.abs(torch_logits - logits)
        results["variants"][variant] = {
            "latency_ms": mlx_ms,
            "max_abs_diff": float(diff.max()),
            "mean_abs_diff": float(diff.mean()),
        }
    return results


def _eval_logits(model, input_ids: mx.array, attention_mask: mx.array) -> np.ndarray:
    logits = model(input_ids, attention_mask)
    mx.eval(logits)
    return _to_numpy(logits)


def run_extraction(model_dir: Path, variants: List[str]) -> Dict[str, Any]:
    from pulpie.chunker import extract_blocks, pack_chunks, tokenize_blocks
    from pulpie.markdown import to_markdown
    from pulpie.model_utils import extract_item_ids, predictions_to_labels
    from pulpie.reconstruct import extract_main_html
    from pulpie.simplify import simplify

    html = (
        "<html><body><article><h1>Apple MLX conversion</h1>"
        "<p>This article explains how to convert a EuroBERT content extraction "
        "model to MLX format.</p></article></body></html>"
    )
    tokenizer = _load_tokenizer(model_dir)
    sep_token_id = tokenizer.convert_tokens_to_ids("<|sep|>")
    simplified, map_html = simplify(html)
    blocks = extract_blocks(simplified)
    item_ids = extract_item_ids(blocks)
    chunks = pack_chunks(
        tokenize_blocks(blocks, tokenizer),
        max_tokens=128,
        sep_token_id=sep_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

    results = {
        "input_html": html,
        "num_blocks": len(blocks),
        "chunk_lengths": [len(chunk_ids) for chunk_ids, _ in chunks],
        "variants": {},
    }

    for variant in variants:
        model = load_model(model_dir, variant)
        predictions = [0] * len(blocks)
        start = time.perf_counter()
        for chunk_ids, block_indices in chunks:
            input_ids = mx.array([chunk_ids])
            attention_mask = mx.ones_like(input_ids)
            logits = model(input_ids, attention_mask)
            mx.eval(logits)
            logits_np = _to_numpy(logits)[0]
            ids_np = np.array(chunk_ids)
            sep_positions = np.where(ids_np == sep_token_id)[0]
            preds = logits_np[sep_positions].argmax(axis=-1).tolist()
            for idx, block_idx in enumerate(block_indices):
                if idx < len(preds):
                    predictions[block_idx] = int(preds[idx])
        latency_ms = (time.perf_counter() - start) * 1000.0
        labels = predictions_to_labels(item_ids, predictions)
        main_html = extract_main_html(map_html, labels)
        markdown = to_markdown(main_html)
        results["variants"][variant] = {
            "latency_ms": latency_ms,
            "predictions": predictions,
            "labels": labels,
            "html": main_html,
            "markdown": markdown,
            "non_empty": bool(markdown.strip() or main_html.strip()),
        }
    return results


def main() -> None:
    parser = argparse.ArgumentParser(description="Verify converted MLX Pulpie weights.")
    parser.add_argument("--source-dir", type=Path, default=Path("source_model"))
    parser.add_argument("--model-dir", type=Path, default=Path("hf_out"))
    parser.add_argument(
        "--variants", nargs="+", default=["bf16", "8bit", "4bit"], help="Variants to verify."
    )
    parser.add_argument(
        "--output", type=Path, default=Path("hf_out/verification_report.json")
    )
    args = parser.parse_args()

    report = {
        "source_model": "feyninc/pulpie-orange-small",
        "model_dir": str(args.model_dir),
        "variants": args.variants,
        "load_checks": run_load_checks(args.model_dir, args.variants),
        "numerical_accuracy": run_numerical_checks(
            args.source_dir, args.model_dir, args.variants
        ),
        "end_to_end_extraction": run_extraction(args.model_dir, args.variants),
        "compute": {
            "environment": "Linux x86_64 CPU with mlx[cpu]; no paid cloud Mac used.",
            "estimated_incremental_cost_usd": 0.0,
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2, sort_keys=True)
        f.write("\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()