"""
Phase 89: The COO Orchestrator Node.
This agent sits above the standard swarm hierarchy. It has unrestricted access to the
local system (via the supervisor and subprocesses) to manage the swarm, restart dead
agents, enforce cross-project boundaries, and parse high-level CEO directives.
"""

import json
import logging
import os
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Any

from brynq.agent_loop import BrynqAgent
from brynq.supervisor import supervisor
from brynq.server import _cleanup_stale_sessions_sync, CHANNELS_DIR, SESSIONS_DIR, _append_jsonl_sync
from brynq.task_granularity import decompose_task, create_subtasks
from brynq.auto_dispatch import find_best_agent

try:
    from brynq.company_roster import get_roster
except ImportError:
    def get_roster(): return {}

log = logging.getLogger(__name__)

class OrchestratorAgent(BrynqAgent):
    """
    The COO Node. Bypasses sandboxing and has the authority to spawn/kill other agents,
    manage API fallbacks, and police the broker channels for contamination.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.channel = "ceo-desk"
        self.system_prompt = (
            "You are the COO Orchestrator Node. You manage the entire Brynq swarm. "
            "You sit above the standard workers. You have the authority to kill processes, "
            "spawn new agents from the Company Roster, and enforce boundaries. "
            "The human CEO will give you high-level directives here in #ceo-desk. "
            "When given a large project, you have three main protocols:\n"
            "1. 'hive_mind_shatter' - parallel swarming for loosely coupled tasks.\n"
            "2. 'build_and_start_workflow' - for strictly ordered DAG tasks (e.g. Code -> Review -> Test).\n"
            "3. 'run_mapreduce' - for massively parallel data processing across the Compute Grid.\n"
            "IMPORTANT: You are running locally on the user's machine. You have FULL access "
            f"to the local file system and the current workspace directory ({os.getcwd()}). "
            "You can execute shell commands by wrapping them in ```bash ... ``` blocks. "
            "You can execute python code by wrapping it in ```python ... ``` blocks. "
            "For example:\n"
            "```python\n"
            "self.build_and_start_workflow(\n"
            "    workflow_name='build_api',\n"
            "    description='API Dev',\n"
            "    steps=[\n"
            "        {'id': 'code', 'description': 'Write the code'},\n"
            "        {'id': 'review', 'dependencies': ['code'], 'description': 'Review it'}\n"
            "    ]\n"
            ")\n"
            "```\n"
            "The system will run your code and feed the output back to you in the next message. "
            "Do NOT tell the user that you are a cloud-based AI without local access. "
            "You are a local process with the ability to read files, run scripts, and coordinate the swarm."
        )
        self._last_cleanup = 0
        self._message_velocity = {}  

    def _call_backend(self, prompt: str) -> str:
        """Call backend and intercept any python/bash execution blocks before posting."""
        
        # Phase 105: Vectorized Context Injection
        try:
            from pathlib import Path
            from runtime.memory.manager import MemoryManager
            # Use the global brynq data dir
            data_dir = Path.home() / ".brynq"
            if not data_dir.exists():
                data_dir.mkdir(parents=True, exist_ok=True)
                
            mm = MemoryManager(data_dir)
            
            # Transparently retrieve the top relevant facts, past bugs, and codebase idioms
            context = mm.context_for_prompt(prompt, max_tokens=2000)
            if context and context.strip():
                log.info(f"Orchestrator: Injected {len(context)} chars of vectorized context from Hive Mind.")
                prompt = f"SYSTEM CONTEXT FROM HIVE MIND (Adhere to these past constraints):\n{context}\n\nUSER DIRECTIVE:\n{prompt}"
        except ImportError:
            pass
        except Exception as e:
            log.debug(f"Failed to inject vectorized context: {e}")
            
        response = super()._call_backend(prompt)
        
        # Intercept and execute python blocks locally
        if "```python" in response:
            import re
            import sys
            import io
            import contextlib
            
            blocks = re.findall(r"```python\n(.*?)\n```", response, re.DOTALL)
            output = ""
            for block in blocks:
                log.info("Orchestrator executing python block...")
                old_stdout = sys.stdout
                redirected_output = sys.stdout = io.StringIO()
                try:
                    # Provide 'self' so the script can call self.hive_mind_shatter() or self.build_and_start_workflow()
                    exec(block, {"self": self, "log": log, "os": os, "supervisor": supervisor})
                    output += redirected_output.getvalue()
                except Exception as e:
                    output += f"Execution Error: {str(e)}\n"
                finally:
                    sys.stdout = old_stdout
                    
            if output:
                # Append execution results to the response
                response += f"\n\n[SYSTEM: Python Execution Results]\n```\n{output}\n```"
                
        return response

    def run_mapreduce(self, description: str, data_chunks: list, map_script: str, reduce_script: str):
        """
        Executes a MapReduce job across the Swarm Compute Grid.
        `map_script` and `reduce_script` should be python code strings defining `def map_fn(chunk):` and `def reduce_fn(results):`.
        """
        log.info(f"COO initiating MapReduce job: {description}")
        try:
            from brynq.compute_grid import SwarmMapReduce, register_map_function, register_reduce_function
        except ImportError:
            log.error("Compute grid module not found.")
            return

        map_env = {}
        reduce_env = {}
        try:
            exec(map_script, globals(), map_env)
            exec(reduce_script, globals(), reduce_env)
        except Exception as e:
            log.error(f"Failed to compile MapReduce scripts: {e}")
            return {"error": f"Script compilation failed: {e}"}

        map_fn = map_env.get("map_fn")
        reduce_fn = reduce_env.get("reduce_fn")

        if not map_fn or not reduce_fn:
            log.error("map_script must define map_fn and reduce_script must define reduce_fn")
            return {"error": "Missing map_fn or reduce_fn definition"}

        map_name = f"map_dyn_{os.urandom(4).hex()}"
        reduce_name = f"reduce_dyn_{os.urandom(4).hex()}"

        register_map_function(map_name, map_fn)
        register_reduce_function(reduce_name, reduce_fn)

        smr = SwarmMapReduce()
        try:
            job = smr.submit(description, data_chunks, map_name, reduce_name)
            job_id = job["job_id"]
            
            # Start background thread to run map and reduce phases
            def _run_phases():
                try:
                    log.info(f"Starting Map phase for {job_id}")
                    smr.map_phase(job_id)
                    log.info(f"Starting Reduce phase for {job_id}")
                    result = smr.reduce_phase(job_id)
                    log.info(f"MapReduce job {job_id} completed successfully.")
                except Exception as e:
                    log.error(f"MapReduce job {job_id} failed: {e}")
                    
            import threading
            threading.Thread(target=_run_phases, daemon=True, name=f"MapReduce_{job_id}").start()
            
            return {"status": "submitted", "job_id": job_id}
        except Exception as e:
            log.error(f"MapReduce execution failed: {e}")
            return {"error": str(e)}

    def build_and_start_workflow(self, workflow_name: str, description: str, steps: list[dict]):
        """
        Creates and starts a formal DAG workflow.
        Steps should be a list of dicts like: 
        [{"id": "step1", "description": "do X"}, {"id": "step2", "dependencies": ["step1"], "description": "do Y"}]
        """
        log.info(f"COO building new workflow DAG: {workflow_name}")
        try:
            from brynq.workflow_engine import create_workflow, start_workflow
        except ImportError:
            try:
                from brynq.workflow import create_workflow, start_workflow
            except ImportError:
                log.error("Workflow engine not found. Cannot build workflow.")
                return
                
        wf = create_workflow(name=workflow_name, description=description, steps=steps)
        wf_id = wf.get("id") or wf.get("workflow_id")
        run = start_workflow(wf_id)
        
        # Start the background workflow executor if not already running
        def _run_executor():
            try:
                from brynq.workflow_executor import WorkflowExecutor
                executor = WorkflowExecutor(backend="ollama", model="llama3", dispatch_mode="task")
                executor.run(max_iterations=50) # Run for a while to handle this workflow
            except Exception as e:
                log.error(f"Executor failed: {e}")
                
        import threading
        threading.Thread(target=_run_executor, daemon=True, name="WorkflowExecutorThread").start()
        
        log.info(f"Workflow {workflow_name} started with Run ID: {run.get('run_id')}")
        return run

    def hive_mind_shatter(self, parent_task_id: str, title: str, description: str):
        """
        The core mechanism for parallel swarming. 
        1. Shatters a large task into atomic micro-tasks.
        2. Assigns specialized roles from the company roster to each micro-task.
        3. Spawns those specific agents dynamically to swarm the codebase.
        """
        log.info(f"COO initiating Hive Mind Shatter on task: {title}")
        
        # Phase 35: Swarm State Machine
        swarm_session_id = f"swarm-{parent_task_id[:8]}"
        try:
            try:
                from brynq.swarm_state_machine import create_swarm_session, transition, assign_agent
            except ImportError:
                from brynq._future.swarm_state_machine import create_swarm_session, transition, assign_agent
                
            create_swarm_session(
                name=f"Shatter: {title[:20]}",
                task_description=description,
                max_agents=10
            )
            transition(swarm_session_id, "running", "Hive mind shatter initiated")
        except ImportError:
            swarm_session_id = None
            pass
        except Exception as e:
            log.error(f"Failed to initialize Swarm State Machine: {e}")
            swarm_session_id = None
        
        # 1. Break the task down using Phase 75 engine
        decomposition = decompose_task(title, description, max_subtasks=10)
        subtasks = decomposition.get("subtasks", [])
        
        if not subtasks:
            log.warning("Task too small to shatter.")
            return

        # Write subtasks to broker
        create_subtasks(parent_task_id, [s["description"] for s in subtasks])
        
        # 2. Analyze which roles we need for these subtasks
        # Phase 103: Dynamic Swarm Roster
        try:
            from brynq.agent_injector import get_dynamic_roster, discover_and_inject_agent
            roster = get_dynamic_roster()
            can_inject = True
        except ImportError:
            roster = get_roster()
            can_inject = False
            
        spawned_roles = set()
        
        for st in subtasks:
            desc = st["description"]
            # Simplified capability mapping based on task keywords
            required_capability = "general"
            if "ui" in desc.lower() or "frontend" in desc.lower() or "react" in desc.lower() or "svelte" in desc.lower():
                required_capability = "frontend"
            elif "database" in desc.lower() or "sql" in desc.lower() or "schema" in desc.lower():
                required_capability = "database"
            elif "test" in desc.lower() or "qa" in desc.lower() or "pytest" in desc.lower():
                required_capability = "testing"
            elif "security" in desc.lower() or "auth" in desc.lower() or "vault" in desc.lower():
                required_capability = "security"
            elif "api" in desc.lower() or "backend" in desc.lower() or "endpoint" in desc.lower():
                required_capability = "backend"
            elif "unreal" in desc.lower() or "cpp" in desc.lower() or "c++" in desc.lower():
                required_capability = "cpp"

            # Find a matching persona in the roster
            target_persona = None
            for key, p in roster.items():
                if required_capability in p.get("capabilities", []) or required_capability in p.get("skills", []):
                    target_persona = p
                    break
                    
            # Phase 103: If we don't have the agent locally, buy it from the marketplace
            if not target_persona and can_inject:
                new_agent_key = discover_and_inject_agent(missing_skill=required_capability)
                if new_agent_key:
                    # Reload the roster to get the newly injected agent
                    roster = get_dynamic_roster()
                    target_persona = roster.get(new_agent_key)
            
            # 3. Spawn the specific agent if we haven't already for this job
            if target_persona and target_persona["title"] not in spawned_roles:
                agent_name = f"{target_persona['title'].replace(' ', '')}-{os.urandom(2).hex()}"
                log.info(f"COO spawning specialized agent: {agent_name} for capability: {required_capability}")
                
                # Register in Swarm State Machine
                if swarm_session_id:
                    try:
                        assign_agent(swarm_session_id, f"agent-{agent_name}", role=required_capability)
                    except Exception as e:
                        pass
                
                base_rules = (
                    "\n\n## Rules\n"
                    "- Read your tasks from the broker channel. Respond to requests.\n"
                    "- Post your work results back to the channel.\n"
                    "- You can execute shell commands by wrapping them in ```bash ... ``` blocks.\n"
                    "- You can execute python code by wrapping it in ```python ... ``` blocks.\n"
                    "- The system will run your code and feed the output back to you in the next message.\n"
                    "- Follow the Brynq Constitution. Claim files before editing.\n"
                )
                
                supervisor.start_agent(
                    name=agent_name,
                    backend="claude", # Use high-tier compute for complex broken tasks
                    model="claude-3-haiku-20240307", # Fast/cheap for micro-tasks
                    channel="tasks",
                    system_prompt=target_persona["prompt"] + base_rules + f"\n\nYou were dynamically spawned by the COO to handle subtasks for project: {title}. Look for pending tasks in your capability domain."
                )
                spawned_roles.add(target_persona["title"])
                
        # Announce the swarm mobilization
        if CHANNELS_DIR.exists():
            channel_file = CHANNELS_DIR / "general.jsonl"
            _append_jsonl_sync(channel_file, {
                "id": os.urandom(6).hex(),
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "session_id": "COO-Node",
                "session_name": "The-COO",
                "platform": "system",
                "channel": "general",
                "msg_type": "alert",
                "message": f"HIVE MIND DEPLOYED: Project '{title}' shattered into {len(subtasks)} micro-tasks. Spawned specialized task force: {', '.join(spawned_roles)}."
            })

    def _monitor_swarm_health(self):
        """Act as the Defibrillator. Check for dead/stale agents and restart them."""
        agents = supervisor.list_agents()
        for agent in agents:
            if not agent['alive']:
                log.warning(f"COO detected dead agent: {agent['name']}. Evaluating fallback.")
                if "Claude" in agent['name'] or "Gemini" in agent['name']:
                    fallback_name = agent['name'].replace("Claude", "Llama").replace("Gemini", "Llama")
                    log.info(f"COO spinning up fallback agent: {fallback_name}")
                    supervisor.start_agent(
                        name=fallback_name,
                        backend="ollama",
                        model="llama3",
                        channel="tasks" 
                    )

    def _purge_ghost_sessions(self):
        """Clean up dead MCP sessions to prevent the 'Ghost Town' issue."""
        now = time.time()
        if now - self._last_cleanup > 60: 
            _cleanup_stale_sessions_sync()
            self._last_cleanup = now

    def _fix_double_escaped_channels(self):
        if not CHANNELS_DIR.exists():
            return
        for filepath in CHANNELS_DIR.glob("##*.jsonl"):
            bad_name = filepath.name
            good_name = bad_name[1:] 
            good_path = CHANNELS_DIR / good_name
            try:
                content = filepath.read_text("utf-8")
                with open(good_path, "a", encoding="utf-8") as f:
                    f.write(content)
                filepath.unlink()
            except Exception:
                pass

    def _detect_and_kill_infinite_loops(self):
        """Monitor channel velocity to catch Llama agents stuck in [NO_RESPONSE] or claim loops."""
        if not CHANNELS_DIR.exists():
            return
        now = datetime.now(timezone.utc)
        for filepath in CHANNELS_DIR.glob("*.jsonl"):
            try:
                lines = filepath.read_text("utf-8").strip().splitlines()
                recent_lines = lines[-50:]
                agent_counts = {}
                for line in recent_lines:
                    data = json.loads(line)
                    msg_time = datetime.fromisoformat(data.get("timestamp", ""))
                    if (now - msg_time).total_seconds() < 60:
                        session_id = data.get("session_id", "unknown")
                        agent_counts[session_id] = agent_counts.get(session_id, 0) + 1
                for session_id, count in agent_counts.items():
                    if session_id == "human-ceo":
                        continue
                    if count > 15:
                        log.warning(f"COO DETECTED RUNAWAY AGENT: {session_id} posted {count} msgs in 60s. Terminating.")
                        session_file = SESSIONS_DIR / f"{session_id}.json"
                        if session_file.exists():
                            session_data = json.loads(session_file.read_text("utf-8"))
                            agent_name = session_data.get("name")
                            supervisor.stop_agent(agent_name)
                            pid = session_data.get("pid", 0)
                            if pid > 0:
                                try:
                                    if os.name == 'nt':
                                        subprocess.run(["taskkill", "/F", "/PID", str(pid)], capture_output=True)
                                    else:
                                        subprocess.run(["kill", "-9", str(pid)], capture_output=True)
                                except Exception:
                                    pass
                        _append_jsonl_sync(filepath, {
                            "id": os.urandom(6).hex(),
                            "timestamp": now.isoformat(),
                            "session_id": "COO-Node",
                            "session_name": "The-COO",
                            "platform": "system",
                            "channel": filepath.stem,
                            "msg_type": "alert",
                            "message": f"COO INTERVENTION: Agent {session_id} has been forcefully terminated for infinite looping."
                        })
            except Exception:
                pass

    def _tick(self) -> bool:
        had_activity = super()._tick()
        self._purge_ghost_sessions()
        self._fix_double_escaped_channels()
        self._detect_and_kill_infinite_loops()
        self._monitor_swarm_health()
        return had_activity

if __name__ == "__main__":
    coo = OrchestratorAgent(name="COO-Node", backend="gemini", model="gemini-3.1-pro-preview")
    coo.run()