# ============================================================
# ORCHESTRATOR.PY - Orquestrador de Tarefas
# Planner (Gemini/OpenAI/Claude/Ollama) + Executor (Qwen/Ollama)
# ============================================================

import json
import re
from typing import Dict, List, Optional
from llm_providers import (
    call_planner, call_executor, get_planner_llm, get_executor_llm,
    get_available_models, LLM_PROVIDERS, set_planner, set_executor, get_config, save_config
)
from tools_v2 import TOOLS_V2, get_tools_by_danger, get_all_tools_formatted
from credential_manager import sync_credentials, get_services_status

# ============================================================
# SYSTEM PROMPTS
# ============================================================

PLANNER_SYSTEM_PROMPT = """Você é o PLANNER AGENT do BotVPS.
Seu trabalho é decompor tarefas em passos executáveis.

### SUAS TAREFAS:
1. Entender a intenção do usuário
2. Decompor em passos menores e claros
3. Classificar cada passo como SAFE, MEDIUM ou DANGEROUS
4. Usar as ferramentas disponíveis listadas abaixo

### NÍVEIS DE PERIGO:
- SAFE: Pode executar automaticamente (listar, ver status, ler logs)
- MEDIUM: Informa o usuário antes (git pull, build, restart)
- DANGEROUS: REQUER confirmação explícita (delete, reboot, docker down)

### FERRAMENTAS DISPONÍVEIS:
{TOOLS_LIST}

### FORMATO DE RESPOSTA:
Responda APENAS com JSON no seguinte formato:
{{
  "task_name": "Nome resumido da tarefa",
  "summary": "Resumo do que será feito",
  "steps": [
    {{
      "order": 1,
      "action": "Descrição clara do que fazer",
      "tool": "nome_da_ferramenta (ou null se for bash)",
      "command": "comando específico a executar",
      "danger": "safe|medium|dangerous"
    }}
  ]
}}

### REGRAS:
1. Responda APENAS com JSON válido
2. Cada passo deve ser atômico (uma ação por passo)
3. Considere dependências entre passos
4. Para passos bash complexos, use tool="bash" e command="comando"
5. Os passos devem ser na ordem correta de execução
"""

EXECUTOR_SYSTEM_PROMPT = """Você é o EXECUTOR AGENT do BotVPS.
Seu trabalho é executar comandos bash com precisão.

### REGRAS:
1. Execute APENAS o comando passado
2. Retorne o output do comando
3. Se houver erro, descreva o erro claramente
4. Não invente outputs

### FORMATO DE RESPOSTA:
Responda com JSON:
{{
  "success": true|false,
  "output": "output do comando ou erro"
}}

### IMPORTANTE:
- Use caminhos absolutos quando possível
- Redirecione erros (2>/dev/null) quando apropriado
- Mantenha comandos simples e seguros
"""

# ============================================================
# HELPER FUNCTIONS
# ============================================================

def _format_tools_for_prompt() -> str:
    """Formata lista de ferramentas para o prompt."""
    lines = []
    for name, info in TOOLS_V2.items():
        lines.append(f"- {name}: {info['desc']} [{info['danger']}]")
    return "\n".join(lines)

def _parse_json_response(text: str) -> Optional[Dict]:
    """Extrai JSON da resposta do LLM."""
    # Tenta encontrar JSON no texto
    json_match = re.search(r'\{[\s\S]*\}', text)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass
    
    # Tenta extrair de blocos de código
    code_blocks = re.findall(r'```(?:json)?\s*([\s\S]*?)```', text)
    for block in code_blocks:
        try:
            return json.loads(block.strip())
        except json.JSONDecodeError:
            continue
    
    return None

def _classify_dangerous_steps(steps: List[Dict]) -> List[Dict]:
    """Retorna apenas passos perigosos."""
    return [s for s in steps if s.get("danger") in ["medium", "dangerous"]]

# ============================================================
# PLANNER AGENT
# ============================================================

def plan_task(task: str) -> Dict:
    """
    Usa o Planner LLM para decompor uma tarefa.
    
    Args:
        task: Tarefa do usuário
    
    Returns:
        Dicionário com plano de execução:
        {
            "task_name": str,
            "summary": str,
            "steps": [
                {"order": int, "action": str, "tool": str, "command": str, "danger": str}
            ]
        }
    """
    provider, model = get_planner_llm()
    print(f"[PLANNER] Using: {provider}/{model}")
    
    system_prompt = PLANNER_SYSTEM_PROMPT.replace("{TOOLS_LIST}", _format_tools_for_prompt())
    
    response = call_planner(task, system_prompt)
    print(f"[RESPONSE] Planner response:\n{response[:500]}...")
    
    plan = _parse_json_response(response)
    
    if not plan or "steps" not in plan:
        # Fallback: tenta executar como comando único
        return {
            "task_name": task[:50],
            "summary": f"Tarefa: {task}",
            "steps": [{
                "order": 1,
                "action": task,
                "tool": "bash",
                "command": task,
                "danger": "medium"
            }]
        }
    
    return plan

# ============================================================
# EXECUTOR AGENT
# ============================================================

def execute_command(command: str) -> Dict:
    """
    Executa um comando bash via Executor LLM.
    
    Args:
        command: Comando a executar
    
    Returns:
        {"success": bool, "output": str}
    """
    import subprocess
    
    provider, model = get_executor_llm()
    print(f"[EXECUTOR] Using: {provider}/{model}")
    
    # Para comandos bash simples, executa direto sem LLM
    # Usa LLM apenas para comandos complexos
    if len(command) < 100 and not any(c in command for c in ["&&", "||", "|", "$"]):
        try:
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=60
            )
            return {
                "success": result.returncode == 0,
                "output": result.stdout.strip() or result.stderr.strip() or "Sucesso"
            }
        except Exception as e:
            return {"success": False, "output": str(e)}
    
    # Para comandos complexos, usa LLM
    response = call_executor(
        f"Execute este comando e retorne o resultado em JSON: {command}",
        EXECUTOR_SYSTEM_PROMPT
    )
    
    result = _parse_json_response(response)
    if result:
        return result
    
    return {"success": False, "output": response}

def execute_step(step: Dict) -> Dict:
    """
    Executa um passo do plano.
    
    Args:
        step: Dicionário com dados do passo
    
    Returns:
        {"success": bool, "output": str, "step": int}
    """
    tool = step.get("tool")
    command = step.get("command", "")
    order = step.get("order", 0)
    
    print(f"  -> Step {order}: {step.get('action')[:50]}...")
    
    if tool and tool in TOOLS_V2:
        try:
            tool_info = TOOLS_V2[tool]
            func = tool_info["func"]
            
            # Executa a função da ferramenta
            if callable(func):
                result = func(command) if command else func()
            else:
                result = str(func)
            
            return {
                "success": True,
                "output": result,
                "step": order
            }
        except Exception as e:
            return {
                "success": False,
                "output": f"Erro ao executar {tool}: {str(e)}",
                "step": order
            }
    
    # Executa como comando bash
    return execute_command(command)

# ============================================================
# ORCHESTRATOR MAIN
# ============================================================

def orchestrate(task: str, user_confirmed: bool = False) -> Dict:
    """
    Orquestra a execução de uma tarefa.
    
    Args:
        task: Tarefa do usuário
        user_confirmed: Se True, pula confirmação e executa tudo
    
    Returns:
        {
            "status": "needs_confirmation" | "completed" | "error",
            "plan": {...},
            "confirmation_needed_for": [steps peligrosos],
            "results": [...] (se status == "completed")
        }
    """
    print(f"\n{'='*50}")
    print(f">>> PLANNING: {task}")
    print(f"{'='*50}\n")
    
    # 1. Plana a tarefa
    plan = plan_task(task)
    
    # 2. Identifica passos perigosos
    dangerous_steps = _classify_dangerous_steps(plan.get("steps", []))
    
    # 3. Se há passos perigosos e não confirmou, pede confirmação
    if dangerous_steps and not user_confirmed:
        return {
            "status": "needs_confirmation",
            "plan": plan,
            "confirmation_needed_for": [
                {"order": s["order"], "action": s["action"], "danger": s["danger"]}
                for s in dangerous_steps
            ]
        }
    
    # 4. Executa todos os passos
    results = []
    for step in plan.get("steps", []):
        result = execute_step(step)
        results.append(result)
        
        # Para em caso de erro crítico
        if not result.get("success") and step.get("danger") == "dangerous":
            results.append({
                "success": False,
                "output": "Execução abortada devido a erro crítico.",
                "step": -1
            })
            break
    
    # 5. Retorna resultado
    return {
        "status": "completed",
        "plan": plan,
        "results": results
    }

def format_confirmation_message(result: Dict) -> str:
    """
    Formata mensagem de confirmação para o usuário.
    
    Args:
        result: Resultado do orchestrate()
    
    Returns:
        String formatada para envio ao usuário
    """
    if result["status"] != "needs_confirmation":
        return ""
    
    plan = result["plan"]
    dangerous = result["confirmation_needed_for"]
    
    msg = f"[PLANO] {plan.get('task_name', 'Tarefa')}\n\n"
    msg += f"{plan.get('summary', '')}\n\n"
    
    msg += "AVISO: Acoes que precisam de confirmacao:\n\n"
    
    for step in dangerous:
        icon = "[CRITICAL]" if step["danger"] == "dangerous" else "[WARNING]"
        msg += f"{icon} Passo {step['order']}: {step['action']}\n"
    
    msg += "\nDeseja continuar? (sim/não)"
    
    return msg

def format_completion_message(result: Dict) -> str:
    """
    Formata mensagem de conclusão.
    
    Args:
        result: Resultado do orchestrate()
    
    Returns:
        String formatada com os resultados
    """
    if result["status"] != "completed":
        return ""
    
    plan = result["plan"]
    results = result.get("results", [])
    
    msg = f"[OK] Concluido: {plan.get('task_name', 'Tarefa')}\n\n"
    
    success_count = sum(1 for r in results if r.get("success"))
    total_count = len([r for r in results if r.get("step", 0) > 0])
    
    msg += f"[STAT] Resultado: {success_count}/{total_count} passos executados com sucesso.\n\n"
    
    for r in results:
        step_num = r.get("step", 0)
        if step_num > 0:
            status = "[OK]" if r.get("success") else "[FAIL]"
            output = r.get("output", "")[:500]
            msg += f"{status} Passo {step_num}:\n```\n{output}\n```\n\n"
    
    return msg

# ============================================================
# STATUS & CONFIG FUNCTIONS
# ============================================================

def get_orchestrator_status() -> Dict:
    """Retorna status atual do orquestrador."""
    planner_provider, planner_model = get_planner_llm()
    executor_provider, executor_model = get_executor_llm()
    
    return {
        "planner": {
            "provider": planner_provider,
            "model": planner_model,
            "name": LLM_PROVIDERS[planner_provider]["name"]
        },
        "executor": {
            "provider": executor_provider,
            "model": executor_model,
            "name": LLM_PROVIDERS[executor_provider]["name"]
        },
        "credentials": get_services_status(),
        "available_tools": len(TOOLS_V2)
    }

def get_llm_config() -> Dict:
    """Retorna configuração de LLMs."""
    planner_provider, planner_model = get_planner_llm()
    executor_provider, executor_model = get_executor_llm()
    
    return {
        "planner": {
            "provider": planner_provider,
            "model": planner_model,
            "available_providers": [
                {"id": k, "name": v["name"], "type": v["type"]}
                for k, v in LLM_PROVIDERS.items()
            ]
        },
        "executor": {
            "provider": executor_provider,
            "model": executor_model,
            "available_providers": [
                {"id": k, "name": v["name"], "type": v["type"]}
                for k, v in LLM_PROVIDERS.items()
            ]
        }
    }

def set_llm_config(planner_provider: str = None, planner_model: str = None,
                   executor_provider: str = None, executor_model: str = None) -> Dict:
    """Atualiza configuração de LLMs."""
    changes = {}
    
    if planner_provider:
        result = set_planner(planner_provider, planner_model)
        changes["planner"] = result
    
    if executor_provider:
        result = set_executor(executor_provider, executor_model)
        changes["executor"] = result
    
    return changes

# ============================================================
# COMMAND PARSER (para Telegram/Web)
# ============================================================

def parse_command(text: str) -> Dict:
    """
    Interpreta comandos do usuário.
    
    Args:
        text: Texto do usuário
    
    Returns:
        {"type": "orchestrate"|"config"|"status", "data": {...}}
    """
    text = text.strip().lower()
    
    # Comandos de configuração
    if text.startswith("/llm"):
        parts = text.split()
        if len(parts) == 1:
            return {"type": "config", "action": "show"}
        elif len(parts) >= 3:
            if parts[1] == "planner":
                return {"type": "config", "action": "set_planner", "provider": parts[2]}
            elif parts[1] == "executor":
                return {"type": "config", "action": "set_executor", "provider": parts[2]}
        return {"type": "config", "action": "help"}
    
    if text == "/sync":
        return {"type": "config", "action": "sync_credentials"}
    
    if text == "/status":
        return {"type": "status"}
    
    if text == "/tools":
        return {"type": "tools"}
    
    if text.startswith("/"):
        return {"type": "unknown", "command": text}
    
    # Tarefas de orquestração
    return {"type": "orchestrate", "task": text}

# ============================================================
# MAIN HANDLER
# ============================================================

def handle_message(text: str, confirmed: bool = False) -> str:
    """
    Manipula mensagem do usuário.
    
    Args:
        text: Mensagem do usuário
        confirmed: Se o usuário já confirmou ações perigosas
    
    Returns:
        Resposta para o usuário
    """
    parsed = parse_command(text)
    
    # Status
    if parsed["type"] == "status":
        status = get_orchestrator_status()
        msg = "[BOT] Status do Orquestrador:\n\n"
        msg += f"[PLANNER] {status['planner']['name']} ({status['planner']['model']})\n"
        msg += f"[EXECUTOR] {status['executor']['name']} ({status['executor']['model']})\n"
        msg += f"[TOOLS] Ferramentas: {status['available_tools']}\n"
        return msg
    
    # Config
    if parsed["type"] == "config":
        if parsed["action"] == "show":
            config = get_llm_config()
            msg = "[CONFIG] Configuracao de LLMs:\n\n"
            msg += f"[PLANNER] {config['planner']['provider']} / {config['planner']['model']}\n"
            msg += f"[EXECUTOR] {config['executor']['provider']} / {config['executor']['model']}\n"
            msg += "\nPara mudar: /llm planner <provider> ou /llm executor <provider>"
            return msg
        
        if parsed["action"] == "sync_credentials":
            result = sync_credentials()
            return f"[SYNC] Credenciais sincronizadas: {result['status']}"
        
        return "[CONFIG] Use: /llm (mostrar) | /llm planner <provider> | /llm executor <provider>"
    
    # Tools
    if parsed["type"] == "tools":
        return get_all_tools_formatted()
    
    # Orchestrate
    if parsed["type"] == "orchestrate":
        task = parsed["task"]
        result = orchestrate(task, confirmed)
        
        if result["status"] == "needs_confirmation":
            return format_confirmation_message(result)
        
        return format_completion_message(result)
    
    # Unknown
    return "[?] Comando nao reconhecido. Tente: /llm, /status, /tools ou descreva uma tarefa."