import os
import json
import httpx
import gradio as gr

# ---- Configure via Hugging Face Spaces → Settings → Repository secrets ----
# REQUIRED:
#   LETTA_BASE_URL  e.g. https://cloud.letta.ai/api/v1   (or your tenant base that ends with /v1 or /api/v1)
#   LETTA_API_KEY   your Letta Cloud API key
# OPTIONAL:
#   LETTA_AGENT_ID  default agent id (defaults to "default")
#   LETTA_MODEL     e.g. "openai/gpt-4-nano" or "openai/gpt-4o-mini" (optional; agent may decide if omitted)
#   LETTA_SESSION_ID  a stable id for memory (optional; you can also expose a textbox)

LETTA_BASE_URL  = (os.getenv("LETTA_BASE_URL"))
LETTA_API_KEY   = (os.getenv("LETTA_API_KEY"))
LETTA_AGENT_ID  = (os.getenv("LETTA_AGENT_ID", "default"))
# LETTA_MODEL     = (os.getenv("LETTA_MODEL")
DEFAULT_SESSION = (os.getenv("LETTA_SESSION_ID", ""))


def _config_error() -> str | None:
    if not LETTA_BASE_URL:
        return "Missing LETTA_BASE_URL (e.g., https://cloud.letta.ai/api/v1)"
    if not LETTA_BASE_URL.startswith("https://"):
        return "LETTA_BASE_URL must start with https://"
    if not LETTA_API_KEY:
        return "Missing LETTA_API_KEY"
    return None


def _stream_letta(messages, max_tokens, temperature, top_p, session_id: str):
    """
    Streams tokens from Letta Cloud (OpenAI-compatible SSE).
    """
    headers = {
        "Authorization": f"Bearer {LETTA_API_KEY}",
        "Content-Type": "application/json",
    }

    # Minimal OpenAI-compatible payload with Letta agent hints
    payload = {
        "messages": messages,                 # [{"role": "...", "content": "..."}]
        "stream": True,
        "max_tokens": int(max_tokens),
        "temperature": float(temperature),
        "top_p": float(top_p),
        # Letta-specific context (pass via 'extra' to avoid strict OpenAI schema issues on some stacks)
        "extra": {
            "agent_id": LETTA_AGENT_ID,
            "session_id": session_id or DEFAULT_SESSION
        }
    }

    # If you want to pin a specific OpenAI-backed model, set LETTA_MODEL
    if LETTA_MODEL:
        payload["model"] = LETTA_MODEL

    endpoint = f"{LETTA_BASE_URL}/chat/completions"  # base must already include /v1 or /api/v1

    # Stream Server-Sent Events (SSE) lines
    with httpx.stream("POST", endpoint, headers=headers, json=payload, timeout=120) as resp:
        if resp.status_code >= 400:
            # Surface backend error in the UI
            yield f"Letta error {resp.status_code}: {resp.text}"
            return

        partial = ""
        for raw in resp.iter_lines():
            if not raw:
                continue
            line = raw.decode("utf-8", errors="ignore")
            if not line.startswith("data: "):
                continue
            data = line[6:].strip()
            if data == "[DONE]":
                break
            try:
                chunk = json.loads(data)
                # OpenAI delta shape
                delta = chunk["choices"][0]["delta"]
                token = delta.get("content", "")
                if token:
                    partial += token
                    yield partial
            except Exception:
                # ignore keepalives or malformed lines
                continue


def respond(
    message,
    history: list[dict[str, str]],   # ChatInterface type="messages": [{"role","content"}, ...]
    system_message,
    max_tokens,
    temperature,
    top_p,
    _hf_token: gr.OAuthToken,        # unused; kept to match the default template signature
    session_id="",
):
    # Validate configuration once per call
    err = _config_error()
    if err:
        yield f"⚠️ Configuration error: {err}"
        return

    # Build messages for OpenAI-compatible API
    msgs = []
    if system_message:
        msgs.append({"role": "system", "content": system_message})
    if history:
        msgs.extend(history)  # already in [{"role","content"}] format when type="messages"
    msgs.append({"role": "user", "content": message or ""})

    try:
        for chunk in _stream_letta(msgs, max_tokens, temperature, top_p, session_id=session_id or DEFAULT_SESSION):
            yield chunk
    except httpx.ConnectError as e:
        yield f"❌ Network error reaching Letta: {e}"
    except Exception as e:
        yield f"❌ Unexpected error: {e}"


# ----------------- UI -----------------
chat = gr.ChatInterface(
    respond,
    type="messages",  # history comes in [{"role","content"}] which we forward as-is
    additional_inputs=[
        gr.Textbox(value="You are a concise, helpful assistant.", label="System message"),
        gr.Slider(1, 4096, value=512, step=1, label="Max new tokens"),
        gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
        gr.Textbox(value=DEFAULT_SESSION, label="Session ID (memory)"),
    ],
)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    with gr.Sidebar():
        gr.Markdown("### Letta Cloud Chat")
        gr.Markdown("Uses OpenAI-compatible streaming via Letta Cloud.")
    chat.render()

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)