"""
Worker: Executes tasks in the background with guaranteed completion.

Workers run async, emit progress events, and always complete (success or failure).
The completion queue guarantees the user gets notified.
"""

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

from loguru import logger

from kyber.agent.task_registry import Task, TaskRegistry, TaskStatus
from kyber.agent.tools.registry import ToolRegistry
from kyber.agent.tools.filesystem import ReadFileTool, WriteFileTool, EditFileTool, ListDirTool
from kyber.agent.tools.shell import ExecTool
from kyber.agent.tools.web import WebSearchTool, WebFetchTool
from kyber.providers.base import LLMProvider


class Worker:
    """
    Executes a single task with tools.
    
    Guarantees:
    - Always completes (success, failure, or timeout)
    - Emits progress updates
    - Pushes to completion queue when done
    """
    
    def __init__(
        self,
        task: Task,
        provider: LLMProvider,
        workspace: Path,
        registry: TaskRegistry,
        completion_queue: asyncio.Queue,
        model: str | None = None,
        brave_api_key: str | None = None,
        exec_timeout: int = 60,
        task_timeout: int = 300,  # 5 minute max per task
    ):
        self.task = task
        self.provider = provider
        self.workspace = workspace
        self.registry = registry
        self.completion_queue = completion_queue
        self.model = model or provider.get_default_model()
        self.brave_api_key = brave_api_key
        self.exec_timeout = exec_timeout
        self.task_timeout = task_timeout
        
        self.tools = self._build_tools()
    
    def _build_tools(self) -> ToolRegistry:
        """Build the tool registry for this worker."""
        tools = ToolRegistry()
        tools.register(ReadFileTool())
        tools.register(WriteFileTool())
        tools.register(EditFileTool())
        tools.register(ListDirTool())
        tools.register(ExecTool(
            working_dir=str(self.workspace),
            timeout=self.exec_timeout,
        ))
        tools.register(WebSearchTool(api_key=self.brave_api_key))
        tools.register(WebFetchTool())
        return tools
    
    def _build_system_prompt(self) -> str:
        """Build focused system prompt for the worker."""
        return f"""You are a task executor. Your ONLY job is to complete the assigned task using tools.

## Your Task
{self.task.description}

## Rules
1. Use tools to do the actual work - don't just describe what you would do
2. Be thorough but efficient
3. When done, provide a clear summary of what you accomplished
4. Include concrete evidence: file paths created, command outputs, etc.

## Available Tools
- read_file: Read file contents
- write_file: Create or overwrite files
- edit_file: Edit existing files
- list_dir: List directory contents
- exec: Run shell commands
- web_search: Search the web
- web_fetch: Fetch web page content

## Workspace
{self.workspace}

Complete the task, then summarize what you did."""
    
    async def run(self) -> None:
        """
        Execute the task. Guaranteed to complete or fail.
        Always pushes to completion queue.
        """
        try:
            self.registry.mark_started(self.task.id)
            logger.info(f"Worker started: {self.task.label} [{self.task.id}]")
            
            # Run with timeout
            result = await asyncio.wait_for(
                self._execute(),
                timeout=self.task_timeout,
            )
            
            self.registry.mark_completed(self.task.id, result)
            logger.info(f"Worker completed: {self.task.label} [{self.task.id}]")
            
        except asyncio.TimeoutError:
            error = f"Task timed out after {self.task_timeout}s"
            self.registry.mark_failed(self.task.id, error)
            logger.error(f"Worker timeout: {self.task.label} [{self.task.id}]")
            
        except Exception as e:
            error = str(e)
            self.registry.mark_failed(self.task.id, error)
            logger.error(f"Worker failed: {self.task.label} [{self.task.id}] - {e}")
        
        finally:
            # ALWAYS push to completion queue - guaranteed delivery
            await self.completion_queue.put(self.task)
    
    async def _execute(self) -> str:
        """Execute the task loop."""
        messages = [
            {"role": "system", "content": self._build_system_prompt()},
            {"role": "user", "content": self.task.description},
        ]
        
        max_iterations = self.task.max_iterations
        final_result: str | None = None
        
        for iteration in range(1, max_iterations + 1):
            self.registry.update_progress(self.task.id, iteration=iteration)
            
            response = await self.provider.chat(
                messages=messages,
                tools=self.tools.get_definitions(),
                model=self.model,
            )
            
            if response.has_tool_calls:
                # Add assistant message with tool calls
                tool_call_dicts = [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.name,
                            "arguments": json.dumps(tc.arguments),
                        },
                    }
                    for tc in response.tool_calls
                ]
                messages.append({
                    "role": "assistant",
                    "content": response.content or "",
                    "tool_calls": tool_call_dicts,
                })
                
                # Execute each tool
                for tool_call in response.tool_calls:
                    self.registry.update_progress(
                        self.task.id,
                        current_action=tool_call.name,
                    )
                    
                    logger.debug(f"Worker [{self.task.id}] executing: {tool_call.name}")
                    result = await self.tools.execute(tool_call.name, tool_call.arguments)
                    
                    self.registry.update_progress(
                        self.task.id,
                        current_action="",
                        action_completed=tool_call.name,
                    )
                    
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": tool_call.name,
                        "content": result,
                    })
            else:
                # No tool calls = done
                final_result = response.content
                break
        
        return final_result or "Task completed but no summary was generated."


class WorkerPool:
    """
    Manages concurrent worker execution.
    
    Spawns workers as async tasks, tracks them, and ensures
    completion notifications are delivered.
    """
    
    def __init__(
        self,
        provider: LLMProvider,
        workspace: Path,
        registry: TaskRegistry,
        model: str | None = None,
        brave_api_key: str | None = None,
        max_concurrent: int = 5,
    ):
        self.provider = provider
        self.workspace = workspace
        self.registry = registry
        self.model = model
        self.brave_api_key = brave_api_key
        self.max_concurrent = max_concurrent
        
        self.completion_queue: asyncio.Queue[Task] = asyncio.Queue()
        self._running: dict[str, asyncio.Task] = {}
    
    def spawn(self, task: Task) -> None:
        """Spawn a worker for the task."""
        worker = Worker(
            task=task,
            provider=self.provider,
            workspace=self.workspace,
            registry=self.registry,
            completion_queue=self.completion_queue,
            model=self.model,
            brave_api_key=self.brave_api_key,
        )
        
        async_task = asyncio.create_task(worker.run())
        self._running[task.id] = async_task
        
        # Cleanup when done
        def _on_done(_: asyncio.Task) -> None:
            self._running.pop(task.id, None)
        
        async_task.add_done_callback(_on_done)
        logger.info(f"Spawned worker for task: {task.label} [{task.id}]")
    
    @property
    def active_count(self) -> int:
        """Number of currently running workers."""
        return len(self._running)
    
    async def get_completion(self) -> Task:
        """Wait for and return the next completed task."""
        return await self.completion_queue.get()
    
    def get_completion_nowait(self) -> Task | None:
        """Get a completed task if available, else None."""
        try:
            return self.completion_queue.get_nowait()
        except asyncio.QueueEmpty:
            return None
