"""Run performance benchmarks (decode TPS, end-to-end latency, and memory).

Usage:
    python benchmarks/run_performance.py \
        --model-path AutomatosX/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8 \
        --image test_image.jpg \
        --output benchmarks/results/performance_results.json
"""

from __future__ import annotations

import argparse
import importlib.metadata
import json
import math
import platform
import sys
import time
from numbers import Integral, Real
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))


def get_memory_usage_mb() -> float:
    """Get the process peak resident-set size in MB."""
    try:
        import resource
        usage = resource.getrusage(resource.RUSAGE_SELF)
        # macOS reports bytes; Linux and most BSD-derived Python builds report
        # KiB. The benchmark is Apple-focused, but avoid nonsensical values
        # when its reporting helpers are run elsewhere.
        divisor = 1024 * 1024 if sys.platform == "darwin" else 1024
        return usage.ru_maxrss / divisor
    except Exception:
        return 0.0


def get_system_info() -> dict:
    """Collect system information for reproducibility."""
    info = {
        "platform": platform.platform(),
        "processor": platform.processor(),
        "python_version": platform.python_version(),
        "machine": platform.machine(),
    }
    for distribution, key in (("mlx", "mlx_version"), ("mlx-vlm", "mlx_vlm_version")):
        try:
            info[key] = importlib.metadata.version(distribution)
        except importlib.metadata.PackageNotFoundError:
            info[key] = None

    # Try to get chip info on macOS
    try:
        import subprocess
        result = subprocess.run(
            ["sysctl", "-n", "machdep.cpu.brand_string"],
            capture_output=True, text=True, timeout=5
        )
        if result.returncode == 0:
            info["chip"] = result.stdout.strip()
    except Exception:
        pass

    # Try to get memory info
    try:
        import subprocess
        result = subprocess.run(
            ["sysctl", "-n", "hw.memsize"],
            capture_output=True, text=True, timeout=5
        )
        if result.returncode == 0:
            info["total_memory_gb"] = int(result.stdout.strip()) / (1024**3)
    except Exception:
        pass

    return info


def run_performance_benchmark(
    model_path: str,
    image_path: str,
    prompt: str = "<image>document parsing.",
    max_tokens: int = 4096,
    num_warmup: int = 1,
    num_runs: int = 3,
    served_revision: str | None = None,
) -> dict:
    """Run performance benchmark measuring TPS, memory, and latency."""
    if not Path(image_path).is_file():
        raise FileNotFoundError(f"Test image not found: {image_path}")
    if not isinstance(prompt, str) or prompt.count("<image>") != 1:
        raise ValueError("prompt must contain exactly one literal '<image>' token")
    if not isinstance(max_tokens, int) or isinstance(max_tokens, bool) or max_tokens < 1:
        raise ValueError("max_tokens must be a positive integer")
    if not isinstance(num_warmup, int) or isinstance(num_warmup, bool) or num_warmup < 0:
        raise ValueError("num_warmup must be a non-negative integer")
    if not isinstance(num_runs, int) or isinstance(num_runs, bool) or num_runs < 1:
        raise ValueError("num_runs must be a positive integer")
    if served_revision is not None and (
        not isinstance(served_revision, str) or not served_revision.strip()
    ):
        raise ValueError("served_revision must be a non-empty string or None")

    from mlx_vlm import load, generate

    system_info = get_system_info()
    total_memory = system_info.get("total_memory_gb")
    memory_label = f"{total_memory:.0f} GB" if isinstance(total_memory, (int, float)) else "Unknown memory"
    print(f"System: {system_info.get('chip', 'Unknown')} | {memory_label}")

    # Measure model load time
    load_start = time.perf_counter()
    model, processor = load(model_path, revision=served_revision)
    load_time = time.perf_counter() - load_start
    mem_after_load = get_memory_usage_mb()

    print(f"Model load time: {load_time:.2f}s")
    print(f"Memory after load: {mem_after_load:.0f} MB")

    # Warmup runs
    print(f"\nWarmup ({num_warmup} run(s))...")
    for _ in range(num_warmup):
        generate(model, processor, prompt=prompt, image=[image_path],
                 max_tokens=128, temperature=0.0, verbose=False)

    import mlx.core as mx
    mx.clear_cache()

    # Benchmark runs
    print(f"Benchmarking ({num_runs} runs, max_tokens={max_tokens})...")
    runs = []

    for i in range(num_runs):
        active_memory_before_mb = mx.get_active_memory() / 1e6
        if not math.isfinite(active_memory_before_mb) or active_memory_before_mb < 0:
            raise ValueError("MLX returned invalid active memory")
        mx.reset_peak_memory()
        start = time.perf_counter()

        response = generate(
            model, processor,
            prompt=prompt,
            image=[image_path],
            max_tokens=max_tokens,
            temperature=0.0,
            verbose=False,
        )

        elapsed = time.perf_counter() - start
        if not math.isfinite(elapsed) or elapsed <= 0:
            raise ValueError("Performance timer returned an invalid elapsed time")
        process_peak_mb = get_memory_usage_mb()
        if not math.isfinite(process_peak_mb) or process_peak_mb < 0:
            raise ValueError("Process peak memory is invalid")

        # Extract token count
        if isinstance(response, str):
            text = response
            tokens = len(text.split())
            token_count_source = "whitespace estimate"
        elif hasattr(response, "text"):
            text = response.text
            if not isinstance(text, str):
                raise TypeError("mlx-vlm returned a non-string text field")
            raw_tokens = getattr(
                response,
                "generation_tokens",
                getattr(response, "tokens_generated", None),
            )
            if raw_tokens is None:
                tokens = len(text.split())
                token_count_source = "whitespace estimate"
            else:
                tokens = raw_tokens
                token_count_source = "mlx-vlm token count"
        elif isinstance(response, dict):
            text = response.get("text", "")
            if not isinstance(text, str):
                raise TypeError("mlx-vlm returned a non-string text field")
            raw_tokens = response.get(
                "generation_tokens", response.get("tokens_generated")
            )
            if raw_tokens is None:
                tokens = len(text.split())
                token_count_source = "whitespace estimate"
            else:
                tokens = raw_tokens
                token_count_source = "mlx-vlm token count"
        else:
            text = str(response)
            tokens = len(text.split())
            token_count_source = "whitespace estimate"
        if (
            not isinstance(tokens, Integral)
            or isinstance(tokens, bool)
            or not 0 < tokens <= max_tokens
        ):
            raise ValueError("mlx-vlm returned an invalid generation token count")
        tokens = int(tokens)

        # Prefer mlx-vlm's decode-only rate; the fallback is end-to-end and is
        # explicitly identified in each run.
        reported_tps = (
            response.get("generation_tps", 0.0)
            if isinstance(response, dict)
            else getattr(response, "generation_tps", 0.0)
        )
        if reported_tps:
            if (
                not isinstance(reported_tps, Real)
                or isinstance(reported_tps, bool)
                or not math.isfinite(float(reported_tps))
                or reported_tps <= 0
            ):
                raise ValueError("mlx-vlm returned invalid generation throughput")
            tps = float(reported_tps)
            tps_source = "mlx-vlm generation_tps"
        else:
            tps = tokens / elapsed if elapsed > 0 and tokens > 0 else 0
            tps_source = "end-to-end estimate"

        raw_peak_memory = (
            response.get("peak_memory", 0.0)
            if isinstance(response, dict)
            else getattr(response, "peak_memory", 0.0)
        )
        if (
            not isinstance(raw_peak_memory, Real)
            or isinstance(raw_peak_memory, bool)
            or not math.isfinite(float(raw_peak_memory))
            or raw_peak_memory < 0
        ):
            raise ValueError("mlx-vlm returned invalid peak memory")
        reported_peak_gb = float(raw_peak_memory)
        mlx_peak_mb = (
            reported_peak_gb * 1000
            if reported_peak_gb > 0
            else mx.get_peak_memory() / 1e6
        )
        if not math.isfinite(mlx_peak_mb) or mlx_peak_mb < 0:
            raise ValueError("MLX returned invalid peak memory")
        run_result = {
            "run": i + 1,
            "elapsed_seconds": elapsed,
            "tokens_generated": tokens,
            "tokens_generated_source": token_count_source,
            "tokens_per_second": tps,
            "tokens_per_second_source": tps_source,
            "peak_memory_mb": mlx_peak_mb,
            "memory_delta_mb": max(0.0, mlx_peak_mb - active_memory_before_mb),
            "process_peak_rss_mb": process_peak_mb,
            "prompt_tokens": int(getattr(response, "prompt_tokens", 0) or 0),
            "prompt_tokens_per_second": float(
                getattr(response, "prompt_tps", 0.0) or 0.0
            ),
            "finish_reason": getattr(response, "finish_reason", None),
        }
        runs.append(run_result)
        print(f"  Run {i+1}: {elapsed:.2f}s | {tokens} tokens | {tps:.1f} tok/s | "
              f"{mlx_peak_mb:.0f} MB")

    # Aggregate
    import numpy as np
    tps_values = [r["tokens_per_second"] for r in runs]
    elapsed_values = [r["elapsed_seconds"] for r in runs]
    mem_values = [r["peak_memory_mb"] for r in runs]

    results = {
        "model_path": Path(model_path).name if Path(model_path).is_dir() else model_path,
        "served_revision": served_revision,
        "image_path": Path(image_path).name,
        "prompt": prompt,
        "max_tokens": max_tokens,
        "system": system_info,
        "model_load_time_seconds": load_time,
        "memory_after_load_mb": mem_after_load,
        "num_warmup": num_warmup,
        "num_runs": num_runs,
        "mean_tps": float(np.mean(tps_values)) if tps_values else 0,
        "std_tps": float(np.std(tps_values)) if tps_values else 0,
        "mean_elapsed_seconds": float(np.mean(elapsed_values)),
        "mean_peak_memory_mb": float(np.mean(mem_values)),
        "runs": runs,
    }

    return results


def main():
    parser = argparse.ArgumentParser(description="Run OCR performance benchmark")
    parser.add_argument("--model-path", required=True)
    parser.add_argument("--image", required=True, help="Test image path")
    parser.add_argument("--output", type=Path, default=Path("benchmarks/results/performance_results.json"))
    parser.add_argument("--prompt", default="<image>document parsing.")
    parser.add_argument("--max-tokens", type=int, default=4096)
    parser.add_argument("--warmup", type=int, default=1)
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument(
        "--served-revision",
        default=None,
        help="Immutable model revision used for this benchmark",
    )
    args = parser.parse_args()

    print("=" * 60)
    print("Unlimited-OCR Performance Benchmark")
    print("=" * 60)

    results = run_performance_benchmark(
        model_path=args.model_path,
        image_path=args.image,
        prompt=args.prompt,
        max_tokens=args.max_tokens,
        num_warmup=args.warmup,
        num_runs=args.runs,
        served_revision=args.served_revision,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8")

    print(f"\n{'='*60}")
    print(f"Mean TPS: {results['mean_tps']:.1f} | "
          f"Peak Memory: {results['mean_peak_memory_mb']:.0f} MB")
    print(f"Saved to: {args.output}")


if __name__ == "__main__":
    main()