"""
Project David — vLLM Raw Token Stream Profiler
===============================================
Targets /v1/completions (NOT /v1/chat/completions).
This bypasses vLLM's chat template + tool-call parser entirely.
You get the raw generated text exactly as the model emits it.

Usage:
------
1. Start vLLM server (in a separate terminal):
       vllm serve Qwen/Qwen2.5-0.5B-Instruct --port 8000

2. Run this script:
       python profile_vllm_raw.py

Swap MODEL_ID and rendered prompt to cycle through families.
"""

import json

import httpx

VLLM_BASE = "http://localhost:8000"

# ── Model profiles — swap these to cycle through families ─────────────────────

PROFILES = [
    {
        "name": "Qwen2.5-0.5B-Instruct",
        "model_id": "Qwen/Qwen2.5-0.5B-Instruct",
        "prompt": (
            "<|im_start|>system\n"
            "You are a helpful assistant with access to tools.\n\n"
            "## Tools\n\n"
            '{"type":"function","function":{"name":"get_weather","description":"Get weather.","parameters":{"type":"object","properties":{"location":{"type":"string"}},"required":["location"]}}}\n'
            "<|im_end|>\n"
            "<|im_start|>user\nWhat is the weather in London?\n<|im_end|>\n"
            "<|im_start|>assistant\n"
        ),
    },
    {
        "name": "Mistral-7B-Instruct-v0.3",
        "model_id": "mistralai/Mistral-7B-Instruct-v0.3",
        "prompt": (
            "<s>[INST] You have access to the following tools:\n"
            '[{"type":"function","function":{"name":"get_weather","parameters":{"type":"object","properties":{"location":{"type":"string"}},"required":["location"]}}}]\n\n'
            "What is the weather in London? [/INST]"
        ),
    },
    {
        "name": "Llama-3.2-1B-Instruct",
        "model_id": "meta-llama/Llama-3.2-1B-Instruct",
        "prompt": (
            "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
            "You are a helpful assistant with access to tools.\n"
            'Tools: [{"name": "get_weather", "parameters": {"location": {"type": "string"}}}]'
            "<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"
            "What is the weather in London?<|eot_id|>"
            "<|start_header_id|>assistant<|end_header_id|>\n\n"
        ),
    },
    # Add more families as needed:
    # Gemma3, Phi-3, DeepSeek-V3, Command-R...
]


# ══════════════════════════════════════════════════════════════════════════════
# Health check
# ══════════════════════════════════════════════════════════════════════════════


def check_server():
    try:
        r = httpx.get(f"{VLLM_BASE}/health", timeout=5)
        if r.status_code == 200:
            print("✅ vLLM server is up\n")
            return True
    except Exception:
        pass
    print("❌ vLLM server not reachable at", VLLM_BASE)
    print("   Start it with: vllm serve <model_id> --port 8000")
    return False


def list_loaded_models():
    try:
        r = httpx.get(f"{VLLM_BASE}/v1/models", timeout=5)
        models = r.json().get("data", [])
        print("Loaded models:")
        for m in models:
            print(f"  - {m['id']}")
        print()
        return [m["id"] for m in models]
    except Exception as e:
        print(f"Could not list models: {e}")
        return []


# ══════════════════════════════════════════════════════════════════════════════
# Core profiler — hits /v1/completions (raw text, no chat wrapper)
# ══════════════════════════════════════════════════════════════════════════════


def profile_vllm_raw(profile: dict):
    name = profile["name"]
    model_id = profile["model_id"]
    prompt = profile["prompt"]

    print("═" * 60)
    print(f"PROFILING: {name}")
    print(f"Endpoint:  {VLLM_BASE}/v1/completions")
    print("═" * 60)
    print("\nRendered prompt sent to model:")
    print(prompt)
    print()

    payload = {
        "model": model_id,
        "prompt": prompt,
        "max_tokens": 300,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": False},
        # Ensure special tokens are included in output
        "skip_special_tokens": False,
    }

    raw_chunks = []
    chunk_index = 0

    print("RAW SSE STREAM (/v1/completions):\n")

    with httpx.Client(timeout=120) as http:
        with http.stream(
            "POST",
            f"{VLLM_BASE}/v1/completions",
            json=payload,
            headers={"Content-Type": "application/json"},
        ) as resp:

            print(f"HTTP {resp.status_code}\n")

            if resp.status_code != 200:
                print(f"ERROR: {resp.read().decode()}")
                return

            for line in resp.iter_lines():
                if not line:
                    continue
                if not line.startswith("data:"):
                    print(f"SSE meta: {repr(line)}")
                    continue

                raw = line[5:].strip()
                if raw == "[DONE]":
                    print("\n[DONE]")
                    break

                try:
                    parsed = json.loads(raw)
                    # /v1/completions delta is in choices[0].text
                    text = parsed["choices"][0].get("text", "")
                    finish = parsed["choices"][0].get("finish_reason")

                    raw_chunks.append(text)
                    print(f"[{chunk_index:03d}] {repr(text)}", end="")
                    if finish:
                        print(f"  ← finish_reason={finish}", end="")
                    print()
                    chunk_index += 1

                except Exception as e:
                    print(f"  parse error: {e} | raw={repr(raw)}")

    full_output = "".join(raw_chunks)

    # ── Structural analysis ──────────────────────────────────────────────────

    print("\n" + "─" * 60)
    print("FULL OUTPUT:")
    print(full_output)

    print("\n" + "─" * 60)
    print("MARKER SCAN:")

    markers = [
        "<tool_call>",
        "</tool_call>",  # Qwen
        "[TOOL_CALLS]",  # Mistral
        "<|python_tag|>",  # Llama
        "```python",
        "```json",  # Gemma / Cohere
        "Action:",  # Cohere
        "<|tool▁calls▁begin|>",  # DeepSeek
        "<|tool_call|>",  # Phi
        "<think>",
        "</think>",  # Qwen3 reasoning
        '{"name":',
        '"arguments":',
        '"parameters":',
    ]

    for m in markers:
        idx = full_output.find(m)
        if idx != -1:
            ctx = full_output[max(0, idx - 5) : idx + len(m) + 30]
            print(f"  ✓ {repr(m):30s} at char {idx:4d} → {repr(ctx)}")
        else:
            print(f"  ✗ {repr(m)}")

    print("\n" + "─" * 60)
    print("CHUNK BOUNDARY ANALYSIS:")
    print("(Where do structural markers land relative to chunk edges?)\n")

    running = ""
    for i, chunk in enumerate(raw_chunks):
        prev = len(running)
        running += chunk
        for m in ["<tool_call>", "[TOOL_CALLS]", "<|python_tag|>", "<think>"]:
            pos = running.find(m)
            if pos != -1 and pos >= prev:
                print(
                    f"  Marker {repr(m)} first seen in chunk [{i:03d}]: {repr(chunk)}"
                )
                surrounding = raw_chunks[max(0, i - 1) : i + 3]
                print(f"  Context chunks: {surrounding}")

    # ── Save ────────────────────────────────────────────────────────────────

    out = {
        "model": name,
        "endpoint": "vllm_v1_completions",
        "full_output": full_output,
        "chunk_count": len(raw_chunks),
        "chunks": raw_chunks,
        "markers": {
            m: full_output.find(m) for m in markers if full_output.find(m) != -1
        },
    }

    fname = f"vllm_profile_{name.replace('/', '_').replace('-', '_')}.json"
    with open(fname, "w", encoding="utf-8") as f:
        json.dump(out, f, indent=2, ensure_ascii=False)
    print(f"\n✅ Saved: {fname}\n")


# ══════════════════════════════════════════════════════════════════════════════
# Run
# ══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    if not check_server():
        exit(1)

    loaded = list_loaded_models()

    for profile in PROFILES:
        # Only profile models that are actually loaded in the server
        if loaded and not any(profile["model_id"] in m for m in loaded):
            print(f"⚠️  Skipping {profile['name']} — not loaded in vLLM server")
            print(f"   Run: vllm serve {profile['model_id']} --port 8000\n")
            continue
        profile_vllm_raw(profile)

    print("✅ All vLLM profiles complete.")