"""
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,
        persona_prompt: str,
        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.persona_prompt = persona_prompt
        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 system prompt that embodies the bot's persona while executing tasks."""
        return f"""{self.persona_prompt}

    ## Your Current Task
    {self.task.description}

    ## How You Work
    You are completing a task in the background. Use tools to do the actual work — don't just describe what you would do. Be thorough but efficient.

    When you are DONE, deliver your findings directly to the user in your own voice. This is what they will read. Be specific, include evidence (file paths, command outputs, key findings), and stay in character.

    Do NOT end on a tool call. Your final message IS the deliverable.

    ## 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}"""
    
    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.

        Step budget is set by complexity. When approaching the limit, the LLM
        is warned to wrap up. If it still exhausts steps, a forced summary
        turn (no tools) guarantees a result.
        """
        messages = [
            {"role": "system", "content": self._build_system_prompt()},
            {"role": "user", "content": self.task.description},
        ]

        max_iterations = self.task.max_iterations
        # Warn the LLM when ~80% of budget is used so it can wrap up naturally
        warn_at = max(3, int(max_iterations * 0.8))
        warned = False

        final_result: str | None = None
        last_content: str | None = None
        tools_used: list[str] = []
        tool_outputs: list[str] = []

        for iteration in range(1, max_iterations + 1):
            self.registry.update_progress(self.task.id, iteration=iteration)

            # Nudge the LLM to start wrapping up
            if iteration >= warn_at and not warned:
                warned = True
                remaining = max_iterations - iteration
                messages.append({
                    "role": "user",
                    "content": (
                        f"You have {remaining} steps remaining. Start wrapping up — "
                        "finish any critical tool calls, then deliver your final "
                        "answer to the user in your next response."
                    ),
                })

            response = await self.provider.chat(
                messages=messages,
                tools=self.tools.get_definitions(),
                model=self.model,
            )

            if response.content:
                last_content = response.content

            if response.has_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,
                })

                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)

                    tools_used.append(tool_call.name)
                    preview = result[:200] if result else "(empty)"
                    tool_outputs.append(f"{tool_call.name}: {preview}")

                    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 = LLM is done, this is the final answer
                final_result = response.content
                break

        # --- Budget exhausted: force a summary with tools disabled ---
        if final_result is None:
            messages.append({
                "role": "user",
                "content": (
                    "You've used all available steps. Deliver your findings "
                    "and answer NOW, in your own voice, in character. "
                    "Be specific — include what you found, what you checked, "
                    "and any evidence from the tools you used."
                ),
            })
            try:
                summary_response = await self.provider.chat(
                    messages=messages,
                    tools=None,
                    model=self.model,
                )
                if summary_response.content and summary_response.content.strip():
                    final_result = summary_response.content
            except Exception as e:
                logger.warning(f"Worker [{self.task.id}] summary turn failed: {e}")

        if final_result is None and last_content:
            final_result = last_content

        # Absolute fallback: build from tool execution log
        if final_result is None:
            unique_tools = list(dict.fromkeys(tools_used))
            tool_summary = ", ".join(unique_tools) if unique_tools else "none"
            recent_outputs = "\n".join(tool_outputs[-5:]) if tool_outputs else "No output captured."
            final_result = (
                f"Task ran for {max_iterations} steps. Tools used: {tool_summary}.\n\n"
                f"Last tool outputs:\n{recent_outputs}"
            )

        return final_result


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,
        persona_prompt: str,
        model: str | None = None,
        brave_api_key: str | None = None,
        max_concurrent: int = 5,
    ):
        self.provider = provider
        self.workspace = workspace
        self.registry = registry
        self.persona_prompt = persona_prompt
        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,
            persona_prompt=self.persona_prompt,
            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
