import os
import requests
import json
from dotenv import load_dotenv
from typing import List, Dict, Any

# Carregar variáveis de ambiente
load_dotenv()

class GrokAgent:
    def __init__(self, model: str = "grok-4", fast_mode: bool = True):
        """
        Inicializa o agente Grok AI.

        :param model: 'grok-4' ou 'grok-3'
        :param fast_mode: True para respostas rápidas (fast), False para modo completo
        """
        self.api_key = os.getenv("XAI_API_KEY")
        if not self.api_key:
            raise ValueError("XAI_API_KEY não encontrada no .env")

        self.model = model
        self.fast_mode = fast_mode
        self.base_url = "https://api.x.ai/v1/chat/completions"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def _build_messages(self, user_prompt: str, history: List[Dict] = None) -> List[Dict]:
        """Constrói a lista de mensagens com histórico opcional."""
        system_prompt = (
            "Você é Grok, uma IA útil, direta e com senso de humor. "
            "Responda de forma clara, concisa e inteligente. "
            "Use raciocínio lógico e evite enrolação."
        )

        messages = [{"role": "system", "content": system_prompt}]

        if history:
            messages.extend(history)

        messages.append({"role": "user", "content": user_prompt})
        return messages

    def ask(self, prompt: str, history: List[Dict] = None) -> str:
        """
        Envia uma pergunta ao Grok e retorna a resposta.

        :param prompt: Pergunta do usuário
        :param history: Histórico de conversa (opcional)
        :return: Resposta do Grok
        """
        payload = {
            "model": self.model,
            "messages": self._build_messages(prompt, history),
            "temperature": 0.5 if self.fast_mode else 0.7,
            "max_tokens": 512 if self.fast_mode else 2048,
            "stream": False
        }

        try:
            response = requests.post(self.base_url, headers=self.headers, json=payload)
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"].strip()
        except requests.exceptions.RequestException as e:
            return f"Erro na API: {e}"
        except Exception as e:
            return f"Erro inesperado: {e}"

    def chat(self):
        """Modo interativo de chat com Grok (fast mode ativado)."""
        print("🤖 Grok AI Fast 4 - Digite 'sair' para encerrar.\n")
        history = []

        while True:
            user_input = input("Você: ").strip()
            if user_input.lower() in ["sair", "exit", "quit"]:
                print("Até logo! 🚀")
                break
            if not user_input:
                continue

            response = self.ask(user_input, history)
            print(f"Grok: {response}\n")

            # Atualiza histórico
            history.append({"role": "user", "content": user_input})
            history.append({"role": "assistant", "content": response})

            # Limita histórico para evitar excesso
            if len(history) > 20:
                history = history[-20:]

# === EXECUÇÃO ===
if __name__ == "__main__":
    # Cria agente Grok-4 em modo rápido
    grok = GrokAgent(model="grok-4", fast_mode=True)

    # Modo chat interativo
    grok.chat()