import gradio as gr
import spaces
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import threading
import time
import json

# ─── Phase Configuration ───
PHASE = "Phase 6: Ultimate Combined (ZeroGPU)"
MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
MODEL_CONFIG = {
    "phase": PHASE,
    "model_name": MODEL_NAME,
    "torch_dtype": "float16",
    "quantization": "none",
    "optimization": "sdpa-fa-greedy-static-kv",
    "hardware": "zero-a10g",
    "max_new_tokens": 512,
    "do_sample": False,
    "cache_implementation": "static",
}

# ─── Load model with SDPA attention ───
print("Loading model with SDPA attention...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
    low_cpu_mem_usage=True,
)
print("Model loaded successfully with SDPA! (Ultimate Combined: FA backend + Greedy + Static KV)", flush=True)

# Track whether static cache works on this environment
_static_cache_available = None


@spaces.GPU
def generate_response(message, history_tuples=None):
    """Core generation logic, returns response + metrics."""
    global _static_cache_available

    # Move model to GPU (ZeroGPU provides GPU only inside @spaces.GPU)
    model.to("cuda")

    messages = []
    if history_tuples:
        for user_msg, assistant_msg in history_tuples:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    )
    # apply_chat_template may return a tensor or BatchEncoding depending on version
    if hasattr(input_ids, "input_ids"):
        input_ids = input_ids.input_ids
    input_ids = input_ids.to("cuda")

    input_tokens = input_ids.shape[1]

    start_time = time.time()
    with torch.no_grad():
        # Force SDPA to use FlashAttention backend (Phase 5a technique)
        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
            # Try static cache first, fall back to greedy-only if it fails (Phase 5b technique)
            if _static_cache_available is not False:
                try:
                    outputs = model.generate(
                        input_ids,
                        max_new_tokens=MODEL_CONFIG["max_new_tokens"],
                        do_sample=False,
                        cache_implementation="static",
                        pad_token_id=tokenizer.eos_token_id,
                    )
                    _static_cache_available = True
                except Exception as e:
                    print(f"Static cache failed ({type(e).__name__}: {e}), falling back to greedy only", flush=True)
                    _static_cache_available = False
                    outputs = model.generate(
                        input_ids,
                        max_new_tokens=MODEL_CONFIG["max_new_tokens"],
                        do_sample=False,
                        pad_token_id=tokenizer.eos_token_id,
                    )
            else:
                outputs = model.generate(
                    input_ids,
                    max_new_tokens=MODEL_CONFIG["max_new_tokens"],
                    do_sample=False,
                    pad_token_id=tokenizer.eos_token_id,
                )
    inference_time = time.time() - start_time

    output_tokens = outputs.shape[1] - input_tokens
    response = tokenizer.decode(outputs[0][input_tokens:], skip_special_tokens=True)
    tokens_per_sec = round(output_tokens / inference_time, 2) if inference_time > 0 else 0

    cache_status = "static" if _static_cache_available else "dynamic (fallback)"

    return {
        "response": response,
        "inference_time_s": round(inference_time, 2),
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "tokens_per_sec": tokens_per_sec,
        "model_config": {**MODEL_CONFIG, "cache_actual": cache_status},
    }


def _run_generate(kwargs):
    """Run model.generate in a thread with FA backend. No fallback — avoids corrupting the streamer."""
    with torch.no_grad():
        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
            model.generate(**kwargs)


@spaces.GPU
def generate_streaming(message, history_tuples=None):
    """Streaming generation — yields partial text chunks, then final metrics JSON."""
    global _static_cache_available

    model.to("cuda")

    messages = []
    if history_tuples:
        for user_msg, assistant_msg in history_tuples:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    )
    if hasattr(input_ids, "input_ids"):
        input_ids = input_ids.input_ids
    input_ids = input_ids.to("cuda")

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    generate_kwargs = {
        "input_ids": input_ids,
        "max_new_tokens": MODEL_CONFIG["max_new_tokens"],
        "do_sample": False,
        "pad_token_id": tokenizer.eos_token_id,
        "streamer": streamer,
    }

    # Only use static cache if KNOWN to work (True). When None (unknown/first call),
    # skip it — fallback in _run_generate would corrupt the streamer's skip_prompt state.
    if _static_cache_available is True:
        generate_kwargs["cache_implementation"] = "static"

    start_time = time.time()
    thread = threading.Thread(target=_run_generate, args=(generate_kwargs,))
    thread.start()

    generated_text = ""
    for chunk in streamer:
        generated_text += chunk
        yield chunk

    thread.join()
    inference_time = time.time() - start_time

    output_tokens = len(tokenizer.encode(generated_text))
    tokens_per_sec = round(output_tokens / inference_time, 2) if inference_time > 0 else 0
    cache_status = "static" if _static_cache_available else "dynamic (fallback)"

    yield json.dumps({
        "__metrics__": True,
        "inference_time_s": round(inference_time, 2),
        "output_tokens": output_tokens,
        "tokens_per_sec": tokens_per_sec,
        "model_config": {**MODEL_CONFIG, "cache_actual": cache_status},
    })


def parse_history(history):
    """Convert Gradio 5 history format to tuples."""
    if not history:
        return None
    tuples = []
    i = 0
    while i < len(history):
        item = history[i]
        if isinstance(item, dict):
            if item.get("role") == "user":
                user_msg = item.get("content", "")
                asst_msg = ""
                if i + 1 < len(history):
                    next_item = history[i + 1]
                    if isinstance(next_item, dict) and next_item.get("role") == "assistant":
                        asst_msg = next_item.get("content", "")
                        i += 1
                tuples.append((user_msg, asst_msg))
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            tuples.append(tuple(item))
        i += 1
    return tuples if tuples else None


# ─── Gradio Chat (for HF Spaces UI) ───
def chat(message, history):
    history_tuples = parse_history(history)
    result = generate_response(message, history_tuples)
    cache_info = result["model_config"].get("cache_actual", "unknown")
    timing = f"\n\n---\n*Inference: {result['inference_time_s']}s | {result['tokens_per_sec']} t/s | Ultimate: SDPA+FA+Greedy+{cache_info} cache*"
    return result["response"] + timing


# ─── API Endpoint (for React app + benchmark) ───
def api_chat(message, history_json="[]"):
    try:
        if not history_json or history_json.strip() == "":
            history_json = "[]"
        history = json.loads(history_json) if isinstance(history_json, str) else history_json
        history_tuples = [tuple(h) for h in history] if history else None
        result = generate_response(message, history_tuples)
        return json.dumps(result)
    except Exception as e:
        import traceback
        return json.dumps({"error": str(e), "traceback": traceback.format_exc()})


# ─── Streaming API Endpoint ───
def api_chat_stream(message, history_json="[]"):
    """Streaming API — yields text chunks, then final metrics JSON."""
    try:
        if not history_json or history_json.strip() == "":
            history_json = "[]"
        history = json.loads(history_json) if isinstance(history_json, str) else history_json
        history_tuples = [tuple(h) for h in history] if history else None
        for chunk in generate_streaming(message, history_tuples):
            yield chunk
    except Exception as e:
        import traceback
        yield json.dumps({"__error__": True, "error": str(e), "traceback": traceback.format_exc()})


# ─── Build Gradio App ───
with gr.Blocks() as demo:
    gr.Markdown(f"# Phi-3 Mini Chatbot ({PHASE})")
    gr.Markdown("Chat UI + API endpoint for benchmarking | SDPA + FlashAttention backend + Greedy decoding + Static KV cache")

    with gr.Tab("Chat"):
        chatbot = gr.ChatInterface(fn=chat)

    with gr.Tab("API"):
        gr.Markdown("""
### API Endpoints

**Non-streaming** (`/gradio_api/call/api_chat`):
```
POST /gradio_api/call/api_chat
{"data": ["your question", "[]"]}
→ returns {"event_id": "..."}

GET /gradio_api/call/api_chat/{event_id}
→ SSE stream with data: [json_result]
```

**Streaming** (`/gradio_api/call/api_chat_stream`):
```
POST /gradio_api/call/api_chat_stream
{"data": ["your question", "[]"]}
→ returns {"event_id": "..."}

GET /gradio_api/call/api_chat_stream/{event_id}
→ SSE stream with data: ["token_chunk"] per token, final chunk has __metrics__
```
        """)
        msg_input = gr.Textbox(label="Message", placeholder="Type your question...")
        history_input = gr.Textbox(label="History (JSON)", value="[]", visible=False)
        api_output = gr.Textbox(label="API Response (JSON)", lines=10)
        api_btn = gr.Button("Call API (non-streaming)")
        api_btn.click(
            fn=api_chat,
            inputs=[msg_input, history_input],
            outputs=api_output,
            api_name="api_chat",
        )

        stream_output = gr.Textbox(label="Streaming Response", lines=10)
        stream_btn = gr.Button("Call API (streaming)")
        stream_btn.click(
            fn=api_chat_stream,
            inputs=[msg_input, history_input],
            outputs=stream_output,
            api_name="api_chat_stream",
        )

if __name__ == "__main__":
    demo.launch()