"""
Orchestrator: The new agent architecture with guaranteed execution.

Key principles:
1. LLM declares intent, system executes (no hallucination possible)
2. Omniscient context - LLM always sees current system state
3. Character voice - everything sounds like the bot
4. Guaranteed delivery - completions always reach the user

The user can always chat while tasks run in the background.
"""

import asyncio
import json
from datetime import datetime
from pathlib import Path
from typing import Any

from loguru import logger

from kyber.bus.events import InboundMessage, OutboundMessage
from kyber.bus.queue import MessageBus
from kyber.providers.base import LLMProvider
from kyber.session.manager import SessionManager

from kyber.agent.task_registry import TaskRegistry, Task, TaskStatus
from kyber.agent.worker import WorkerPool
from kyber.agent.voice import CharacterVoice
from kyber.agent.intent import (
    Intent, IntentAction, AgentResponse, 
    RESPOND_TOOL, parse_tool_call
)


# Action-claiming phrases that require intent validation
ACTION_PHRASES = [
    "started", "kicked off", "working on", "in progress",
    "spawned", "running", "executing", "on it",
    "i'll", "i will", "let me", "going to",
]


class Orchestrator:
    """
    The main agent orchestrator.
    
    Handles:
    - Message processing with structured intent
    - Task spawning and tracking
    - Progress updates (batched every 30s)
    - Completion notifications (guaranteed, in character)
    - Concurrent chat while tasks run
    """
    
    def __init__(
        self,
        bus: MessageBus,
        provider: LLMProvider,
        workspace: Path,
        persona_prompt: str,
        model: str | None = None,
        brave_api_key: str | None = None,
        progress_interval: float = 30.0,
    ):
        self.bus = bus
        self.provider = provider
        self.workspace = workspace
        self.persona = persona_prompt
        self.model = model or provider.get_default_model()
        self.brave_api_key = brave_api_key
        self.progress_interval = progress_interval
        
        # Core components
        self.registry = TaskRegistry()
        self.sessions = SessionManager(workspace)
        self.voice = CharacterVoice(persona_prompt, provider)
        self.workers = WorkerPool(
            provider=provider,
            workspace=workspace,
            registry=self.registry,
            model=self.model,
            brave_api_key=brave_api_key,
        )
        
        self._running = False
        self._last_progress_update: dict[str, datetime] = {}
    
    async def run(self) -> None:
        """
        Main run loop. Processes messages and handles completions.
        """
        self._running = True
        logger.info("Orchestrator started")
        
        # Start background tasks
        completion_task = asyncio.create_task(self._completion_loop())
        progress_task = asyncio.create_task(self._progress_loop())
        
        try:
            while self._running:
                try:
                    msg = await asyncio.wait_for(
                        self.bus.consume_inbound(),
                        timeout=1.0,
                    )
                    # Handle each message in its own task (non-blocking)
                    asyncio.create_task(self._handle_message(msg))
                except asyncio.TimeoutError:
                    continue
        finally:
            completion_task.cancel()
            progress_task.cancel()
    
    def stop(self) -> None:
        """Stop the orchestrator."""
        self._running = False
        logger.info("Orchestrator stopping")
    
    async def _handle_message(self, msg: InboundMessage) -> None:
        """Handle a single inbound message."""
        try:
            response = await self._process_message(msg)
            if response:
                await self.bus.publish_outbound(response)
        except Exception as e:
            logger.error(f"Error handling message: {e}")
            await self.bus.publish_outbound(OutboundMessage(
                channel=msg.channel,
                chat_id=msg.chat_id,
                content=f"Sorry, something went wrong: {str(e)}",
            ))
    
    async def _process_message(self, msg: InboundMessage) -> OutboundMessage | None:
        """
        Process a message using structured intent.
        
        Flow:
        1. Inject current system state into context
        2. LLM responds with message + intent
        3. Validate honesty (claims match intent)
        4. Execute intent (spawn task, check status, etc.)
        5. Return response with any injected references
        """
        logger.info(f"Processing message from {msg.channel}:{msg.sender_id}")
        
        session = self.sessions.get_or_create(msg.session_key)
        
        # Build context with omniscient state
        system_state = self.registry.get_context_summary()
        messages = self._build_messages(session, msg.content, system_state)
        
        # Get structured response from LLM
        agent_response = await self._get_structured_response(messages)
        
        if not agent_response:
            return OutboundMessage(
                channel=msg.channel,
                chat_id=msg.chat_id,
                content="Sorry, I'm having trouble responding right now.",
            )
        
        # Validate honesty
        agent_response = await self._validate_honesty(agent_response, messages)
        
        # Execute intent and get final message
        final_message = await self._execute_intent(
            agent_response,
            msg.channel,
            msg.chat_id,
        )
        
        # Save to session
        session.add_message("user", msg.content)
        session.add_message("assistant", final_message)
        self.sessions.save(session)
        
        return OutboundMessage(
            channel=msg.channel,
            chat_id=msg.chat_id,
            content=final_message,
        )
    
    def _build_messages(
        self,
        session: Any,
        current_message: str,
        system_state: str,
    ) -> list[dict[str, Any]]:
        """Build messages with omniscient context."""
        system_prompt = f"""{self.persona}

## Current System State
{system_state}

## How to Respond
You MUST use the 'respond' tool for every response. This tool lets you:
- Send a natural message to the user
- Optionally request an action (spawn_task, check_status, etc.)

When the user asks you to DO something (create, build, fix, write, etc.):
- Set intent.action to "spawn_task"
- Provide a detailed task_description
- Provide a short task_label

When the user asks about status or progress:
- Set intent.action to "check_status"
- Optionally provide task_ref if they mention a specific reference

For pure conversation:
- Set intent.action to "none"

IMPORTANT: Your message should be natural and in-character. The system will
handle adding references - you don't need to make them up."""

        messages = [{"role": "system", "content": system_prompt}]
        
        # Add history
        for h in session.get_history(max_messages=10):
            messages.append(h)
        
        # Add current message
        messages.append({"role": "user", "content": current_message})
        
        return messages
    
    async def _get_structured_response(
        self,
        messages: list[dict[str, Any]],
    ) -> AgentResponse | None:
        """Get structured response from LLM using function calling."""
        try:
            response = await self.provider.chat(
                messages=messages,
                tools=[RESPOND_TOOL],
                model=self.model,
            )
            
            if response.has_tool_calls:
                for tc in response.tool_calls:
                    if tc.name == "respond":
                        return parse_tool_call(tc)
            
            # Fallback: LLM didn't use the tool, treat as pure chat
            if response.content:
                return AgentResponse(
                    message=response.content,
                    intent=Intent(action=IntentAction.NONE),
                )
            
            return None
            
        except Exception as e:
            logger.error(f"Error getting structured response: {e}")
            return None
    
    async def _validate_honesty(
        self,
        response: AgentResponse,
        messages: list[dict[str, Any]],
    ) -> AgentResponse:
        """
        Validate that the LLM's claims match its declared intent.
        
        If the message claims an action but intent doesn't declare it,
        force a retry.
        """
        message_lower = response.message.lower()
        claims_action = any(phrase in message_lower for phrase in ACTION_PHRASES)
        declared_action = response.intent.action == IntentAction.SPAWN_TASK
        
        if claims_action and not declared_action:
            logger.warning(
                f"Honesty violation: claims action but intent is {response.intent.action}"
            )
            
            # Retry with correction
            messages.append({
                "role": "assistant",
                "content": response.message,
            })
            messages.append({
                "role": "user",
                "content": (
                    "You said you would do something, but you didn't set "
                    "intent.action to 'spawn_task'. If you're going to do work, "
                    "you MUST declare it in the intent. Please respond again "
                    "with the correct intent."
                ),
            })
            
            retry = await self._get_structured_response(messages)
            if retry:
                return retry
        
        return response
    
    async def _execute_intent(
        self,
        response: AgentResponse,
        channel: str,
        chat_id: str,
    ) -> str:
        """
        Execute the declared intent and return the final message.
        
        This is where the system takes over - the LLM declared what it wants,
        now we actually do it and inject proof (references).
        """
        message = response.message
        intent = response.intent
        
        if intent.action == IntentAction.SPAWN_TASK:
            # Create and spawn task
            task = self.registry.create(
                description=intent.task_description or message,
                label=intent.task_label or "Task",
                origin_channel=channel,
                origin_chat_id=chat_id,
            )
            self.workers.spawn(task)
            
            # Inject reference into message
            if task.reference not in message:
                message = f"{message}\n\nRef: {task.reference}"
            
            logger.info(f"Spawned task: {task.label} [{task.reference}]")
        
        elif intent.action == IntentAction.CHECK_STATUS:
            # Get status and inject into message
            status = self.registry.get_status_for_ref(intent.task_ref)
            status_voiced = await self.voice.speak_status(status)
            message = status_voiced
        
        elif intent.action == IntentAction.CANCEL_TASK:
            # TODO: Implement task cancellation
            message = f"{message}\n\n(Task cancellation not yet implemented)"
        
        # IntentAction.NONE = pure chat, no modification needed
        
        return message
    
    async def _completion_loop(self) -> None:
        """
        Background loop that handles task completions.
        
        Guarantees: Every completed task gets a notification sent to the user.
        """
        while self._running:
            try:
                # Wait for a completion
                task = await asyncio.wait_for(
                    self.workers.get_completion(),
                    timeout=1.0,
                )
                
                # Generate in-character notification
                if task.status == TaskStatus.COMPLETED:
                    notification = await self.voice.speak_completion(
                        label=task.label,
                        result=task.result or "Done",
                        completion_reference=task.completion_reference or "✅",
                    )
                else:
                    notification = await self.voice.speak_failure(
                        label=task.label,
                        error=task.error or "Unknown error",
                        completion_reference=task.completion_reference or "❌",
                    )
                
                # Push directly to user - guaranteed delivery
                await self.bus.publish_outbound(OutboundMessage(
                    channel=task.origin_channel,
                    chat_id=task.origin_chat_id,
                    content=notification,
                ))
                
                logger.info(f"Sent completion notification for: {task.label}")
                
            except asyncio.TimeoutError:
                continue
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in completion loop: {e}")
    
    async def _progress_loop(self) -> None:
        """
        Background loop that sends batched progress updates.
        
        Every 30 seconds, if there are active tasks with meaningful progress,
        send an update to the user.
        """
        while self._running:
            try:
                await asyncio.sleep(self.progress_interval)
                
                active_tasks = self.registry.get_active_tasks()
                if not active_tasks:
                    continue
                
                # Group by origin (channel:chat_id)
                by_origin: dict[str, list[Task]] = {}
                for task in active_tasks:
                    key = f"{task.origin_channel}:{task.origin_chat_id}"
                    if key not in by_origin:
                        by_origin[key] = []
                    by_origin[key].append(task)
                
                # Send progress update to each origin
                for origin_key, tasks in by_origin.items():
                    # Check if we should send (avoid spamming)
                    last_update = self._last_progress_update.get(origin_key)
                    if last_update:
                        elapsed = (datetime.now() - last_update).total_seconds()
                        if elapsed < self.progress_interval * 0.9:
                            continue
                    
                    # Build summaries
                    summaries = []
                    for task in tasks:
                        if task.status == TaskStatus.RUNNING:
                            if task.current_action:
                                action = task.current_action
                            elif task.actions_completed:
                                action = f"step {task.iteration}/{task.max_iterations}"
                            else:
                                action = f"getting started (step {task.iteration}/{task.max_iterations})"
                            summaries.append(f"{task.label}: {action}")
                    
                    if not summaries:
                        continue
                    
                    # Generate in-character update
                    update = await self.voice.speak_progress(summaries)
                    
                    # Parse origin
                    parts = origin_key.split(":", 1)
                    channel = parts[0]
                    chat_id = parts[1] if len(parts) > 1 else "direct"
                    
                    await self.bus.publish_outbound(OutboundMessage(
                        channel=channel,
                        chat_id=chat_id,
                        content=update,
                    ))
                    
                    self._last_progress_update[origin_key] = datetime.now()
                    logger.debug(f"Sent progress update to {origin_key}")
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in progress loop: {e}")
    
    async def process_direct(self, content: str, session_key: str = "cli:direct") -> str:
        """Process a message directly (for CLI usage)."""
        msg = InboundMessage(
            channel="cli",
            sender_id="user",
            chat_id="direct",
            content=content,
        )
        
        response = await self._process_message(msg)
        return response.content if response else ""
