from __future__ import annotations

import json
import os
import tempfile
import time
from pathlib import Path

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from remora_control import LayerController

MODEL_ID = "LiquidAI/LFM2.5-1.2B-Instruct"
MAX_OUTPUT_TOKENS = 128
DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")

# ZeroGPU's CUDA emulation makes module-level placement correct: a real GPU is
# attached only when `run_experiment` enters the decorator.
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
).to(DEVICE).eval()
LAYER_COUNT = len(model.model.layers)


def parse_layers(text: str) -> list[int]:
    raw = [part.strip() for part in text.split(",") if part.strip()]
    try:
        layers = sorted(set(int(part) for part in raw))
    except ValueError:
        raise gr.Error("Layers must be whole numbers, e.g. '2,5'.")
    invalid = [value for value in layers if value < 0 or value >= LAYER_COUNT]
    if invalid:
        raise gr.Error(f"Layer values must be between 0 and {LAYER_COUNT - 1}; got {invalid}.")
    return layers


@spaces.GPU(duration=120, size="large")
def run_experiment(prompt: str, selected_layers: str, gain: float, max_new_tokens: int, seed: int):
    if not prompt.strip():
        raise gr.Error("Enter a prompt.")
    if max_new_tokens < 1 or max_new_tokens > MAX_OUTPUT_TOKENS:
        raise gr.Error(f"Output tokens must be 1–{MAX_OUTPUT_TOKENS}.")

    layers = parse_layers(selected_layers)
    controller = LayerController(layers, gain)
    seed_value = int(seed)
    generator = torch.Generator(device=DEVICE).manual_seed(seed_value)
    started = time.perf_counter()
    controller.attach(model)
    try:
        messages = [{"role": "user", "content": prompt}]
        rendered = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(rendered, return_tensors="pt").to(DEVICE)
        with torch.inference_mode():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=int(max_new_tokens),
                do_sample=True,
                temperature=0.1,
                top_k=50,
                repetition_penalty=1.05,
                generator=generator,
                use_cache=True,
            )
    finally:
        controller.detach()

    elapsed_s = time.perf_counter() - started
    new_tokens = output_ids[0, inputs.input_ids.shape[1] :]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    trace = {
        "schema": "remora-layer-lab/v1",
        "model": MODEL_ID,
        "layer_count": LAYER_COUNT,
        "selected_layers": layers,
        "gain": gain,
        "seed": seed_value,
        "prompt": prompt,
        "generated_tokens": int(new_tokens.shape[0]),
        "elapsed_s": round(elapsed_s, 4),
        "tokens_per_second": round(float(new_tokens.shape[0]) / elapsed_s, 3) if elapsed_s else None,
        "events": controller.json_events(),
    }
    artifact_fd, artifact_path = tempfile.mkstemp(prefix="remora-trace-", suffix=".json")
    os.close(artifact_fd)  # close the descriptor mkstemp returns; only the path is needed
    Path(artifact_path).write_text(json.dumps(trace, indent=2), encoding="utf-8")
    status = (
        f"{len(layers)} controlled layers; {len(trace['events'])} trace events; "
        f"{trace['generated_tokens']} tokens in {elapsed_s:.2f}s "
        f"({trace['tokens_per_second']} tok/s)."
    )
    return text, status, str(artifact_path)


with gr.Blocks(title="Remora Layer Lab") as demo:
    gr.Markdown(
        "# Remora Layer Lab\n"
        "PyTorch/ZeroGPU research adapter for controlled LFM2.5 decoder-layer experiments. "
        "`gain=0` is an observation-only baseline. This is not the local Rust production lane."
    )
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", lines=6, value="Explain why water forecasting needs uncertainty.")
        with gr.Column():
            selected_layers = gr.Textbox(label=f"Layers (0–{LAYER_COUNT - 1}, comma-separated)", value="2,5")
            gain = gr.Slider(-1.0, 1.0, value=0.0, step=0.05, label="Layer gain")
            max_new_tokens = gr.Slider(1, MAX_OUTPUT_TOKENS, value=32, step=1, label="Max output tokens")
            seed = gr.Number(value=42, precision=0, label="Random seed")
            run = gr.Button("Run controlled experiment", variant="primary")
    answer = gr.Textbox(label="Generated text", lines=8)
    status = gr.Textbox(label="Run statistics")
    trace = gr.File(label="Download trace JSON")
    run.click(run_experiment, [prompt, selected_layers, gain, max_new_tokens, seed], [answer, status, trace])

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