"""
Enhanced Task Orchestrator - Multi-Agent Integration

Extends the existing TaskOrchestrator with Multi-Agent support:
- Session management with checkpoints
- Message bus integration for inter-agent communication
- Quality judge integration via BQAS
- Heartbeat-based liveness

Helper methods (routing, quality checks, recovery) are in
enhanced_task_helpers.py via EnhancedTaskHelpersMixin.
"""

import structlog
from typing import Optional, Dict, Any

from services.task_orchestrator import TaskOrchestrator
from services.enhanced_task_helpers import EnhancedTaskHelpersMixin
from models.task import Task, TaskState

# Import agent-core components
import sys
sys.path.insert(0, '/Users/benjaminadmin/Projekte/breakpilot-pwa/agent-core')

from sessions.session_manager import SessionManager, AgentSession
from sessions.heartbeat import HeartbeatMonitor, HeartbeatClient
from brain.memory_store import MemoryStore
from brain.context_manager import ContextManager
from orchestrator.message_bus import MessageBus
from orchestrator.task_router import TaskRouter

logger = structlog.get_logger(__name__)


class EnhancedTaskOrchestrator(EnhancedTaskHelpersMixin, TaskOrchestrator):
    """
    Enhanced TaskOrchestrator with Multi-Agent support.

    Extends the existing TaskOrchestrator to integrate with:
    - Session management for persistence and recovery
    - Message bus for inter-agent communication
    - Quality judge for response validation
    - Memory store for long-term learning

    Internal helper methods are provided by EnhancedTaskHelpersMixin.
    """

    def __init__(
        self,
        redis_client=None,
        db_pool=None,
        namespace: str = "breakpilot"
    ):
        """
        Initialize the enhanced orchestrator.

        Args:
            redis_client: Async Redis/Valkey client
            db_pool: Async PostgreSQL connection pool
            namespace: Namespace for isolation
        """
        super().__init__()

        # Initialize agent-core components
        self.session_manager = SessionManager(
            redis_client=redis_client,
            db_pool=db_pool,
            namespace=namespace
        )

        self.memory_store = MemoryStore(
            redis_client=redis_client,
            db_pool=db_pool,
            namespace=namespace
        )

        self.context_manager = ContextManager(
            redis_client=redis_client,
            db_pool=db_pool,
            namespace=namespace
        )

        self.message_bus = MessageBus(
            redis_client=redis_client,
            db_pool=db_pool,
            namespace=namespace
        )

        self.heartbeat = HeartbeatMonitor(
            timeout_seconds=30,
            check_interval_seconds=5,
            max_missed_beats=3
        )

        self.task_router = TaskRouter()

        # Track active sessions by voice session ID
        self._voice_sessions: Dict[str, AgentSession] = {}
        self._heartbeat_clients: Dict[str, HeartbeatClient] = {}

        logger.info("Enhanced TaskOrchestrator initialized with agent-core")

    async def start(self) -> None:
        """Starts the enhanced orchestrator"""
        await self.message_bus.start()
        await self.heartbeat.start_monitoring()

        # Subscribe to messages directed at this orchestrator
        await self.message_bus.subscribe(
            "voice-orchestrator",
            self._handle_agent_message
        )

        logger.info("Enhanced TaskOrchestrator started")

    async def stop(self) -> None:
        """Stops the enhanced orchestrator"""
        # Stop all heartbeat clients
        for client in self._heartbeat_clients.values():
            await client.stop()
        self._heartbeat_clients.clear()

        await self.heartbeat.stop_monitoring()
        await self.message_bus.stop()

        logger.info("Enhanced TaskOrchestrator stopped")

    async def create_session(
        self,
        voice_session_id: str,
        user_id: str = "",
        metadata: Optional[Dict[str, Any]] = None
    ) -> AgentSession:
        """
        Creates a new agent session for a voice session.

        Args:
            voice_session_id: The voice session ID
            user_id: Optional user ID
            metadata: Additional metadata

        Returns:
            The created AgentSession
        """
        # Create session via session manager
        session = await self.session_manager.create_session(
            agent_type="voice-orchestrator",
            user_id=user_id,
            context={"voice_session_id": voice_session_id},
            metadata=metadata
        )

        # Create conversation context
        self.context_manager.create_context(
            session_id=session.session_id,
            system_prompt=self._get_system_prompt(),
            max_messages=50
        )

        # Start heartbeat for this session
        heartbeat_client = HeartbeatClient(
            session_id=session.session_id,
            monitor=self.heartbeat,
            interval_seconds=10
        )
        await heartbeat_client.start()

        # Register heartbeat for monitoring
        self.heartbeat.register(session.session_id, "voice-orchestrator")

        # Store references
        self._voice_sessions[voice_session_id] = session
        self._heartbeat_clients[session.session_id] = heartbeat_client

        logger.info(
            "Created agent session",
            session_id=session.session_id[:8],
            voice_session_id=voice_session_id
        )

        return session

    async def get_session(
        self,
        voice_session_id: str
    ) -> Optional[AgentSession]:
        """Gets the agent session for a voice session"""
        return self._voice_sessions.get(voice_session_id)

    async def end_session(self, voice_session_id: str) -> None:
        """
        Ends an agent session.

        Args:
            voice_session_id: The voice session ID
        """
        session = self._voice_sessions.get(voice_session_id)
        if not session:
            return

        # Stop heartbeat
        if session.session_id in self._heartbeat_clients:
            await self._heartbeat_clients[session.session_id].stop()
            del self._heartbeat_clients[session.session_id]

        # Unregister from heartbeat monitor
        self.heartbeat.unregister(session.session_id)

        # Mark session as completed
        session.complete()
        await self.session_manager.update_session(session)

        # Clean up
        del self._voice_sessions[voice_session_id]

        logger.info(
            "Ended agent session",
            session_id=session.session_id[:8],
            duration_seconds=session.get_duration().total_seconds()
        )

    async def queue_task(self, task: Task) -> None:
        """
        Queue a task with session checkpointing.

        Extends parent to add checkpoint for recovery.
        """
        # Get session for this task
        session = self._voice_sessions.get(task.session_id)

        if session:
            # Checkpoint before queueing
            session.checkpoint("task_queued", {
                "task_id": task.id,
                "task_type": task.type.value,
                "parameters": task.parameters
            })
            await self.session_manager.update_session(session)

        # Call parent implementation
        await super().queue_task(task)

    async def process_task(self, task: Task) -> None:
        """
        Process a task with enhanced routing and quality checks.

        Extends parent to:
        - Route complex tasks to specialized agents
        - Run quality checks via BQAS
        - Store results in memory for learning
        """
        session = self._voice_sessions.get(task.session_id)

        if session:
            session.checkpoint("task_processing", {
                "task_id": task.id
            })

        # Check if this task should be routed to a specialized agent
        if self._needs_specialized_agent(task):
            await self._route_to_agent(task, session)
        else:
            # Use parent implementation for simple tasks
            await super().process_task(task)

        # Run quality check on result
        if task.result_ref and self._needs_quality_check(task):
            await self._run_quality_check(task, session)

        # Store in memory for learning
        if task.state == TaskState.READY and task.result_ref:
            await self._store_task_result(task)

        if session:
            session.checkpoint("task_completed", {
                "task_id": task.id,
                "state": task.state.value
            })
            await self.session_manager.update_session(session)
