#!/usr/bin/env python3
"""
Teste da API do LM Studio para completar código Python
"""

import json
import urllib.request
import urllib.error

# Configuração
API_URL = "http://192.168.1.30:1234/v1/chat/completions"
MODEL = "openai/gpt-oss-120b"

# Código incompleto para a IA completar
CODIGO_INCOMPLETO = """
def calcular_media(numeros):
    # TODO: implementar
    pass

def encontrar_maior(lista):
    # TODO: implementar
    pass
"""

# Prompt para a IA
prompt = f"""Complete o código Python abaixo implementando as funções.
Adicione docstrings e type hints. Seja conciso.

{CODIGO_INCOMPLETO}

Retorne apenas o código completo, sem explicações."""

# Preparar requisição
data = {
    "model": MODEL,
    "messages": [
        {"role": "system", "content": "Você é um assistente de programação especializado em Python."},
        {"role": "user", "content": prompt}
    ],
    "temperature": 0.3,
    "max_tokens": 800
}

print("🤖 Testando LM Studio API...")
print(f"📡 Servidor: {API_URL}")
print(f"🧠 Modelo: {MODEL}")
print("\n" + "="*60)
print("📝 CÓDIGO ORIGINAL (incompleto):")
print("="*60)
print(CODIGO_INCOMPLETO)
print("="*60)

try:
    # Fazer requisição
    req = urllib.request.Request(
        API_URL,
        data=json.dumps(data).encode('utf-8'),
        headers={'Content-Type': 'application/json'}
    )

    print("\n⏳ Aguardando resposta da IA (pode levar 10-30s para modelo 120B)...\n")

    with urllib.request.urlopen(req, timeout=60) as response:
        result = json.loads(response.read().decode('utf-8'))

    # Extrair resposta
    codigo_completo = result['choices'][0]['message']['content']

    print("="*60)
    print("✨ CÓDIGO COMPLETADO PELA IA:")
    print("="*60)
    print(codigo_completo)
    print("="*60)

    # Salvar resultado
    with open('/Users/neog/teste-ia-resultado.py', 'w') as f:
        f.write(codigo_completo)

    print("\n✅ Teste concluído com sucesso!")
    print("📄 Código salvo em: /Users/neog/teste-ia-resultado.py")

    # Estatísticas
    usage = result.get('usage', {})
    print(f"\n📊 Estatísticas:")
    print(f"   • Tokens de entrada: {usage.get('prompt_tokens', 'N/A')}")
    print(f"   • Tokens de saída: {usage.get('completion_tokens', 'N/A')}")
    print(f"   • Total: {usage.get('total_tokens', 'N/A')}")

except urllib.error.URLError as e:
    print(f"\n❌ Erro de conexão: {e}")
    print("\n💡 Verifique:")
    print("   1. LM Studio está rodando em 192.168.1.30:1234")
    print("   2. Modelo 'openai/gpt-oss-120b' está carregado")
    print("   3. CORS está configurado no LM Studio")
except Exception as e:
    print(f"\n❌ Erro: {e}")