"""
bridge.py
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
WebSocket Bridge — Ponte completa entre o site e o PC do usuário.

Responsabilidades:
  • Conectar ao servidor via WebSocket seguro (wss://)
  • Enviar mapa completo do ambiente ao conectar
  • Receber comandos estruturados do servidor
  • Executar via motor híbrido (UI Automation → fallback visual)
  • Retornar resultados e screenshots em tempo real
  • Reconexão automática em caso de queda
  • Heartbeat para manter a conexão viva
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
"""

import asyncio
import json
import time
import uuid
import platform
import websockets
import uiautomation as auto
import pyautogui
import mss
import base64
import io
import subprocess
import psutil

from pathlib import Path
from typing import Optional
from loguru import logger
from PIL import Image

from environment_mapper import (
    mapear_ambiente_completo,
    mapear_processos,
    mapear_janelas_ativas,
    encontrar_elemento_na_janela,
    normalizar_coordenadas,
)
from app_detector import (
    DetectorApp,
    mapear_tipos_apps_abertos,
)
from security_manager import (
    verificar_seguranca_completa,
    verificar_pode_automatizar_app,
    e_administrador,
)


# ══════════════════════════════════════════════════════
# MOTOR DE EXECUÇÃO HÍBRIDO COMPLETO
# ══════════════════════════════════════════════════════

class MotorExecucaoCompleto:
    """
    Motor que combina todas as abordagens de automação:
      1ª tentativa → UI Automation API (Win32/WPF/UWP)
      2ª tentativa → Playwright/CDP  (Browser/Electron)
      3ª tentativa → Visão computacional (Screenshot → Grok → coords)
      Último recurso → pyautogui (coordenadas brutas)
    """

    def __init__(self):
        self.detector = DetectorApp()

    # ── Abrir aplicativo ────────────────────────────
    def abrir_app(self, app: str) -> dict:
        """
        Abre um app. Verifica se já está aberto antes.
        Tenta pelo nome, apelido ou caminho direto.
        """
        from environment_mapper import encontrar_executavel_app

        # Verificar se já está rodando
        for proc in psutil.process_iter(["name"]):
            try:
                if app.lower().replace(".exe", "") in proc.info["name"].lower():
                    logger.info(f"App '{app}' já está em execução")
                    return {"status": "ja_aberto", "app": app}
            except:
                continue

        # Descobrir executável
        exe = encontrar_executavel_app(app) or app

        try:
            subprocess.Popen(exe, shell=True)
            time.sleep(2)
            logger.success(f"App '{app}' aberto")
            return {"status": "ok", "app": app, "exe": exe}
        except Exception as e:
            logger.error(f"Erro ao abrir '{app}': {e}")
            return {"status": "erro", "detalhe": str(e)}

    # ── Executar ação via UI Automation ─────────────
    def executar_via_uia(self, janela_nome: str, acao: dict) -> dict:
        """Motor principal — UI Automation API."""
        try:
            janela = auto.WindowControl(Name=janela_nome, searchDepth=1)
            if not janela.Exists(maxSearchSeconds=5):
                return {"status": "janela_nao_encontrada", "janela": janela_nome}

            tipo = acao.get("tipo", "")
            logger.debug(f"UIA — janela='{janela_nome}' tipo='{tipo}'")

            if tipo == "clicar_botao":
                el = janela.ButtonControl(Name=acao["elemento"])
                if not el.Exists(maxSearchSeconds=3):
                    return {"status": "elemento_nao_encontrado",
                            "elemento": acao["elemento"]}
                el.Click()
                return {"status": "ok", "acao": tipo}

            elif tipo == "digitar_texto":
                campo_nome = acao.get("elemento", "")
                campo = (janela.EditControl(Name=campo_nome)
                         if campo_nome else janela.EditControl())
                campo.Click()
                campo.SetValue(acao["texto"])
                return {"status": "ok", "acao": tipo}

            elif tipo == "clicar_menu":
                menu = janela.MenuControl(Name=acao["menu"])
                menu.Click()
                time.sleep(0.35)
                item = janela.MenuItemControl(Name=acao["item"])
                item.Click()
                return {"status": "ok", "acao": tipo}

            elif tipo == "selecionar_lista":
                lista = janela.ListControl(Name=acao.get("elemento", ""))
                item  = lista.ListItemControl(Name=acao["valor"])
                item.Click()
                return {"status": "ok", "acao": tipo}

            elif tipo == "pressionar_tecla":
                janela.SendKeys(acao["tecla"])
                return {"status": "ok", "acao": tipo}

            elif tipo == "marcar_checkbox":
                cb = janela.CheckBoxControl(Name=acao["elemento"])
                if cb.GetTogglePattern().ToggleState != 1:
                    cb.Click()
                return {"status": "ok", "acao": tipo}

            elif tipo == "selecionar_combobox":
                cb = janela.ComboBoxControl(Name=acao.get("elemento", ""))
                cb.Click()
                time.sleep(0.2)
                item = cb.ListItemControl(Name=acao["valor"])
                item.Click()
                return {"status": "ok", "acao": tipo}

            elif tipo == "rolar_scroll":
                direction = acao.get("direcao", "baixo")
                el = (janela.Control(Name=acao["elemento"])
                      if acao.get("elemento") else janela)
                rect = el.BoundingRectangle
                cx = rect.left + rect.width() // 2
                cy = rect.top  + rect.height() // 2
                delta = -120 if direction == "baixo" else 120
                pyautogui.scroll(delta, x=cx, y=cy)
                return {"status": "ok", "acao": tipo}

            elif tipo == "ler_texto":
                el = janela.Control(Name=acao.get("elemento", ""))
                texto = el.Name or ""
                try:
                    padrao = el.GetValuePattern()
                    if padrao:
                        texto = padrao.Value or texto
                except:
                    pass
                return {"status": "ok", "acao": tipo, "texto": texto}

            elif tipo == "verificar_elemento":
                el     = janela.Control(Name=acao["elemento"])
                existe = el.Exists(maxSearchSeconds=2)
                return {"status": "ok", "existe": existe,
                        "elemento": acao["elemento"]}

            else:
                return {"status": "tipo_desconhecido", "tipo": tipo}

        except Exception as e:
            logger.warning(f"UIA falhou: {e}")
            return {
                "status":        "falha_ui_automation",
                "erro":          str(e),
                "requer_visao":  True
            }

    # ── Fallback visual ──────────────────────────────
    def executar_visual(self, x: int, y: int, acao: str = "clicar") -> dict:
        """
        Executa ação por coordenadas de tela (último recurso).
        Coordenadas vêm do Grok Vision após análise do screenshot.
        """
        try:
            # Normalizar para DPI
            coords = normalizar_coordenadas(x, y)
            xn = coords["x_normalizado"]
            yn = coords["y_normalizado"]

            if acao == "clicar":
                pyautogui.moveTo(xn, yn, duration=0.25)
                pyautogui.click()
            elif acao == "duplo_clique":
                pyautogui.doubleClick(xn, yn)
            elif acao == "clicar_direito":
                pyautogui.rightClick(xn, yn)
            elif acao == "digitar":
                pyautogui.click(xn, yn)
                pyautogui.typewrite(
                    coords.get("texto", ""), interval=0.05
                )

            logger.success(f"Visual — {acao} em ({xn}, {yn})")
            return {"status": "ok", "acao": f"visual_{acao}",
                    "coords": [xn, yn]}
        except Exception as e:
            logger.error(f"Fallback visual falhou: {e}")
            return {"status": "falha_total", "erro": str(e)}

    # ── Capturar tela ────────────────────────────────
    def capturar_tela(self, monitor: int = 1) -> str:
        """Captura screenshot e retorna em base64 PNG."""
        try:
            with mss.mss() as sct:
                idx  = min(monitor, len(sct.monitors) - 1)
                shot = sct.grab(sct.monitors[idx])
                img  = Image.frombytes("RGB", shot.size,
                                       shot.bgra, "raw", "BGRX")
                buf  = io.BytesIO()
                img.save(buf, format="PNG", optimize=True)
                encoded = base64.b64encode(buf.getvalue()).decode()
                logger.debug(f"Screenshot capturada ({len(encoded)//1024}KB)")
                return encoded
        except Exception as e:
            logger.error(f"Erro ao capturar tela: {e}")
            return ""

    # ── Executar comando completo ────────────────────
    def executar_comando(self, cmd: dict) -> dict:
        """
        Ponto de entrada principal para qualquer comando.
        Decide a abordagem baseado no tipo do app.
        """
        tipo_cmd = cmd.get("tipo", "")

        if tipo_cmd == "abrir_app":
            return self.abrir_app(cmd.get("app", ""))

        elif tipo_cmd == "executar_acao":
            janela = cmd.get("janela", "")
            acao   = cmd.get("acao",   {})

            # Detectar tipo do app para escolher a melhor abordagem
            tipo_app = self._detectar_tipo_janela(janela)

            if tipo_app in ("WIN32", "WPF", "UWP"):
                resultado = self.executar_via_uia(janela, acao)
            else:
                # Electron, Java, desconhecido → precisará de screenshot
                resultado = {"status": "requer_visao_primario",
                             "requer_visao": True}

            # Se UIA falhou ou app requer visão → sinalizar para o servidor
            if resultado.get("requer_visao"):
                resultado["screenshot"] = self.capturar_tela()
                resultado["tipo_app"]   = tipo_app

            return resultado

        elif tipo_cmd == "fallback_visual":
            return self.executar_visual(
                cmd.get("x", 0),
                cmd.get("y", 0),
                cmd.get("acao_visual", "clicar")
            )

        return {"status": "tipo_desconhecido", "tipo": tipo_cmd}

    def _detectar_tipo_janela(self, nome_janela: str) -> str:
        """Detecta o tipo de uma janela aberta pelo nome."""
        try:
            janela = auto.WindowControl(Name=nome_janela, searchDepth=1)
            if not janela.Exists(maxSearchSeconds=2):
                return "UNKNOWN"

            pid  = janela.ProcessId
            proc = psutil.Process(pid)
            info = self.detector.detectar(
                nome_janela   = nome_janela,
                classe_janela = janela.ClassName or "",
                nome_processo = proc.name(),
                exe_path      = proc.exe() or ""
            )
            return info.tipo
        except:
            return "UNKNOWN"


# ══════════════════════════════════════════════════════
# WEBSOCKET HANDLER — SERVIDOR LOCAL
# ══════════════════════════════════════════════════════

motor = MotorExecucaoCompleto()


async def handler(websocket):
    """
    Gerencia a conexão de um servidor remoto.
    Envia o mapa completo do ambiente ao conectar.
    """
    cliente_id = str(uuid.uuid4())[:8]
    logger.info(f"[{cliente_id}] Servidor conectado: {websocket.remote_address}")

    # ── Envio inicial do ambiente completo ──────────
    try:
        logger.info(f"[{cliente_id}] Mapeando ambiente completo...")

        ambiente  = mapear_ambiente_completo(com_elementos_filhos=False)
        tipos     = mapear_tipos_apps_abertos()
        seguranca = verificar_seguranca_completa()

        payload_inicial = {
            "tipo":      "ambiente_inicial",
            "client_id": cliente_id,
            "ambiente":  ambiente,
            "tipos_apps": tipos,
            "seguranca": {
                "e_admin":             seguranca.e_admin,
                "uac_ativo":           seguranca.uac_ativo,
                "nivel_integridade":   seguranca.nivel_integridade,
                "antivirus":           [av["nome"] for av in seguranca.antivirus_detectado],
                "alertas":             seguranca.alertas,
                "recomendacoes":       seguranca.recomendacoes,
            }
        }

        await websocket.send(json.dumps(payload_inicial, ensure_ascii=False))
        logger.success(f"[{cliente_id}] Ambiente enviado ao servidor")

    except Exception as e:
        logger.error(f"[{cliente_id}] Erro ao enviar ambiente: {e}")
        await websocket.send(json.dumps({
            "tipo":  "erro_ambiente",
            "erro":  str(e)
        }))

    # ── Loop principal de comandos ──────────────────
    try:
        async for raw in websocket:
            try:
                cmd    = json.loads(raw)
                tipo   = cmd.get("tipo", "")
                req_id = cmd.get("req_id", str(uuid.uuid4())[:8])

                logger.info(f"[{cliente_id}][{req_id}] Comando: {tipo}")
                inicio = time.time()

                # ── Executar ação (abrir app / automação)
                if tipo in ("abrir_app", "executar_acao", "fallback_visual"):
                    resultado = motor.executar_comando(cmd)
                    resultado["req_id"]  = req_id
                    resultado["duracao"] = round(time.time() - inicio, 3)
                    await websocket.send(json.dumps(resultado, ensure_ascii=False))

                # ── Capturar tela sob demanda
                elif tipo == "capturar_tela":
                    screenshot = motor.capturar_tela(cmd.get("monitor", 1))
                    await websocket.send(json.dumps({
                        "tipo":      "screenshot",
                        "req_id":    req_id,
                        "imagem":    screenshot,
                        "timestamp": time.strftime("%H:%M:%S")
                    }))

                # ── Re-mapear ambiente
                elif tipo == "mapear_ambiente":
                    com_filhos = cmd.get("com_elementos_filhos", False)
                    ambiente   = mapear_ambiente_completo(com_filhos)
                    await websocket.send(json.dumps({
                        "tipo":    "ambiente",
                        "req_id":  req_id,
                        "dados":   ambiente
                    }, ensure_ascii=False))

                # ── Mapear janelas (snapshot rápido)
                elif tipo == "mapear_janelas":
                    janelas = mapear_janelas_ativas(
                        com_elementos_filhos = cmd.get("com_filhos", False)
                    )
                    await websocket.send(json.dumps({
                        "tipo":    "janelas",
                        "req_id":  req_id,
                        "janelas": janelas
                    }, ensure_ascii=False))

                # ── Encontrar elemento específico
                elif tipo == "encontrar_elemento":
                    elemento = encontrar_elemento_na_janela(
                        cmd.get("janela",   ""),
                        cmd.get("elemento", ""),
                        cmd.get("tipo_elemento")
                    )
                    await websocket.send(json.dumps({
                        "tipo":     "elemento",
                        "req_id":   req_id,
                        "encontrado": elemento is not None,
                        "dados":    elemento
                    }, ensure_ascii=False))

                # ── Verificar segurança de um processo
                elif tipo == "verificar_permissao":
                    perm = verificar_pode_automatizar_app(
                        cmd.get("processo", "")
                    )
                    await websocket.send(json.dumps({
                        "tipo":   "permissao",
                        "req_id": req_id,
                        "dados":  perm
                    }))

                # ── Processos em tempo real
                elif tipo == "listar_processos":
                    processos = mapear_processos()
                    await websocket.send(json.dumps({
                        "tipo":      "processos",
                        "req_id":    req_id,
                        "processos": processos
                    }, ensure_ascii=False))

                # ── Tipos de apps abertos
                elif tipo == "tipos_apps":
                    tipos_apps = mapear_tipos_apps_abertos()
                    await websocket.send(json.dumps({
                        "tipo":       "tipos_apps",
                        "req_id":     req_id,
                        "apps":       tipos_apps
                    }, ensure_ascii=False))

                # ── Ping / keepalive
                elif tipo == "ping":
                    await websocket.send(json.dumps({
                        "tipo":   "pong",
                        "req_id": req_id,
                        "ts":     time.time(),
                        "admin":  e_administrador()
                    }))

                else:
                    logger.warning(f"[{cliente_id}] Comando desconhecido: {tipo}")
                    await websocket.send(json.dumps({
                        "status":  "tipo_desconhecido",
                        "tipo":    tipo,
                        "req_id":  req_id
                    }))

            except json.JSONDecodeError as e:
                logger.error(f"[{cliente_id}] JSON inválido: {e}")
                await websocket.send(json.dumps({
                    "status": "erro_json",
                    "detalhe": str(e)
                }))

            except Exception as e:
                logger.error(f"[{cliente_id}] Erro ao processar comando: {e}")
                await websocket.send(json.dumps({
                    "status":  "erro_interno",
                    "detalhe": str(e)
                }))

    except websockets.exceptions.ConnectionClosed as e:
        logger.warning(f"[{cliente_id}] Conexão encerrada: {e.code}")
    except Exception as e:
        logger.error(f"[{cliente_id}] Erro inesperado: {e}")

    logger.info(f"[{cliente_id}] Handler encerrado")


# ══════════════════════════════════════════════════════
# SERVIDOR WEBSOCKET LOCAL
# ══════════════════════════════════════════════════════

async def iniciar_servidor_local(porta: int = 8765,
                                  host: str  = "localhost"):
    """Inicia o servidor WebSocket que escuta conexões do relay no servidor."""
    logger.info(f"Iniciando servidor WebSocket local em ws://{host}:{porta}")

    async with websockets.serve(
        handler,
        host,
        porta,
        ping_interval   = 30,   # heartbeat a cada 30s
        ping_timeout    = 10,   # timeout de 10s para pong
        max_size        = 50 * 1024 * 1024,  # 50MB (para screenshots)
        compression     = None
    ):
        logger.success(f"✅ WebSocket local ativo em ws://{host}:{porta}")
        await asyncio.Future()  # roda para sempre


# ══════════════════════════════════════════════════════
# CLIENTE WEBSOCKET — CONECTA AO SERVIDOR REMOTO
# (modo alternativo: o agente conecta no servidor, não ao contrário)
# ══════════════════════════════════════════════════════

async def conectar_ao_servidor(server_url: str,
                                reconectar: bool = True,
                                intervalo_reconexao: int = 5):
    """
    Modo alternativo: agente conecta no servidor remoto.
    Útil quando o servidor não consegue alcançar o cliente (NAT, firewall).
    """
    while True:
        try:
            logger.info(f"Conectando ao servidor: {server_url}")

            async with websockets.connect(
                server_url,
                ping_interval = 20,
                ping_timeout  = 10,
                max_size      = 50 * 1024 * 1024,
            ) as ws:
                logger.success(f"✅ Conectado ao servidor: {server_url}")

                # Enviar ambiente inicial
                ambiente = mapear_ambiente_completo(com_elementos_filhos=False)
                seg      = verificar_seguranca_completa()

                await ws.send(json.dumps({
                    "tipo":     "registrar_agente",
                    "ambiente": ambiente,
                    "seguranca": {
                        "e_admin":   seg.e_admin,
                        "alertas":   seg.alertas,
                    }
                }, ensure_ascii=False))

                # Loop de comandos
                async for raw in ws:
                    cmd       = json.loads(raw)
                    req_id    = cmd.get("req_id", "")
                    resultado = motor.executar_comando(cmd)
                    resultado["req_id"] = req_id
                    await ws.send(json.dumps(resultado, ensure_ascii=False))

        except (websockets.exceptions.ConnectionClosed,
                OSError, ConnectionRefusedError) as e:
            logger.warning(f"Conexão perdida: {e}")

            if not reconectar:
                break

            logger.info(f"Reconectando em {intervalo_reconexao}s...")
            await asyncio.sleep(intervalo_reconexao)

        except Exception as e:
            logger.error(f"Erro inesperado: {e}")
            if not reconectar:
                break
            await asyncio.sleep(intervalo_reconexao)


# ══════════════════════════════════════════════════════
# ENTRY POINT
# ══════════════════════════════════════════════════════

def iniciar(porta: int = 8765, server_url: str = ""):
    """
    Inicia o bridge.
    Se server_url for fornecido → conecta no servidor remoto.
    Caso contrário → inicia servidor local aguardando conexões.
    """
    if server_url:
        asyncio.run(conectar_ao_servidor(server_url))
    else:
        asyncio.run(iniciar_servidor_local(porta))


if __name__ == "__main__":
    iniciar()
