"""
Visão de ecrã + Grok Vision (xAI) — uso pelo agente HTTP local (Orb Logic).

Variáveis de ambiente: XAI_API_KEY ou GROK_API_KEY; opcional XAI_VISION_MODEL (default grok-2-vision-latest).

Dependências: pip install openai pillow pyautogui
"""

import base64
import json
import os
import time
from pathlib import Path
from typing import Optional, Callable
from openai import OpenAI
from PIL import ImageGrab
import pyautogui

_VISION_CACHE = Path(__file__).resolve().parent / ".ad06_cache" / "vision"


def default_vision_screenshot_dir() -> Path:
    _VISION_CACHE.mkdir(parents=True, exist_ok=True)
    return _VISION_CACHE


def resolve_xai_api_key(explicit: str | None = None) -> str | None:
    if explicit and str(explicit).strip():
        return str(explicit).strip()
    return os.environ.get("XAI_API_KEY") or os.environ.get("GROK_API_KEY")


# ==============================================================================
# CLASSE PRINCIPAL
# ==============================================================================

class VisionLoop:
    """
    Sistema de visão do agente — conecta a tela ao modelo Grok Vision.

    Fluxo:
        1. Captura screenshot da tela (ou região)
        2. Envia ao Grok Vision com contexto e pergunta
        3. Recebe análise detalhada
        4. Opcionalmente executa a ação sugerida

    Uso simples:
        vision = VisionLoop(api_key="sua_chave")
        resposta = vision.ver_e_responder("O que está na tela?")
        print(resposta)
    """

    def __init__(
        self,
        api_key: str | None = None,
        model: str | None = None,
        screenshot_dir: str | Path | None = None,
        max_tokens: int = 1500,
        debug: bool = False,
    ):
        key = resolve_xai_api_key(api_key)
        if not key:
            raise ValueError(
                "Chave xAI em falta: defina a variável de ambiente XAI_API_KEY ou GROK_API_KEY."
            )
        self.client = OpenAI(api_key=key, base_url="https://api.x.ai/v1")
        self.model = model or os.environ.get("XAI_VISION_MODEL", "grok-2-vision-latest")
        self.max_tokens = max_tokens
        self.debug = debug

        self.screenshot_dir = Path(screenshot_dir) if screenshot_dir else default_vision_screenshot_dir()
        self.screenshot_dir.mkdir(parents=True, exist_ok=True)

        # Histórico de visão — o agente "lembra" o que viu
        self._vision_history: list[dict] = []
        self._last_analysis: Optional[dict] = None

    # ==========================================================================
    # CAPTURA DE TELA
    # ==========================================================================

    def capture_screen(self, region: Optional[tuple] = None) -> str:
        """
        Captura a tela e retorna como base64.

        Args:
            region : Tupla (x, y, largura, altura) ou None para tela inteira

        Returns:
            String base64 da imagem
        """
        if region:
            x, y, w, h = region
            img = ImageGrab.grab(bbox=(x, y, x + w, y + h))
        else:
            img = ImageGrab.grab()

        if self.debug:
            ts = int(time.time())
            img.save(self.screenshot_dir / f"vision_{ts}.png")

        # Converte para base64
        import io
        buffer = io.BytesIO()
        img.save(buffer, format="PNG")
        return base64.b64encode(buffer.getvalue()).decode()

    def save_screenshot(self, name: str = "screen") -> str:
        """
        Tira screenshot e salva em arquivo.

        Returns:
            Caminho do arquivo salvo
        """
        ts = int(time.time())
        path = str(self.screenshot_dir / f"{name}_{ts}.png")
        img = ImageGrab.grab()
        img.save(path)
        return path

    # ==========================================================================
    # VISÃO — O AGENTE "VÊ" A TELA
    # ==========================================================================

    def ver_tela(
        self,
        pergunta: str,
        region: Optional[tuple] = None,
        contexto_extra: str = "",
    ) -> str:
        """
        Função principal: tira screenshot e pergunta ao Grok o que está vendo.

        Args:
            pergunta      : O que o agente deve responder sobre a tela
            region        : Região específica da tela (ou None para tela inteira)
            contexto_extra: Contexto adicional para o modelo

        Returns:
            Resposta do Grok descrevendo o que vê

        Exemplo:
            resposta = vision.ver_tela("O que você está vendo na tela?")
            resposta = vision.ver_tela("Tem algum botão de 'Salvar' visível?")
            resposta = vision.ver_tela("Qual é o título da janela aberta?")
        """
        img_b64 = self.capture_screen(region)

        system_prompt = """Você é o sistema de visão de um agente AI desktop.
Sua função é analisar screenshots da tela do computador e descrever com precisão o que está visível.

Ao analisar uma tela, sempre informe:
- Qual aplicativo/janela está em foco
- Textos visíveis importantes
- Botões, menus e elementos clicáveis disponíveis
- Estado geral da interface (carregando, erro, sucesso, etc.)
- Qualquer informação relevante para o usuário

Seja preciso e objetivo. Se o usuário pedir para clicar em algo, forneça as coordenadas (x, y) aproximadas."""

        user_content = [
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/png;base64,{img_b64}"}
            },
            {
                "type": "text",
                "text": f"{contexto_extra}\n\nPergunta: {pergunta}".strip()
            }
        ]

        response = self.client.chat.completions.create(
            model=self.model,
            max_tokens=self.max_tokens,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content}
            ]
        )

        resultado = response.choices[0].message.content

        # Salva no histórico
        self._vision_history.append({
            "timestamp": time.time(),
            "pergunta": pergunta,
            "resposta": resultado,
            "region": region
        })

        return resultado

    def descrever_tela_completa(self) -> str:
        """
        Faz uma descrição completa de tudo que está na tela.
        Útil para o agente entender o estado atual do sistema.

        Returns:
            Descrição detalhada da tela

        Exemplo:
            descricao = vision.descrever_tela_completa()
            print(descricao)
            # "A tela mostra o navegador Chrome aberto na página do Google.
            #  Há uma barra de pesquisa no centro, um botão 'Pesquisa Google'
            #  e 'Estou com sorte'. No canto superior direito há um botão 'Entrar'..."
        """
        return self.ver_tela(
            "Faça uma descrição detalhada e completa de tudo que está visível na tela. "
            "Liste todos os elementos: janelas abertas, textos, botões, menus, ícones, "
            "e qualquer informação relevante. Organize por área da tela."
        )

    def identificar_elementos_clicaveis(self) -> dict:
        """
        Identifica todos os elementos clicáveis na tela com suas coordenadas.

        Returns:
            Dict com elementos clicáveis e suas posições

        Exemplo:
            elementos = vision.identificar_elementos_clicaveis()
            # {"botoes": [{"nome": "OK", "x": 450, "y": 300}], "links": [...]}
        """
        resposta = self.ver_tela(
            """Identifique todos os elementos clicáveis visíveis na tela.
            Para cada elemento, forneça:
            - Nome/texto do elemento
            - Tipo (botão, link, ícone, campo de texto, menu, etc.)
            - Coordenadas aproximadas (x, y) no centro do elemento
            
            Responda APENAS em JSON no formato:
            {
                "botoes": [{"nome": "...", "x": 0, "y": 0}],
                "links": [{"nome": "...", "x": 0, "y": 0}],
                "campos_texto": [{"nome": "...", "x": 0, "y": 0}],
                "menus": [{"nome": "...", "x": 0, "y": 0}],
                "outros": [{"nome": "...", "x": 0, "y": 0}]
            }"""
        )

        try:
            # Limpa o JSON da resposta
            json_str = resposta
            if "```json" in json_str:
                json_str = json_str.split("```json")[1].split("```")[0]
            elif "```" in json_str:
                json_str = json_str.split("```")[1].split("```")[0]
            return json.loads(json_str.strip())
        except Exception:
            return {"raw_response": resposta}

    def encontrar_elemento(self, descricao: str) -> Optional[tuple]:
        """
        Encontra um elemento específico na tela e retorna suas coordenadas.

        Args:
            descricao : Descrição do elemento a encontrar

        Returns:
            Tupla (x, y) do elemento ou None se não encontrado

        Exemplo:
            pos = vision.encontrar_elemento("botão de fechar janela")
            pos = vision.encontrar_elemento("campo de pesquisa do Google")
            pos = vision.encontrar_elemento("ícone da lixeira na área de trabalho")
        """
        resposta = self.ver_tela(
            f"""Encontre o seguinte elemento na tela: "{descricao}"
            
            Se encontrar, responda APENAS com as coordenadas no formato JSON:
            {{"encontrado": true, "x": 000, "y": 000, "descricao": "..."}}
            
            Se NÃO encontrar:
            {{"encontrado": false, "motivo": "..."}}"""
        )

        try:
            json_str = resposta
            if "```" in json_str:
                json_str = json_str.split("```")[1].split("```")[0]
                if json_str.startswith("json"):
                    json_str = json_str[4:]
            data = json.loads(json_str.strip())
            if data.get("encontrado"):
                return (data["x"], data["y"])
        except Exception:
            pass
        return None

    def ler_texto_na_tela(self, region: Optional[tuple] = None) -> str:
        """
        Lê e retorna todo o texto visível na tela via visão do Grok.
        Mais preciso que OCR tradicional para textos complexos.

        Args:
            region : Região específica (ou None para tela inteira)

        Returns:
            Todo o texto visível na tela
        """
        return self.ver_tela(
            "Liste TODO o texto visível na tela, exatamente como aparece, "
            "preservando a estrutura e formatação. Não adicione interpretações, "
            "apenas transcreva o texto que está visível.",
            region=region
        )

    def verificar_estado(self, descricao_esperada: str) -> dict:
        """
        Verifica se a tela está no estado esperado.

        Args:
            descricao_esperada : O que deveria estar na tela

        Returns:
            Dict com {"correto": bool, "diferenca": str, "sugestao": str}

        Exemplo:
            estado = vision.verificar_estado("página de login do Gmail")
            if not estado["correto"]:
                print(f"Problema: {estado['diferenca']}")
        """
        resposta = self.ver_tela(
            f"""Verifique se a tela atual corresponde ao estado esperado.
            
            Estado esperado: "{descricao_esperada}"
            
            Responda em JSON:
            {{
                "correto": true/false,
                "o_que_esta_na_tela": "...",
                "diferenca": "...",
                "sugestao": "..."
            }}"""
        )

        try:
            json_str = resposta
            if "```" in json_str:
                json_str = json_str.split("```")[1].split("```")[0]
                if json_str.startswith("json"):
                    json_str = json_str[4:]
            return json.loads(json_str.strip())
        except Exception:
            return {"correto": False, "raw": resposta}

    # ==========================================================================
    # PERCEPTION-ACTION LOOP — O CICLO PRINCIPAL DO AGENTE
    # ==========================================================================

    def executar_tarefa(
        self,
        objetivo: str,
        max_steps: int = 10,
        callback_status: Optional[Callable] = None,
    ) -> dict:
        """
        Loop principal de Percepção → Raciocínio → Ação.
        O agente vê a tela, decide o que fazer e executa — repetindo até concluir.

        Args:
            objetivo       : O que o agente deve fazer (em linguagem natural)
            max_steps      : Número máximo de passos para evitar loop infinito
            callback_status: Função chamada a cada passo com status (opcional)
                             Assinatura: callback_status(passo, descricao, acao)

        Returns:
            Dict com resultado: {"concluido": bool, "passos": [...], "resultado": str}

        Exemplo:
            resultado = vision.executar_tarefa("Abra o Bloco de Notas e escreva 'Olá Mundo'")
            resultado = vision.executar_tarefa("Feche todas as janelas abertas")
        """
        passos = []
        historico_acoes = []

        system_prompt = f"""Você é um agente AI que controla o computador.
        
OBJETIVO: {objetivo}

A cada passo você vai:
1. Ver um screenshot da tela atual
2. Analisar o que está vendo
3. Decidir a próxima ação para avançar no objetivo
4. Retornar a ação em JSON

AÇÕES DISPONÍVEIS:
- click: {{"acao": "click", "x": 000, "y": 000, "descricao": "..."}}
- double_click: {{"acao": "double_click", "x": 000, "y": 000, "descricao": "..."}}
- right_click: {{"acao": "right_click", "x": 000, "y": 000, "descricao": "..."}}
- type: {{"acao": "type", "texto": "...", "descricao": "..."}}
- press_key: {{"acao": "press_key", "tecla": "Enter", "descricao": "..."}}
- scroll: {{"acao": "scroll", "x": 000, "y": 000, "direcao": "down", "quantidade": 3, "descricao": "..."}}
- wait: {{"acao": "wait", "segundos": 2, "descricao": "..."}}
- concluido: {{"acao": "concluido", "resultado": "Tarefa concluída com sucesso porque..."}}
- falhou: {{"acao": "falhou", "motivo": "..."}}

Responda APENAS com o JSON da próxima ação. Sem texto extra."""

        for passo in range(1, max_steps + 1):
            img_b64 = self.capture_screen()

            # Monta histórico de ações para o agente saber o que já fez
            historico_str = ""
            if historico_acoes:
                historico_str = f"\n\nAções já executadas:\n" + "\n".join(
                    f"  Passo {i+1}: {a}" for i, a in enumerate(historico_acoes)
                )

            response = self.client.chat.completions.create(
                model=self.model,
                max_tokens=500,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/png;base64,{img_b64}"}
                            },
                            {
                                "type": "text",
                                "text": f"Passo {passo}/{max_steps}. Estado atual da tela acima.{historico_str}\n\nQual é a próxima ação?"
                            }
                        ]
                    }
                ]
            )

            resposta = response.choices[0].message.content.strip()

            # Extrai JSON da resposta
            try:
                json_str = resposta
                if "```" in json_str:
                    json_str = json_str.split("```")[1].split("```")[0]
                    if json_str.startswith("json"):
                        json_str = json_str[4:]
                acao = json.loads(json_str.strip())
            except Exception:
                acao = {"acao": "falhou", "motivo": f"Resposta inválida: {resposta}"}

            descricao = acao.get("descricao", acao.get("acao", ""))
            passos.append({"passo": passo, "acao": acao})
            historico_acoes.append(f"{acao.get('acao', '?')}: {descricao}")

            if callback_status:
                callback_status(passo, descricao, acao)

            # Executa a ação
            tipo = acao.get("acao", "")

            if tipo == "concluido":
                return {
                    "concluido": True,
                    "passos": passos,
                    "resultado": acao.get("resultado", "Tarefa concluída")
                }

            elif tipo == "falhou":
                return {
                    "concluido": False,
                    "passos": passos,
                    "resultado": acao.get("motivo", "Tarefa falhou")
                }

            elif tipo == "click":
                pyautogui.click(acao["x"], acao["y"])

            elif tipo == "double_click":
                pyautogui.doubleClick(acao["x"], acao["y"])

            elif tipo == "right_click":
                pyautogui.rightClick(acao["x"], acao["y"])

            elif tipo == "type":
                pyautogui.typewrite(acao["texto"], interval=0.05)

            elif tipo == "press_key":
                tecla = acao["tecla"].lower()
                pyautogui.press(tecla)

            elif tipo == "scroll":
                x, y = acao.get("x", 500), acao.get("y", 400)
                qtd = acao.get("quantidade", 3)
                direcao = acao.get("direcao", "down")
                pyautogui.moveTo(x, y)
                pyautogui.scroll(-qtd if direcao == "down" else qtd)

            elif tipo == "wait":
                time.sleep(acao.get("segundos", 1))

            # Pequena pausa entre ações
            time.sleep(0.8)

        return {
            "concluido": False,
            "passos": passos,
            "resultado": f"Limite de {max_steps} passos atingido sem concluir"
        }

    # ==========================================================================
    # MONITORAMENTO CONTÍNUO
    # ==========================================================================

    def monitorar_tela(
        self,
        condicao: str,
        timeout: float = 30.0,
        interval: float = 1.0,
        callback: Optional[Callable] = None,
    ) -> bool:
        """
        Monitora a tela até uma condição ser verdadeira.

        Args:
            condicao  : Descrição da condição a aguardar
            timeout   : Tempo máximo em segundos
            interval  : Intervalo entre verificações
            callback  : Função chamada quando a condição for verdadeira

        Returns:
            True se a condição foi atendida, False se deu timeout

        Exemplo:
            vision.monitorar_tela("download concluído aparece na tela", timeout=60)
            vision.monitorar_tela("botão OK está visível", timeout=10)
        """
        start = time.time()

        while time.time() - start < timeout:
            resposta = self.ver_tela(
                f"""A seguinte condição está verdadeira na tela atual?
                Condição: "{condicao}"
                
                Responda APENAS: {{"verdadeiro": true}} ou {{"verdadeiro": false}}"""
            )

            try:
                json_str = resposta
                if "```" in json_str:
                    json_str = json_str.split("```")[1].split("```")[0]
                data = json.loads(json_str.strip())
                if data.get("verdadeiro"):
                    if callback:
                        callback()
                    return True
            except Exception:
                pass

            time.sleep(interval)

        return False

    def capturar_mudanca(
        self,
        region: Optional[tuple] = None,
        timeout: float = 10.0,
    ) -> Optional[str]:
        """
        Aguarda a tela mudar e descreve o que mudou.

        Args:
            region  : Região a monitorar
            timeout : Tempo máximo de espera

        Returns:
            Descrição do que mudou ou None se não houve mudança
        """
        screenshot_antes = self.capture_screen(region)
        start = time.time()

        while time.time() - start < timeout:
            time.sleep(0.5)
            screenshot_depois = self.capture_screen(region)

            if screenshot_antes != screenshot_depois:
                # Pede ao Grok para descrever a mudança
                resposta = self.client.chat.completions.create(
                    model=self.model,
                    max_tokens=500,
                    messages=[{
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/png;base64,{screenshot_depois}"}
                            },
                            {
                                "type": "text",
                                "text": "A tela acabou de mudar. Descreva o que está visível agora e o que parece ter mudado."
                            }
                        ]
                    }]
                )
                return resposta.choices[0].message.content

            screenshot_antes = screenshot_depois

        return None

    # ==========================================================================
    # MEMÓRIA DE VISÃO
    # ==========================================================================

    def get_historico(self, ultimos: int = 5) -> list[dict]:
        """
        Retorna o histórico das últimas análises visuais.

        Args:
            ultimos : Quantas entradas retornar

        Returns:
            Lista com histórico de visão
        """
        return self._vision_history[-ultimos:]

    def get_ultima_analise(self) -> Optional[dict]:
        """Retorna a última análise visual feita."""
        if self._vision_history:
            return self._vision_history[-1]
        return None

    def limpar_historico(self) -> None:
        """Limpa o histórico de visão."""
        self._vision_history.clear()

    # ==========================================================================
    # INTEGRAÇÃO COM GUI CONTROLLER
    # ==========================================================================

    def ver_e_clicar(self, descricao_elemento: str) -> bool:
        """
        Combina visão + ação: encontra um elemento e clica nele.

        Args:
            descricao_elemento : Descrição do que clicar

        Returns:
            True se encontrou e clicou

        Exemplo:
            vision.ver_e_clicar("botão Fechar no canto superior direito")
            vision.ver_e_clicar("ícone do Chrome na barra de tarefas")
        """
        pos = self.encontrar_elemento(descricao_elemento)
        if pos:
            pyautogui.click(pos[0], pos[1])
            return True
        return False

    def ver_e_digitar(self, descricao_campo: str, texto: str) -> bool:
        """
        Encontra um campo de texto via visão e digita nele.

        Args:
            descricao_campo : Descrição do campo (ex: "campo de pesquisa")
            texto           : Texto a digitar

        Returns:
            True se encontrou e digitou
        """
        pos = self.encontrar_elemento(descricao_campo)
        if pos:
            pyautogui.click(pos[0], pos[1])
            time.sleep(0.3)
            pyautogui.typewrite(texto, interval=0.05)
            return True
        return False


# ==============================================================================
# EXEMPLO DE USO
# ==============================================================================

if __name__ == "__main__":
    API_KEY = resolve_xai_api_key()
    if not API_KEY:
        raise SystemExit("Defina XAI_API_KEY ou GROK_API_KEY")
    vision = VisionLoop(api_key=API_KEY, debug=True)

    print("=== TESTE DO SISTEMA DE VISÃO ===\n")

    # 1. Descrição completa da tela
    print("1. Descrevendo a tela...")
    descricao = vision.descrever_tela_completa()
    print(descricao)
    print()

    # 2. Encontrar um elemento
    print("2. Procurando a barra de tarefas...")
    pos = vision.encontrar_elemento("barra de tarefas do Windows")
    print(f"Posição encontrada: {pos}")
    print()

    # 3. Executar uma tarefa completa com o loop
    print("3. Executando tarefa com loop de percepção-ação...")
    def status_callback(passo, descricao, acao):
        print(f"  Passo {passo}: {descricao}")

    resultado = vision.executar_tarefa(
        objetivo="Abra o Bloco de Notas",
        max_steps=5,
        callback_status=status_callback
    )
    print(f"Resultado: {resultado['resultado']}")
