"""R-SWA (Reference Sliding Window Attention) cache validation.

Instruments MLX KV-cache behavior at increasing output lengths to verify
whether the model's R-SWA mechanism properly bounds cache growth.

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

from __future__ import annotations

import argparse
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"))

from unlimited_ocr.engine import SlidingWindowNoRepeatNGramProcessor  # noqa: E402


# Output lengths to test
TEST_LENGTHS = [512, 2048, 4096, 8192, 16384]


class MinGeneratedTokensProcessor:
    """Suppress EOS until a requested number of new tokens has been sampled.

    ``mlx-vlm`` passes prompt tokens to logits processors as well as generated
    tokens.  The first observed token count is therefore recorded as the prompt
    length, keeping this stress-test control independent of prompt tokenization.
    A fresh processor must be constructed for every generation call.
    """

    def __init__(self, min_new_tokens: int, eos_token_ids: list[int]):
        if (
            not isinstance(min_new_tokens, int)
            or isinstance(min_new_tokens, bool)
            or min_new_tokens < 1
        ):
            raise ValueError("min_new_tokens must be a positive integer")
        normalized_ids = sorted({
            int(token_id)
            for token_id in eos_token_ids
            if isinstance(token_id, Integral)
            and not isinstance(token_id, bool)
            and token_id >= 0
        })
        if not normalized_ids:
            raise ValueError("At least one integer EOS token ID is required")
        self.min_new_tokens = min_new_tokens
        self.eos_token_ids = normalized_ids
        self._prompt_token_count: int | None = None

    def __call__(self, tokens, logits):
        sequence = tokens.tolist() if hasattr(tokens, "tolist") else list(tokens)
        if sequence and isinstance(sequence[0], list):
            sequence = sequence[0]
        if self._prompt_token_count is None:
            self._prompt_token_count = len(sequence)
        generated_tokens = max(0, len(sequence) - self._prompt_token_count)
        if generated_tokens < self.min_new_tokens:
            logits[:, self.eos_token_ids] = float("-inf")
        return logits


def collect_eos_token_ids(model, processor) -> list[int]:
    """Collect EOS IDs from model config and the loaded tokenizer wrapper."""
    values = []
    model_config = getattr(model, "config", None)
    values.append(getattr(model_config, "eos_token_id", None))

    tokenizer = getattr(processor, "tokenizer", processor)
    values.append(getattr(tokenizer, "eos_token_ids", None))
    values.append(getattr(tokenizer, "eos_token_id", None))
    stopping_criteria = getattr(tokenizer, "stopping_criteria", None)
    values.append(getattr(stopping_criteria, "eos_token_ids", None))

    normalized: set[int] = set()
    for value in values:
        candidates = value if isinstance(value, (list, tuple, set)) else [value]
        for candidate in candidates:
            if (
                isinstance(candidate, Integral)
                and not isinstance(candidate, bool)
                and candidate >= 0
            ):
                normalized.add(int(candidate))
    if not normalized:
        raise ValueError("Could not determine EOS token IDs from model or processor")
    return sorted(normalized)


def detect_repetition_rate(text: str, ngram_size: int = 35, threshold: int = 3) -> float:
    """Detect fraction of repeated n-grams in output."""
    if not isinstance(text, str):
        raise TypeError("text must be a string")
    if not isinstance(ngram_size, int) or isinstance(ngram_size, bool) or ngram_size < 1:
        raise ValueError("ngram_size must be a positive integer")
    if not isinstance(threshold, int) or isinstance(threshold, bool) or threshold < 1:
        raise ValueError("threshold must be a positive integer")
    words = text.split()
    if len(words) < ngram_size:
        return 0.0

    from collections import Counter
    ngrams = [tuple(words[i:i+ngram_size]) for i in range(len(words) - ngram_size + 1)]
    counts = Counter(ngrams)
    repeated = sum(1 for c in counts.values() if c > threshold)
    return repeated / len(counts) if counts else 0.0


def analyze_cache_behavior(results: list[dict]) -> tuple[bool | None, str]:
    """Conservatively classify cache growth from successful length-limited runs.

    Returns ``None`` when the generated token counts do not span enough lengths
    to distinguish a plateau from ordinary noise or linear growth.
    """
    successful = []
    for result in results:
        if not isinstance(result, dict) or result.get("status") != "success":
            continue
        tokens = result.get("tokens_generated")
        memory = result.get("memory_growth_mb")
        if (
            not isinstance(tokens, Integral)
            or isinstance(tokens, bool)
            or tokens <= 0
            or result.get("tokens_generated_reliable", True) is not True
            or not isinstance(memory, Real)
            or isinstance(memory, bool)
            or not math.isfinite(float(memory))
            or memory < 0
        ):
            continue
        successful.append(result)
    # Collapse duplicate actual token counts; raising max_tokens is not useful
    # evidence when the model naturally stops at the same length each time.
    by_tokens: dict[int, float] = {}
    for result in successful:
        tokens = int(result["tokens_generated"])
        memory = float(result["memory_growth_mb"])
        by_tokens[tokens] = max(by_tokens.get(tokens, 0.0), memory)
    points = sorted(by_tokens.items())

    if len(points) < 3:
        return None, "Need at least three distinct generated lengths"
    if points[-1][0] < points[0][0] * 2:
        return None, "Generated token counts span less than 2x"

    # Peak-memory readings can fluctuate slightly, so enforce their cumulative
    # maximum before comparing early and late slopes.
    monotonic_memory = []
    running_peak = 0.0
    for tokens, memory in points:
        running_peak = max(running_peak, float(memory))
        monotonic_memory.append((tokens, running_peak))

    midpoint = len(monotonic_memory) // 2
    first_tokens, first_memory = monotonic_memory[0]
    middle_tokens, middle_memory = monotonic_memory[midpoint]
    last_tokens, last_memory = monotonic_memory[-1]
    early_span = middle_tokens - first_tokens
    late_span = last_tokens - middle_tokens
    if early_span <= 0 or late_span <= 0:
        return None, "Generated lengths are not sufficiently distinct"

    total_memory_change = last_memory - first_memory
    # A wide token span with <=64 MB of additional peak memory is direct
    # plateau evidence even if the early slope is effectively zero.
    if last_tokens >= first_tokens * 4 and total_memory_change <= 64.0:
        return True, "Peak memory stayed within 64 MB across a >=4x token span"

    early_slope = (middle_memory - first_memory) / early_span
    late_slope = (last_memory - middle_memory) / late_span
    if early_slope <= 0:
        return None, "Memory variation is too small to estimate a growth trend"

    if late_slope <= early_slope * 0.25:
        return True, "Late memory-growth slope is <=25% of the early slope"
    return False, "Memory growth does not plateau at longer generated lengths"


def run_rswa_validation(
    model_path: str,
    image_path: str,
    prompt: str = "<image>document parsing.",
    test_lengths: list[int] | None = None,
    force_min_tokens: int | None = None,
) -> dict:
    """Run R-SWA cache validation at increasing output lengths.

    Measures:
    - Peak memory at each output length
    - Decode TPS in first vs last quarter
    - Cache growth behavior
    - Repetition rate
    - Completion status
    """
    if test_lengths is None:
        test_lengths = list(TEST_LENGTHS)
    else:
        test_lengths = list(test_lengths)
    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 test_lengths or any(
        not isinstance(length, int) or isinstance(length, bool) or length < 1
        for length in test_lengths
    ):
        raise ValueError("test_lengths must contain positive integers")
    if len(set(test_lengths)) != len(test_lengths):
        raise ValueError("test_lengths must not contain duplicates")
    if test_lengths != sorted(test_lengths):
        raise ValueError("test_lengths must be in increasing order")
    if len(test_lengths) < 3:
        raise ValueError("test_lengths must contain at least three lengths")
    if 8192 not in test_lengths:
        raise ValueError("test_lengths must include 8192 for the required 8K check")
    if test_lengths[-1] < test_lengths[0] * 2:
        raise ValueError("test_lengths must span at least 2x")
    if force_min_tokens is not None and (
        not isinstance(force_min_tokens, int)
        or isinstance(force_min_tokens, bool)
        or force_min_tokens < 1
    ):
        raise ValueError("force_min_tokens must be a positive integer or None")

    from mlx_vlm import load, generate
    import mlx.core as mx

    print(f"Loading model: {model_path}")
    model, processor = load(model_path)
    eos_token_ids = (
        collect_eos_token_ids(model, processor)
        if force_min_tokens is not None
        else []
    )

    baseline_memory = mx.get_active_memory() / 1e6
    if not math.isfinite(baseline_memory) or baseline_memory < 0:
        raise ValueError("MLX returned invalid active memory")
    print(f"Baseline memory: {baseline_memory:.0f} MB")

    results = []

    for max_tokens in test_lengths:
        print(f"\n--- Testing max_tokens={max_tokens} ---")

        mx.clear_cache()
        mem_before = mx.get_active_memory() / 1e6
        if not math.isfinite(mem_before) or mem_before < 0:
            raise ValueError("MLX returned invalid active memory")
        mx.reset_peak_memory()
        start = time.perf_counter()

        try:
            forced_minimum = (
                min(force_min_tokens, max_tokens)
                if force_min_tokens is not None
                else None
            )
            logits_processors = [
                SlidingWindowNoRepeatNGramProcessor(35, 128)
            ]
            if forced_minimum is not None:
                logits_processors.append(
                    MinGeneratedTokensProcessor(forced_minimum, eos_token_ids)
                )
            response = generate(
                model, processor,
                prompt=prompt,
                image=[image_path],
                max_tokens=max_tokens,
                temperature=0.0,
                logits_processors=logits_processors,
                verbose=False,
            )
            elapsed = time.perf_counter() - start
            peak_memory = (
                response.get("peak_memory")
                if isinstance(response, dict)
                else getattr(response, "peak_memory", None)
            )
            if peak_memory is not None and (
                not isinstance(peak_memory, Real)
                or isinstance(peak_memory, bool)
                or not math.isfinite(float(peak_memory))
                or peak_memory < 0
            ):
                raise ValueError("mlx-vlm returned invalid peak_memory")
            mem_after = (
                float(peak_memory) * 1000
                if peak_memory is not None and peak_memory > 0
                else mx.get_peak_memory() / 1e6
            )
            if not math.isfinite(mem_after) or mem_after < 0:
                raise ValueError("MLX returned invalid peak memory")

            # Extract text
            if isinstance(response, str):
                text = response
                tokens_generated = len(text.split())
                token_count_source = "whitespace estimate"
                tokens_reliable = False
            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_generated = len(text.split())
                    token_count_source = "whitespace estimate"
                    tokens_reliable = False
                else:
                    tokens_generated = raw_tokens
                    token_count_source = "mlx-vlm token count"
                    tokens_reliable = True
            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_generated = len(text.split())
                    token_count_source = "whitespace estimate"
                    tokens_reliable = False
                else:
                    tokens_generated = raw_tokens
                    token_count_source = "mlx-vlm token count"
                    tokens_reliable = True
            else:
                text = str(response)
                tokens_generated = len(text.split())
                token_count_source = "whitespace estimate"
                tokens_reliable = False

            # Compute metrics
            if (
                not isinstance(tokens_generated, Integral)
                or isinstance(tokens_generated, bool)
                or not 0 <= tokens_generated <= max_tokens
            ):
                raise ValueError("mlx-vlm returned an invalid generation token count")
            tokens_generated = int(tokens_generated)
            reported_tps = (
                response.get("generation_tps", 0.0)
                if isinstance(response, dict)
                else getattr(response, "generation_tps", 0.0)
            )
            if isinstance(reported_tps, bool) or (
                reported_tps is not None
                and (
                    not isinstance(reported_tps, Real)
                    or not math.isfinite(float(reported_tps))
                    or reported_tps < 0
                )
            ):
                raise ValueError("mlx-vlm returned invalid generation_tps")
            tps = (
                float(reported_tps)
                if reported_tps
                else tokens_generated / elapsed if elapsed > 0 else 0.0
            )
            repetition = detect_repetition_rate(text)
            memory_growth = max(0.0, mem_after - mem_before)
            finish_reason = (
                response.get("finish_reason")
                if isinstance(response, dict)
                else getattr(response, "finish_reason", None)
            )
            if finish_reason is not None:
                completed = finish_reason != "length"
            elif tokens_reliable:
                completed = tokens_generated < max_tokens
            else:
                completed = None

            result = {
                "max_tokens": max_tokens,
                "tokens_generated": tokens_generated,
                "tokens_generated_source": token_count_source,
                "tokens_generated_reliable": tokens_reliable,
                "elapsed_seconds": elapsed,
                "mean_tps": tps,
                "tps_source": (
                    "mlx-vlm generation_tps" if reported_tps else "end-to-end estimate"
                ),
                "peak_memory_mb": mem_after,
                "memory_growth_mb": memory_growth,
                "repetition_rate": repetition,
                "completed_naturally": completed,
                "finish_reason": finish_reason,
                "forced_minimum_tokens": forced_minimum,
                "status": "success",
            }

            print(f"  Tokens: {tokens_generated} | TPS: {tps:.1f} | "
                  f"Memory: {mem_after:.0f} MB (+{memory_growth:.0f}) | "
                  f"Repetition: {repetition:.4f}")

        except Exception as e:
            elapsed = time.perf_counter() - start
            result = {
                "max_tokens": max_tokens,
                "status": "error",
                "error": str(e),
                "elapsed_seconds": elapsed,
            }
            print(f"  ERROR: {e}")

        results.append(result)

    # Analyze cache behavior
    successful = [r for r in results if r["status"] == "success"]
    reliable_successful = [
        r for r in successful if r.get("tokens_generated_reliable", True)
    ]
    cache_bounded, cache_reason = analyze_cache_behavior(results)

    # TPS degradation check
    tps_stable: bool | None = None
    distinct_tps_runs = sorted(
        reliable_successful,
        key=lambda result: result.get("tokens_generated", 0),
    )
    if (
        len({r.get("tokens_generated", 0) for r in distinct_tps_runs}) >= 2
        and distinct_tps_runs[-1].get("tokens_generated", 0)
        >= distinct_tps_runs[0].get("tokens_generated", 0) * 2
    ):
        first_tps = distinct_tps_runs[0].get("mean_tps", 0)
        last_tps = distinct_tps_runs[-1].get("mean_tps", 0)
        if first_tps > 0:
            tps_stable = last_tps >= first_tps * 0.7  # Within 30% degradation

    test_8k = next((r for r in results if r.get("max_tokens") == 8192), None)
    test_8k_passed = bool(
        test_8k
        and test_8k.get("status") == "success"
        and test_8k.get("tokens_generated_reliable", True)
        and test_8k.get("tokens_generated", 0) >= int(8192 * 0.95)
    )

    summary = {
        "model_path": Path(model_path).name if Path(model_path).is_dir() else model_path,
        "image_path": Path(image_path).name,
        "prompt": prompt,
        "force_min_tokens": force_min_tokens,
        "generation_settings": {
            "temperature": 0.0,
            "no_repeat_ngram_size": 35,
            "ngram_window": 128,
        },
        "forced_eos_token_ids": eos_token_ids,
        "platform": platform.platform(),
        "processor": platform.processor(),
        "baseline_memory_mb": baseline_memory,
        "test_results": results,
        "analysis": {
            "cache_appears_bounded": cache_bounded,
            "cache_analysis_reason": cache_reason,
            "tps_stable_over_length": tps_stable,
            "max_successful_tokens": max(
                (r["tokens_generated"] for r in reliable_successful), default=0
            ),
            "max_repetition_rate": max(
                (r.get("repetition_rate", 0) for r in successful), default=0
            ),
        },
        "pass_conditions": {
            "cache_bounded": cache_bounded is True,
            "tps_stable": tps_stable is True,
            "8k_test_passed": test_8k_passed,
        },
    }
    summary["passed"] = all(summary["pass_conditions"].values())

    return summary


def main():
    parser = argparse.ArgumentParser(description="R-SWA cache validation")
    parser.add_argument("--model-path", required=True)
    parser.add_argument("--image", required=True, help="Test image (should produce long output)")
    parser.add_argument("--output", type=Path, default=Path("benchmarks/results/rswa_results.json"))
    parser.add_argument("--prompt", default="<image>document parsing.")
    parser.add_argument("--lengths", type=int, nargs="+", default=None,
                        help="Custom output lengths to test")
    parser.add_argument(
        "--force-min-tokens",
        type=int,
        default=None,
        help=(
            "Stress-test only: suppress EOS until this many new tokens "
            "(capped at each --lengths value)"
        ),
    )
    args = parser.parse_args()

    print("=" * 60)
    print("R-SWA Cache Validation — Unlimited-OCR MLX")
    print("=" * 60)

    results = run_rswa_validation(
        model_path=args.model_path,
        image_path=args.image,
        prompt=args.prompt,
        test_lengths=args.lengths,
        force_min_tokens=args.force_min_tokens,
    )

    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 summary
    analysis = results["analysis"]
    pass_cond = results["pass_conditions"]

    print(f"\n{'='*60}")
    print("R-SWA Validation Summary:")
    cache_status = (
        "PASS" if analysis["cache_appears_bounded"] is True
        else "FAIL" if analysis["cache_appears_bounded"] is False
        else "INCONCLUSIVE"
    )
    tps_status = (
        "PASS" if analysis["tps_stable_over_length"] is True
        else "FAIL" if analysis["tps_stable_over_length"] is False
        else "INCONCLUSIVE"
    )
    print(f"  Cache bounded:     {cache_status}")
    print(f"  TPS stable:        {tps_status}")
    print(f"  8K test:           {'PASS' if pass_cond['8k_test_passed'] else 'FAIL'}")
    print(f"  Max tokens:        {analysis['max_successful_tokens']}")
    print(f"  Max repetition:    {analysis['max_repetition_rate']:.4f}")
    print(f"\nSaved to: {args.output}")
    if not results["passed"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()