import os
import json
from huggingface_hub import create_repo, list_models, upload_file, constants
from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
import gradio as gr
from google import genai
from google.genai import types

# --- Globals ---
client = None  # Will hold Gemini client
chat = None    # Will hold Gemini chat session

# --- System prompt for Gemini ---
system_instruction = (
    "You are a helpful assistant that writes, debugs, and pushes code to Hugging Face Spaces. "
    "Treat Hugging Face as a sandbox: create spaces, upload code, and debug via function calls. "
    "Respond in JSON with {success, data, message}."
)

# --- Function declarations for logs (behind the scenes) ---
get_build_logs_decl = {
    "name": "get_build_logs",
    "description": "Fetches build logs for a Space",
    "parameters": {"type":"object","properties":{"repo_id":{"type":"string"}},"required":["repo_id"]}
}
get_container_logs_decl = {
    "name": "get_container_logs",
    "description": "Fetches container logs for a Space",
    "parameters": {"type":"object","properties":{"repo_id":{"type":"string"}},"required":["repo_id"]}
}
tools = [types.Tool(function_declarations=[get_build_logs_decl, get_container_logs_decl])]

# --- Core Hugging Face functions (unchanged) ---
def _fetch_space_logs_level(repo_id: str, level: str):
    jwt_url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
    r = get_session().get(jwt_url, headers=build_hf_headers())
    hf_raise_for_status(r)
    jwt = r.json()["token"]
    logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
    records = []
    with get_session().get(logs_url, headers=build_hf_headers(token=jwt), stream=True) as resp:
        hf_raise_for_status(resp)
        for raw in resp.iter_lines():
            if not raw.startswith(b"data: "): continue
            try:
                event = json.loads(raw[len(b"data: "):].decode())
                records.append({"timestamp": event.get("timestamp"), "message": event.get("data")})
            except: pass
    return records

# --- Backends for creating space & uploading files ---
def create_space_backend(username: str, hf_token: str, repo_name: str, sdk: str) -> str:
    repo_id = f"{username}/{repo_name}"
    create_repo(
        repo_id=repo_id,
        token=hf_token,
        exist_ok=True,
        repo_type="space",
        space_sdk=sdk
    )
    return repo_id

# --- Chat initialization & handlers ---
def init_chat(repo_name: str, sdk: str, api_key: str, hf_profile: gr.OAuthProfile, hf_token: gr.OAuthToken):
    global client, chat
    # Validate inputs
    if not api_key:
        return {"success": False, "data": None, "message": "Missing Gemini API key."}, ""
    if hf_profile is None or hf_token is None:
        return {"success": False, "data": None, "message": "Please sign in with Hugging Face."}, ""
    # Create HF space sandbox
    repo_id = create_space_backend(hf_profile.username, hf_token.token, repo_name, sdk)
    # Set HF token for log streaming
    os.environ["HF_TOKEN"] = hf_token.token
    # Init Gemini client and chat
    client = genai.Client(api_key=api_key)
    chat = client.chats.create(
        model="gemini-2.5-flash-preview-04-17",
        config=types.GenerateContentConfig(
            system_instruction=system_instruction,
            tools=tools,
            temperature=0
        )
    )
    return {"success": True, "data": None, "message": f"Sandbox ready: {repo_id}"}, repo_id


def chatbot_respond(message: str, history: list, repo_id: str, api_key: str):
    global chat, client
    if not chat:
        # Should not happen if initialized properly
        history.append((None, "Error: chat not initialized."))
        return history
    response = chat.send_message(message)
    part = response.candidates[0].content.parts[0]
    # Handle function calls
    if part.function_call:
        fn = part.function_call
        args = json.loads(fn.args)
        if fn.name == "get_build_logs":
            result = _fetch_space_logs_level(repo_id, "build")
        else:
            result = _fetch_space_logs_level(repo_id, "run")
        response2 = chat.send_message("", function_response={fn.name: result})
        reply = response2.candidates[0].content.parts[0].text
    else:
        reply = part.text
    history.append((message, reply))
    return history

# --- UI: Chatbot with Sidebar for Space Creation ---
with gr.Blocks(title="HF Code Sandbox Chat") as demo:
    # Top bar: HF login
    with gr.Row():
        login_btn = gr.LoginButton(variant="huggingface", label="Sign in with HF")
        login_status = gr.Markdown("*Not signed in.*")
        login_btn.click(
            lambda profile: f"✅ Logged in as {profile.username}" if profile else "*Not signed in.*",
            inputs=None, outputs=login_status
        )

    with gr.Row():
        # Sidebar for sandbox setup
        with gr.Column(scale=2):
            gr.Markdown("### 🏗️ Create New Space Sandbox")
            api_key = gr.Textbox(label="Gemini API Key", placeholder="sk-...", type="password")
            repo_name = gr.Textbox(label="Space Name", placeholder="my-sandbox")
            sdk_selector = gr.Radio(label="SDK", choices=["gradio","streamlit"], value="gradio")
            create_btn = gr.Button("Initialize Sandbox")
            create_status = gr.JSON(label="Initialization Status")
            # Hidden store for repo_id
            repo_store = gr.Variable("")
            create_btn.click(
                init_chat,
                inputs=[repo_name, sdk_selector, api_key, login_btn.profile, login_btn.token],
                outputs=[create_status, repo_store]
            )

        # Chat area
        with gr.Column(scale=8):
            chatbot = gr.Chatbot()
            user_input = gr.Textbox(show_label=False, placeholder="Ask the sandbox to write or debug code...")
            user_input.submit(
                chatbot_respond,
                inputs=[user_input, chatbot, repo_store, api_key],
                outputs=chatbot
            )

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