import gradio as gr
import requests
import json
import os

FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")
API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
MODEL_NAME = "accounts/fireworks/models/glm-4p6"


# ---------- Fireworks GLM Request ----------
def query_fireworks(messages, max_tokens=2048, temperature=0.6):
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {FIREWORKS_API_KEY}",
    }
    payload = {
        "model": MODEL_NAME,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "messages": messages,
    }
    response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"Error: {response.text}"


# ---------- Documentation Generator ----------
def generate_docs(code_input_type, code_content, file_obj, github_url, api_key):
    global FIREWORKS_API_KEY
    if api_key:
        FIREWORKS_API_KEY = api_key

    code = ""
    if code_input_type == "Paste Code":
        code = code_content
    elif code_input_type == "Upload File" and file_obj is not None:
        code = file_obj.read().decode("utf-8")
    elif code_input_type == "GitHub Link" and github_url:
        if "raw.githubusercontent.com" not in github_url:
            github_url = github_url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/")
        r = requests.get(github_url)
        if r.status_code == 200:
            code = r.text
        else:
            return f"Unable to fetch file (HTTP {r.status_code})", ""

    if not code.strip():
        return "⚠️ Please provide valid Python code.", ""

    messages = [
        {
            "role": "system",
            "content": (
                "You are an AI documentation assistant that analyzes Python code and generates "
                "multi-layer documentation. Include: 1) Overview, 2) Key classes/functions, "
                "3) Usage examples, 4) Dependencies, 5) Suggested improvements."
            ),
        },
        {"role": "user", "content": code},
    ]

    output = query_fireworks(messages, max_tokens=4096)
    return output, code


# ---------- Code Chat ----------
def chat_with_code(user_input, chat_history, code_context, api_key):
    global FIREWORKS_API_KEY
    if api_key:
        FIREWORKS_API_KEY = api_key

    if not code_context.strip():
        return "⚠️ No code context found. Please generate documentation first.", chat_history

    messages = [
        {
            "role": "system",
            "content": (
                "You are a code analysis assistant. You answer questions about the following Python code. "
                "Explain clearly and refer to specific functions or logic when relevant."
            ),
        },
        {"role": "user", "content": f"Code:\n{code_context}\n\nQuestion:\n{user_input}"},
    ]

    reply = query_fireworks(messages)
    chat_history.append(("You", user_input))
    chat_history.append(("GLM-4.6", reply))
    return "", chat_history


# ---------- Gradio UI ----------

with gr.Blocks(title="Intelligent Documentation Generator Agent") as demo:
    gr.Markdown("# 🧠 Intelligent Documentation Generator Agent (GLM-4.6 on Fireworks AI)")
    gr.Markdown(
        "Generate intelligent documentation or chat with your Python code. "
        "Provide your Fireworks API key to get started."
    )

    with gr.Row():
        api_key_box = gr.Textbox(
            label="🔑 Fireworks API Key",
            placeholder="Enter your Fireworks API Key",
            type="password",
        )

    with gr.Tabs():
        # ---- Documentation Tab ----
        with gr.TabItem("📄 Generate Documentation"):
            code_input_type = gr.Radio(
                ["Paste Code", "Upload File", "GitHub Link"], label="Choose input type", value="Paste Code"
            )

            code_box = gr.Textbox(label="Paste your Python code", lines=15, visible=True)
            file_upload = gr.File(label="Upload Python file", visible=False)
            github_box = gr.Textbox(label="GitHub raw file link", visible=False)

            def toggle_inputs(input_type):
                return (
                    gr.update(visible=input_type == "Paste Code"),
                    gr.update(visible=input_type == "Upload File"),
                    gr.update(visible=input_type == "GitHub Link"),
                )

            code_input_type.change(toggle_inputs, [code_input_type], [code_box, file_upload, github_box])

            generate_btn = gr.Button("🚀 Generate Documentation")

            output_box = gr.Markdown(label="📘 Generated Documentation")
            hidden_code_context = gr.Textbox(visible=False)

            generate_btn.click(
                generate_docs,
                [code_input_type, code_box, file_upload, github_box, api_key_box],
                [output_box, hidden_code_context],
            )

        # ---- Chat Tab ----
        with gr.TabItem("💬 Chat with Code"):
            gr.Markdown("Chat about the Python file you uploaded or analyzed in the previous tab.")
            chatbot = gr.Chatbot(label="GLM-4.6 Code Assistant")
            msg = gr.Textbox(label="Ask a question about the code")
            send_btn = gr.Button("Ask")

            send_btn.click(
                chat_with_code,
                [msg, chatbot, hidden_code_context, api_key_box],
                [msg, chatbot],
            )

demo.launch()