import os
from dotenv import load_dotenv
from typing import Iterator, Tuple
from env_server import KernelOptimization_env, TASKS, app as openenv_app
from openai import OpenAI
from models import Action
import gradio as gr
import traceback
import uvicorn
from fastapi.responses import RedirectResponse

load_dotenv()

def task_baseline_code(task_id: str) -> str:
    return TASKS[task_id]["baseline_code"]

def extract_code(text: str) -> str:
    if "```" not in text:
        return text
    start = text.find("```")
    end = text.rfind("```")
    chunk = text[start + 3 : end]
    if chunk.startswith("cuda") or chunk.startswith("cpp"):
        return chunk.split("\n", 1)[1]
    return chunk

def ui(task_id:str, kernel_code: str, max_steps:int, openai_api_key:str)-> Iterator[Tuple[str,str]]:
    log= []
    env=KernelOptimization_env()
    api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
    if not api_key:
        yield "ERROR: Missing OPENAI_API_KEY", ""
        return

    model = os.getenv("MODEL_NAME", "llama-3.2-3b")
    client = OpenAI(api_key=api_key, base_url=os.getenv("API_BASE_URL", "https://api.openai.com/v1"))
    obs = env.reset(task_id=task_id)["observation"]
    if kernel_code and kernel_code.strip():
        custom_code = kernel_code.strip()
        env.state.best_code = custom_code
        obs["current_best_code"] = custom_code
    best_code = obs["current_best_code"]
    log.append(f"Task: {obs['task_name']}")

    for _ in range(max_steps):
        try:
            prompt = f"Optimize CUDA code:\n{obs['current_best_code']}\nPending checks: {obs['pending_checks']}\nReturn code only."
            res = client.chat.completions.create(
                model=model,
                temperature=0.0,
                messages=[
                    {"role": "system", "content": "Return only optimized CUDA code."},
                    {"role": "user", "content": prompt},
                ],
            )
            raw = (res.choices[0].message.content or "").strip()
            code = extract_code(raw).strip() or obs["current_best_code"]
            step = env.step(Action(optimized_code=code, strategy="ui_proposed"))
            obs = step.observation.model_dump()
            best_code = obs["current_best_code"]
            log.append(f"step={obs['step_count']} reward={step.reward.value:.3f} speedup={obs['current_best_speedup']:.3f}x")
            yield "\n".join(log), best_code
            if step.done:
                break
        except Exception as e:
            yield f"{chr(10).join(log)}\nERROR: {e}\n{traceback.format_exc()}", best_code
            return

with gr.Blocks(title="CUDA Kernel Optimizer") as demo:
    gr.Markdown("CUDA Kernel Optimizer - OpenEnv-aligned workflow")
    task = gr.Dropdown(choices=list(TASKS.keys()), value="vector_add_easy", label="Task")
    kernel_input = gr.Code(
        label="Kernel Code (editable, used as optimization input)",
        language="cpp",
        lines=16,
        value=TASKS["vector_add_easy"]["baseline_code"],
    )
    steps = gr.Slider(minimum=1, maximum=12, value=6, step=1, label="Max Steps")
    key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")
    run = gr.Button("Run Optimization", variant="primary")
    logs = gr.Textbox(label="Logs", lines=14)
    code = gr.Code(label="Best Code", language="cpp", lines=16)
    task.change(task_baseline_code, inputs=[task], outputs=[kernel_input])
    run.click(ui, inputs=[task, kernel_input, steps, key], outputs=[logs, code])

app = gr.mount_gradio_app(openenv_app, demo, path="/ui")


@app.get("/")
def root():
    return RedirectResponse(url="/ui")


if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("PORT", "7860")))