import gradio as gr
import numpy as np, imageio, tempfile, os
import pybullet as p
import openai
from openai import AuthenticationError
from coral.env import BoxPushEnv
from coral.mppi import MPPI
from coral.llm_feedback import LLMFeedback


def run_demo(api_key):
    """Main demo entrypoint for Hugging Face Space"""

    try:
        p.disconnect()  # closes any active simulation session
    except Exception:
        pass

    api_key = api_key.strip()
    if not api_key:
        return None, "⚠️ Please enter a valid OpenAI API key."

    try:
        openai.api_key = api_key
        # Quick ping test for key validation
        openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "ping"}],
            max_tokens=1
        )
    except AuthenticationError:
        return None, "⚠️ Invalid OpenAI API key. Please check and try again."
    except Exception as e:
        return None, f"⚠️ Could not reach OpenAI API: {e}"

    cid = p.connect(p.DIRECT)
    env = BoxPushEnv(cid)
    mppi = MPPI(env)
    llm = LLMFeedback(api_key)

    cam_w, cam_h = 320, 240
    view = p.computeViewMatrix([0.5, -2, 1], [0.5, 0, 0.5], [0, 0, 1])
    proj = p.computeProjectionMatrixFOV(60, cam_w / cam_h, 0.1, 10)
    frames = []

    env.reset()
    history = []

    for step in range(35):
        u = mppi.compute_control()
        box_pos, ee_pos = env.step(u)
        cost = env.state_cost()
        history.append({
            "step": step,
            "box_pos": box_pos.tolist(),
            "ee_pos": ee_pos.tolist(),
            "cost": cost
        })

        _, _, rgb, _, _ = p.getCameraImage(cam_w, cam_h, view, proj, renderer=p.ER_TINY_RENDERER)
        frame = np.array(rgb, dtype=np.uint8).reshape(cam_h, cam_w, 4)[:, :, :3]
        frames.append(frame)

    try:
        code, expl = llm.ask_state_cost_fn("Flip the box using the wall", history, env)
        env.update_state_cost(code)
    except AuthenticationError:
        return None, "⚠️ Your API key was rejected during LLM call. Please try again."
    except Exception as e:
        return None, f"⚠️ Error while querying LLM: {e}"

    tmp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    imageio.mimsave(tmp_path, frames, fps=10)

    p.disconnect()

    return tmp_path, f"✅ LLM updated cost:\n{code}\n\nExplanation:\n{expl}"


gr.Interface(
    fn=run_demo,
    inputs=gr.Textbox(
        label="Enter your OpenAI API Key (without quotation marks)",
        type="password",
        placeholder="sk-..."
    ),
    outputs=["video", "text"],
    title="CoRAL: LLM-guided Box Flipping",
    description=(
        "This demo uses a PyBullet simulation of a Panda robot pushing a box against a wall. "
        "After several steps, an LLM rewrites the cost function guiding the control policy."
    ),
).launch()