"""Codex Orchestrator — manages AI agents, schedules tasks, aggregates results."""
from __future__ import annotations

import asyncio
import logging
import uuid
from datetime import datetime, timezone
from typing import Dict, List, Optional

from mesh.codex.models import (
    CodexAgentInfo, CodexTask, CodexTaskResult, CodexTaskStatus,
    CodexTaskSubmission, CodexTaskType,
)

logger = logging.getLogger(__name__)

# Default schedule intervals (seconds)
ANALYSIS_INTERVAL_S = 86400    # 24 hours (nightly)
DIAGNOSTIC_INTERVAL_S = 3600   # 1 hour
TUNING_INTERVAL_S = 21600      # 6 hours


class CodexOrchestrator:
    """Manages Codex AI agents, schedules periodic tasks, and queues work.

    The orchestrator is the central coordinator for all Codex agents.
    It owns the task queue, schedules periodic runs, and provides
    query access to results.
    """

    def __init__(self, mesh_db, config=None, jepa_trainer=None, coordinator=None):
        self.mesh_db = mesh_db
        self.config = config
        self.jepa_trainer = jepa_trainer
        self.coordinator = coordinator
        self._agents: Dict[str, "CodexAgent"] = {}
        self._task_queue: asyncio.Queue = asyncio.Queue()
        self._running = False
        self._background_tasks: List[asyncio.Task] = []
        self._initialize_agents()

    def _initialize_agents(self):
        """Create all agent instances."""
        from mesh.codex.analyzer import CodeAnalysisAgent
        from mesh.codex.tuner import PerformanceTunerAgent
        from mesh.codex.diagnostics import DiagnosticAgent
        from mesh.codex.test_gen import TestGeneratorAgent

        self._agents["analyzer"] = CodeAnalysisAgent(
            mesh_db=self.mesh_db,
            config=self.config,
        )
        self._agents["tuner"] = PerformanceTunerAgent(
            mesh_db=self.mesh_db,
            config=self.config,
            jepa_trainer=self.jepa_trainer,
            coordinator=self.coordinator,
        )
        self._agents["diagnostics"] = DiagnosticAgent(
            mesh_db=self.mesh_db,
            config=self.config,
        )
        self._agents["test_gen"] = TestGeneratorAgent(
            mesh_db=self.mesh_db,
            config=self.config,
        )

    def get_agent(self, agent_type: str):
        return self._agents.get(agent_type)

    def get_all_agents(self) -> List[CodexAgentInfo]:
        return [agent.info for agent in self._agents.values()]

    async def start(self):
        """Start the orchestrator background loops."""
        if self._running:
            return
        self._running = True
        logger.info("Codex orchestrator starting with %d agents", len(self._agents))

        # Start task processing loop
        task = asyncio.create_task(self._process_loop())
        self._background_tasks.append(task)

        # Start periodic scheduled tasks
        task = asyncio.create_task(self._schedule_loop())
        self._background_tasks.append(task)

        logger.info("Codex orchestrator started")

    async def stop(self):
        """Stop all background loops."""
        self._running = False
        for task in self._background_tasks:
            task.cancel()
        self._background_tasks.clear()
        logger.info("Codex orchestrator stopped")

    async def submit_task(self, submission: CodexTaskSubmission) -> CodexTask:
        """Submit a new task to the queue."""
        task = CodexTask(
            task_id=str(uuid.uuid4()),
            task_type=submission.task_type,
            target=submission.target,
            parameters=submission.parameters,
            status=CodexTaskStatus.PENDING,
            created_at=datetime.now(timezone.utc),
        )

        # Record in DB
        if self.mesh_db:
            self.mesh_db.create_codex_task(
                task_id=task.task_id,
                task_type=task.task_type.value,
                target=task.target,
                parameters=task.parameters,
            )

        await self._task_queue.put(task)
        logger.info("Codex task submitted: %s (%s) target=%s",
                     task.task_id, task.task_type.value, task.target)
        return task

    async def run_agent(self, agent_type: str,
                        target: str = "",
                        parameters: Optional[Dict] = None) -> CodexTaskResult:
        """Trigger an agent manually and wait for the result."""
        agent = self._agents.get(agent_type)
        if not agent:
            return CodexTaskResult(
                task_id="",
                success=False,
                error=f"Unknown agent type: {agent_type}",
            )

        task = CodexTask(
            task_id=str(uuid.uuid4()),
            task_type=self._agent_type_to_task_type(agent_type),
            target=target,
            parameters=parameters or {},
            status=CodexTaskStatus.RUNNING,
            created_at=datetime.now(timezone.utc),
        )

        if self.mesh_db:
            self.mesh_db.create_codex_task(
                task_id=task.task_id,
                task_type=task.task_type.value,
                target=task.target,
                parameters=task.parameters,
            )

        result = await agent.safe_execute(task)
        return result

    def get_results(self, limit: int = 50,
                    task_type: Optional[str] = None) -> List[dict]:
        """Get recent task results from the database."""
        if not self.mesh_db:
            return []
        return self.mesh_db.get_codex_tasks(limit=limit, task_type=task_type)

    def get_findings(self, limit: int = 100,
                     severity: Optional[str] = None) -> List[dict]:
        """Get findings across all tasks."""
        if not self.mesh_db:
            return []
        return self.mesh_db.get_findings(limit=limit, severity=severity)

    async def _process_loop(self):
        """Process queued tasks."""
        while self._running:
            try:
                task = await asyncio.wait_for(self._task_queue.get(), timeout=5.0)
            except asyncio.TimeoutError:
                continue
            except asyncio.CancelledError:
                break

            agent = self._get_agent_for_task(task)
            if agent is None:
                logger.warning("No agent found for task type %s", task.task_type.value)
                if self.mesh_db:
                    self.mesh_db.update_codex_task(
                        task.task_id,
                        status=CodexTaskStatus.FAILED.value,
                        result_json={"error": f"No agent for task type {task.task_type.value}"},
                    )
                continue

            logger.info("Processing task %s with agent %s", task.task_id, agent.agent_type)
            if self.mesh_db:
                self.mesh_db.update_codex_task(task.task_id, status=CodexTaskStatus.RUNNING.value)

            await agent.safe_execute(task)

    async def _schedule_loop(self):
        """Run periodic scheduled tasks."""
        # Wait a bit before starting scheduled tasks
        await asyncio.sleep(30)

        analysis_counter = 0
        diagnostic_counter = 0
        tuning_counter = 0
        tick_s = 60  # Check every minute

        while self._running:
            try:
                await asyncio.sleep(tick_s)
            except asyncio.CancelledError:
                break

            analysis_counter += tick_s
            diagnostic_counter += tick_s
            tuning_counter += tick_s

            # Hourly diagnostics
            if diagnostic_counter >= DIAGNOSTIC_INTERVAL_S:
                diagnostic_counter = 0
                try:
                    await self.run_agent("diagnostics")
                    logger.info("Scheduled diagnostics completed")
                except Exception as e:
                    logger.warning("Scheduled diagnostics failed: %s", e)

            # 6-hour tuning check
            if tuning_counter >= TUNING_INTERVAL_S:
                tuning_counter = 0
                try:
                    await self.run_agent("tuner")
                    logger.info("Scheduled tuning analysis completed")
                except Exception as e:
                    logger.warning("Scheduled tuning analysis failed: %s", e)

            # Nightly code analysis
            if analysis_counter >= ANALYSIS_INTERVAL_S:
                analysis_counter = 0
                try:
                    from pathlib import Path
                    project_root = Path(__file__).resolve().parents[2]
                    await self.run_agent("analyzer", target=str(project_root / "mesh"))
                    logger.info("Scheduled code analysis completed")
                except Exception as e:
                    logger.warning("Scheduled code analysis failed: %s", e)

    def _get_agent_for_task(self, task: CodexTask):
        mapping = {
            CodexTaskType.ANALYZE: "analyzer",
            CodexTaskType.TUNE: "tuner",
            CodexTaskType.DIAGNOSE: "diagnostics",
            CodexTaskType.TEST_GEN: "test_gen",
        }
        agent_type = mapping.get(task.task_type)
        if agent_type:
            return self._agents.get(agent_type)
        # CUSTOM tasks: check parameters for agent_type
        custom_agent = task.parameters.get("agent_type")
        if custom_agent:
            return self._agents.get(custom_agent)
        return None

    def _agent_type_to_task_type(self, agent_type: str) -> CodexTaskType:
        mapping = {
            "analyzer": CodexTaskType.ANALYZE,
            "tuner": CodexTaskType.TUNE,
            "diagnostics": CodexTaskType.DIAGNOSE,
            "test_gen": CodexTaskType.TEST_GEN,
        }
        return mapping.get(agent_type, CodexTaskType.CUSTOM)
