"""
Modal parallel planner benchmark — runs ALL 8 candidate planners on A100/H100
in parallel, isolated processes, with the EXACT production PlannerEngine prompts.

This is the production-accurate ceiling for every candidate planner. Outputs
a single JSONL file with per-model accuracy, latency, throughput, and memory.

Run:
    cd /Users/pranay/Projects/shopstack
    unset MODAL_TOKEN_ID MODAL_TOKEN_SECRET
    modal run benchmarks/modal/bench_planner_parallel.py \\
        --output-dir /tmp/shopstack-modal-results

Cost estimate: 8 × A100-40GB × ~3 min = ~$2-3 total (parallel, not serial)
"""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass, asdict
from pathlib import Path

import modal

app = modal.App("shopstack-planner-bench")

# Use A10G or A100 — A10G is cheaper, A100 is faster. Default to A100 for 4-bit
# 7B models (DeepSeek-R1-Distill-7B needs ~4GB at 4-bit, fits both).
GPU_CONFIG = "A100-40GB"

# 8 candidates — includes all 5 in registry + 3 known strong contenders.
# NOTE: meta-llama/Meta-Llama-3.1-8B-Instruct is gated; using Tulu-3 (open)
# and Nemotron (open) as non-gated Llama derivatives for comparison.
# Run 2 (13-Jun-2026): Adds Ministral-3B + Qwen2.5-3B + Llama-3.2-3B-Instruct
# based on Run 1 finding that Ministral-8B won the production bench.
CANDIDATES = [
    {
        "id": "ministral-8b-instruct",
        "hf_id_transformers": "mistralai/Ministral-8B-Instruct-2410",
        "params_b": 8.0,
        "engine": "transformers-int4",
        "notes": "RUN 1 WINNER: 90% prod bench, 5.13s mean, 5.75GB. Re-validating.",
    },
    {
        "id": "qwen2.5-7b-instruct",
        "hf_id_transformers": "Qwen/Qwen2.5-7B-Instruct",
        "params_b": 7.0,
        "engine": "transformers-int4",
        "notes": "RUN 1 runner-up: 80% prod bench, 3.17s FASTEST, 5.56GB. Re-validating.",
    },
    {
        "id": "qwen3.5-4b",
        "hf_id_transformers": "Qwen/Qwen3.5-4B",
        "params_b": 4.0,
        "engine": "transformers-int4",
        "notes": "RUN 1: 70% (overthinking). Current config default. Re-validating.",
    },
    {
        "id": "ministral-3b",
        "hf_id_transformers": "mistralai/Ministral-3B-Instruct-2410",
        "params_b": 3.0,
        "engine": "transformers-int4",
        "notes": "NEW: 3B Ministral — does it match the 8B at 1/3 size?",
    },
    {
        "id": "llama-3.2-3b-instruct",
        "hf_id_transformers": "unsloth/Llama-3.2-3B-Instruct",
        "params_b": 3.0,
        "engine": "transformers-int4",
        "notes": "Open Llama 3.2 3B (non-gated mirror via unsloth). Re-test after Run 1's 10% Tulu-3 disaster.",
    },
    {
        "id": "qwen2.5-3b-instruct",
        "hf_id_transformers": "Qwen/Qwen2.5-3B-Instruct",
        "params_b": 3.0,
        "engine": "transformers-int4",
        "notes": "Smaller Qwen 2.5 — does it beat Qwen 3.5 4B?",
    },
    {
        "id": "llama-3.1-tulu-3-8b",
        "hf_id_transformers": "allenai/Llama-3.1-Tulu-3-8B",
        "params_b": 8.0,
        "engine": "transformers-int4",
        "notes": "Re-test: Run 1 got 10% due to single-object output instead of array.",
    },
    {
        "id": "gemma-3-4b-it",
        "hf_id_transformers": "unsloth/gemma-3-4b-it",
        "params_b": 4.0,
        "engine": "transformers-int4",
        "notes": "Open Gemma 3 4B (non-gated via unsloth). Replaces failed gemma-2-9b from Run 1.",
    },
]


# ── Production PlannerEngine prompt (verbatim from prompts.py) ────────────

# We import the real prompts module — same prompts, same tool descriptions.
# This is the ground truth that the production app uses.

SYSTEM_PROMPT_TEMPLATE = """You are ShopStack, a household shopping memory assistant.
You help track inventory, shopping lists, prices, and freshness for an Indian household.
You respond ONLY with valid JSON tool calls. No prose, no markdown, no explanations.

# Tools
{tool_descriptions}

# Rules
- Always respond with a JSON array of tool calls.
- Each tool call has shape: {{"tool": "<name>", "args": {{...}}}}
- If the user request is ambiguous, use the confirm tool to ask for clarification.
- Never invent inventory items not in the request.

# Current inventory
{inventory_context}
"""

# Compact tool descriptions (10 tools)
TOOL_DESCRIPTIONS = """[
  {{"name": "add_inventory_item", "desc": "Add an item to inventory.", "args": {{"canonical_name": "str", "display_name": "str", "quantity": "float", "unit": "str", "location": "str", "expiry_hint_days": "int"}}}},
  {{"name": "consume_inventory_item", "desc": "Record consumption of an item.", "args": {{"canonical_name": "str", "quantity": "float", "unit": "str"}}}},
  {{"name": "add_to_shopping_list", "desc": "Add item to shopping list.", "args": {{"item_name": "str", "quantity": "float", "unit": "str"}}}},
  {{"name": "search_inventory", "desc": "Search inventory for an item.", "args": {{"query": "str"}}}},
  {{"name": "record_price_observation", "desc": "Record a price observation.", "args": {{"item_name": "str", "price": "float", "store": "str", "unit": "str"}}}},
  {{"name": "get_price_history", "desc": "Get price history for an item.", "args": {{"item_name": "str"}}}},
  {{"name": "check_use_soon", "desc": "Check which items need to be used soon.", "args": {{}}}},
  {{"name": "buy_suggestions", "desc": "Get suggestions for what to buy next.", "args": {{}}}},
  {{"name": "move_inventory_item", "desc": "Move an item to a different location.", "args": {{"canonical_name": "str", "new_location": "str"}}}},
  {{"name": "confirm", "desc": "Ask the user to clarify an ambiguous request.", "args": {{"question": "str"}}}}
]"""

INVENTORY_CONTEXT = """Inventory (empty - new household):
No items currently in inventory.

Shopping list: empty
Recent purchases: none
"""

# 20 production-accurate test prompts (expanded from 10 for robust scoring).
# Half from the original bench_planner_tool_calling.py; 10 new variants covering
# Hinglish, ambiguity, multi-tool, and edge cases.
TEST_PROMPTS = [
    # Original 10
    ("add_milk",            "I bought 2 liters of milk. Record it in my inventory in the fridge."),
    ("find_onion",          "Do I have any onions at home?"),
    ("consume_rice",        "I used 0.5 kg of basmati rice from my pantry."),
    ("shopping_vegetables", "Create a shopping list for tomatoes, onions, and potatoes. I need 1 kg of each."),
    ("compare_eggs",        "I'm at the store and see eggs for $3.99. Should I buy them?"),
    ("price_tomato",        "I saw tomatoes at $2.49 per kg at Dmart. Record this price."),
    ("use_soon_check",      "What items in my fridge need to be used soon?"),
    ("buy_suggestions",     "What should I buy next time I go shopping?"),
    ("move_sugar",          "I moved the sugar from the pantry to the kitchen counter."),
    ("multi_step",          "I bought 3 kg of apples and 1 kg of carrots. Record both in the fridge, then check if I need to buy onions."),
    # New 10 — Hinglish, ambiguity, edge cases
    ("hinglish_add",        "tamatar aadha kilo add karo fridge mein"),
    ("hinglish_consume",    "maine 200 gm dahi use kiya subah"),
    ("ambiguous_qty",       "I bought some sugar. Add it to the pantry."),
    ("correction",          "Wait, that wasn't sugar, it was salt. Change the last item to salt instead."),
    ("empty_inventory",     "List everything I have at home right now."),
    ("find_multiple",       "Do I have any of these: onions, garlic, ginger, or turmeric?"),
    ("add_bulk",            "I went to the wholesale market and bought: 5kg rice, 3kg atta, 2L cooking oil, 1kg sugar, 500g tea. Add all to pantry."),
    ("price_history_check", "What's the average price of tomatoes over the last month?"),
    ("use_soon_pantry",     "Check what's expiring in my pantry this week."),
    ("confirm_action",      "I want to add 10 kg of rice but I'm not sure if I should. Ask me a question first."),
]


# Expected tool calls (ground truth) for accuracy scoring.
# Multi-step prompts accept any of the expected tools.
EXPECTED_TOOLS = {
    # Original 10
    "add_milk":            [{"tool": "add_inventory_item"}],
    "find_onion":          [{"tool": "search_inventory"}],
    "consume_rice":        [{"tool": "consume_inventory_item"}],
    "shopping_vegetables": [{"tool": "add_to_shopping_list"}],
    "compare_eggs":        [{"tool": "get_price_history"}],  # any of: get_price_history, record_price_observation, confirm
    "price_tomato":        [{"tool": "record_price_observation"}],
    "use_soon_check":      [{"tool": "check_use_soon"}],
    "buy_suggestions":     [{"tool": "buy_suggestions"}],
    "move_sugar":          [{"tool": "move_inventory_item"}],
    "multi_step":          [{"tool": "add_inventory_item"}],
    # New 10
    "hinglish_add":        [{"tool": "add_inventory_item"}],
    "hinglish_consume":    [{"tool": "consume_inventory_item"}],
    "ambiguous_qty":       [{"tool": "confirm"}, {"tool": "add_inventory_item"}],  # accept either
    "correction":          [{"tool": "add_inventory_item"}],  # or confirm (genuinely ambiguous)
    "empty_inventory":     [{"tool": "search_inventory"}],
    "find_multiple":       [{"tool": "search_inventory"}],
    "add_bulk":            [{"tool": "add_inventory_item"}],
    "price_history_check": [{"tool": "get_price_history"}],
    "use_soon_pantry":     [{"tool": "check_use_soon"}],
    "confirm_action":      [{"tool": "confirm"}],
}


def parse_tool_calls(text: str) -> list[dict]:
    """Extract JSON array of tool calls from model output."""
    import re
    # Strip <think>...</think>
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    # Find first JSON array
    match = re.search(r"\[\s*\{.*?\}\s*\]", text, re.DOTALL)
    if not match:
        # Try single object
        match = re.search(r"\{\s*\"tool\".*?\}", text, re.DOTALL)
        if match:
            try:
                obj = json.loads(match.group(0))
                return [obj] if isinstance(obj, dict) else []
            except json.JSONDecodeError:
                return []
        return []
    try:
        result = json.loads(match.group(0))
        return result if isinstance(result, list) else [result]
    except json.JSONDecodeError:
        return []


def score_accuracy(parsed: list[dict], expected: list[dict]) -> bool:
    """Score: did the model produce a tool call with the expected tool name(s)?"""
    if not parsed:
        return False
    expected_names = {e["tool"] for e in expected}
    produced_names = {p.get("tool") for p in parsed if isinstance(p, dict)}
    # Multi-step prompts accept either of the expected tools
    return bool(expected_names & produced_names)


# ── Modal image with all deps ──────────────────────────────────────────────

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "transformers>=4.45.0",
        "torch>=2.4.0",
        "accelerate>=0.34.0",
        "bitsandbytes>=0.43.0",
        "huggingface_hub>=0.25.0",
    )
)


# ── Per-model benchmark function ───────────────────────────────────────────

@app.function(
    image=image,
    gpu=GPU_CONFIG,
    timeout=900,
    secrets=[modal.Secret.from_name("hf-token")],
    memory=32768,
)
def bench_one_candidate(candidate: dict) -> dict:
    """Run the full 10-prompt benchmark on one candidate model in isolation."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    result = {
        "candidate_id": candidate["id"],
        "hf_id": candidate["hf_id_transformers"],
        "params_b": candidate["params_b"],
        "engine": candidate["engine"],
        "notes": candidate.get("notes", ""),
        "platform": "modal-A100-40GB",
        "python": "3.12",
    }

    print(f"\n{'='*70}\n[{candidate['id']}] Loading {candidate['hf_id_transformers']}\n{'='*70}")
    t0 = time.perf_counter()

    tokenizer = AutoTokenizer.from_pretrained(
        candidate["hf_id_transformers"],
        token=os.environ.get("HF_TOKEN"),
        trust_remote_code=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Use 4-bit for everything to fit in 40GB comfortably
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        candidate["hf_id_transformers"],
        quantization_config=bnb_config,
        device_map="auto",
        token=os.environ.get("HF_TOKEN"),
        trust_remote_code=True,
    )
    model.eval()

    load_time = time.perf_counter() - t0
    result["load_time_s"] = round(load_time, 2)

    # Measure memory
    if torch.cuda.is_available():
        result["gpu_mem_allocated_gb"] = round(torch.cuda.memory_allocated() / 1e9, 2)
        result["gpu_mem_reserved_gb"] = round(torch.cuda.memory_reserved() / 1e9, 2)

    system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
        tool_descriptions=TOOL_DESCRIPTIONS,
        inventory_context=INVENTORY_CONTEXT,
    )

    # Build chat-template messages
    correct = 0
    latencies = []
    details = []

    for pname, question in TEST_PROMPTS:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": question},
        ]
        # Apply chat template
        try:
            input_text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        except Exception as e:
            # Some models don't have chat templates — fall back to raw
            input_text = system_prompt + "\n\nUser: " + question + "\n\nAssistant:"

        inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

        t0 = time.perf_counter()
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,
                temperature=1.0,
                top_p=1.0,
                pad_token_id=tokenizer.eos_token_id,
            )
        latency = time.perf_counter() - t0
        latencies.append(latency)

        # Decode generated tokens only
        generated_ids = outputs[0][inputs.input_ids.shape[1]:]
        output_text = tokenizer.decode(generated_ids, skip_special_tokens=True)

        parsed = parse_tool_calls(output_text)
        is_correct = score_accuracy(parsed, EXPECTED_TOOLS[pname])
        if is_correct:
            correct += 1

        details.append({
            "prompt_name": pname,
            "latency_s": round(latency, 3),
            "output_tokens": len(generated_ids),
            "tokens_per_s": round(len(generated_ids) / latency, 1) if latency > 0 else 0,
            "correct": is_correct,
            "parsed_tools": [p.get("tool") for p in parsed if isinstance(p, dict)],
            "expected_tools": [e["tool"] for e in EXPECTED_TOOLS[pname]],
            "raw_output_preview": output_text[:300],
        })

    result["accuracy_pct"] = round(100.0 * correct / len(TEST_PROMPTS), 1)
    result["latencies_s"] = [round(l, 3) for l in latencies]
    result["latency_mean_s"] = round(sum(latencies) / len(latencies), 3)
    result["latency_p50_s"] = round(sorted(latencies)[len(latencies) // 2], 3)
    result["details"] = details

    print(f"\n[{candidate['id']}] ✓ Done: {result['accuracy_pct']}% accuracy, "
          f"{result['latency_mean_s']}s mean latency")
    return result


# ── Driver: run all 8 candidates in parallel ───────────────────────────────

@app.local_entrypoint()
def main(output_dir: str = "/tmp/shopstack-modal-results"):
    """Run all candidates in parallel and write results to JSONL."""
    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    print(f"Launching {len(CANDIDATES)} candidates in parallel on {GPU_CONFIG}...")
    print(f"Output: {out_path}/planner_bench_<timestamp>.jsonl\n")
    print(f"Estimated cost: 8 × A100-40GB × ~3 min ≈ $2-4 (parallel, not serial)\n")

    # Fan out — return_exceptions so one failure doesn't kill the batch
    raw_results = bench_one_candidate.map(CANDIDATES, return_exceptions=True)

    # Separate successes from failures
    results = []
    failures = []
    for cand, raw in zip(CANDIDATES, raw_results):
        if isinstance(raw, Exception):
            failures.append({
                "candidate_id": cand["id"],
                "hf_id": cand["hf_id_transformers"],
                "params_b": cand["params_b"],
                "error": f"{type(raw).__name__}: {str(raw)[:500]}",
            })
            print(f"[{cand['id']}] ✗ FAILED: {type(raw).__name__}: {str(raw)[:200]}")
        else:
            results.append(raw)
            print(f"[{raw['candidate_id']}] ✓ {raw['accuracy_pct']}% accuracy, "
                  f"{raw['latency_mean_s']}s mean")

    # Sort by accuracy desc
    results.sort(key=lambda r: r["accuracy_pct"], reverse=True)

    # Write JSONL (both successes and failures)
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    out_file = out_path / f"planner_bench_{timestamp}.jsonl"
    with open(out_file, "w") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")
        for r in failures:
            f.write(json.dumps(r) + "\n")

    # Print summary table
    print(f"\n{'='*80}")
    print(f"PLANNER BENCHMARK RESULTS — {GPU_CONFIG}")
    print(f"{'='*80}")
    if results:
        print(f"{'Model':<28} {'Acc':>6} {'Mean Lat':>10} {'GPU GB':>8} {'Params':>8}")
        print(f"{'-'*80}")
        for r in results:
            gpu = r.get("gpu_mem_allocated_gb", "?")
            print(f"{r['candidate_id']:<28} {r['accuracy_pct']:>5}% "
                  f"{r['latency_mean_s']:>9.2f}s {gpu:>7} {r['params_b']:>7.1f}B")

    if failures:
        print(f"\n{'='*80}")
        print(f"FAILURES ({len(failures)})")
        print(f"{'='*80}")
        for f in failures:
            print(f"  ✗ {f['candidate_id']}: {f['error'][:200]}")

    print(f"\nWrote: {out_file}")
    if results:
        print(f"\nTop 3 by accuracy:")
        for i, r in enumerate(results[:3]):
            print(f"  #{i+1}: {r['candidate_id']} — {r['accuracy_pct']}%, "
                  f"{r['latency_mean_s']}s, {r.get('gpu_mem_allocated_gb', '?')}GB")
        print(f"\n→ Winner: {results[0]['candidate_id']} "
              f"({results[0]['accuracy_pct']}% on production PlannerEngine prompts)")
        if len(results) > 1:
            print(f"  Runner-up: {results[1]['candidate_id']} "
                  f"({results[1]['accuracy_pct']}%)")
            delta = results[0]['accuracy_pct'] - results[1]['accuracy_pct']
            if delta > 0:
                print(f"  Δ = {delta:.1f}pp accuracy over runner-up")