import gradio as gr
import requests
import os
from datetime import datetime

# Configuração GLM-4.6
HF_TOKEN = os.getenv("HF_TOKEN", "")
MODEL_ID = "zai-org/GLM-4.6-FP8:zai-org"
API_URL = "https://router.huggingface.co/v1"

class CodeAgentGLM:
    def __init__(self):
        self.headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    
    def generate_code(self, prompt, language="python", max_tokens=1024):
        """Gera código usando GLM-4.6"""
        full_prompt = f"""You are an expert code generation AI. Generate high-quality {language} code.

Request: {prompt}

Provide only the code without explanations:"""
        
        payload = {
            "inputs": full_prompt,
            "parameters": {
                "max_new_tokens": max_tokens,
                "temperature": 0.3,
                "top_p": 0.9,
                "do_sample": True
            }
        }
        
        try:
            response = requests.post(API_URL, headers=self.headers, json=payload, timeout=30)
            if response.status_code == 200:
                result = response.json()
                if isinstance(result, list) and len(result) > 0:
                    return result[0].get("generated_text", "").split("Provide only the code without explanations:")[-1].strip()
            return f"Erro: {response.status_code}"
        except Exception as e:
            return f"Erro na chamada: {str(e)}"
    
    def explain_code(self, code):
        """Explica código"""
        prompt = f"Explique brevemente este código:\n\n{code}"
        payload = {
            "inputs": prompt,
            "parameters": {"max_new_tokens": 500, "temperature": 0.5}
        }
        try:
            response = requests.post(API_URL, headers=self.headers, json=payload)
            if response.status_code == 200:
                result = response.json()
                return result[0].get("generated_text", "") if isinstance(result, list) else str(result)
            return f"Erro: {response.status_code}"
        except Exception as e:
            return f"Erro: {str(e)}"

agent = CodeAgentGLM()

# Interface Gradio
with gr.Blocks(title="GLM-4.6 Code Agent", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 💻 GLM-4.6 Code Agent
    **KILOcode Agent Integration** - Geração e explicação de código com GLM-4.6
    
    - 🚀 Segue padrões GPT
    - 📚 Suporta 200K tokens de contexto
    - 💪 Excelente em coding tasks
    """)
    
    with gr.Tabs():
        with gr.Tab("🔨 Code Generation"):
            code_prompt = gr.Textbox(
                label="📝 Descrição do código",
                placeholder="Ex: Crie uma função que calcula fibonacci com memoização",
                lines=3
            )
            
            language = gr.Dropdown(
                choices=["python", "javascript", "java", "cpp", "go", "rust", "typescript, nodeJs"],
                value="python",
                label="🔤 Linguagem"
            )
            
            max_tokens_slider = gr.Slider(
                minimum=256,
                maximum=2048,
                value=1024,
                step=256,
                label="📏 Max tokens"
            )
            
            generate_btn = gr.Button("Gerar Código", variant="primary")
            code_output = gr.Code(language="python", label="💾 Código Gerado")
            
            def generate(prompt, lang, tokens):
                return agent.generate_code(prompt, lang, tokens)
            
            generate_btn.click(generate, [code_prompt, language, max_tokens_slider], code_output)
        
        with gr.Tab("📖 Code Explanation"):
            code_to_explain = gr.Code(
                language="python",
                label="Código para explicar",
                lines=10
            )
            
            explain_btn = gr.Button("Explicar", variant="primary")
            explanation_output = gr.Textbox(label="Explicação", lines=6)
            
            explain_btn.click(agent.explain_code, code_to_explain, explanation_output)
        
        with gr.Tab("⚙️ Configuração"):
            gr.Markdown("""
            ## Configuração do KILOcode Agent
            
            ### Modelo: GLM-4.6
            - **Tamanho**: 357B parâmetros
            - **Contexto**: 200K tokens
            - **Formato**: BF16
            - **Performance**: ⭐⭐⭐⭐⭐ para coding
            
            ### Como usar com KILOcode:
            1. Adicione seu HF_TOKEN nos secrets
            2. Use os endpoints de geração e explicação
            3. Suporta múltiplas linguagens de programação
            """)

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