"""
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.skills import SkillsLoader
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.agent.planner import generate_plan, interpolate_args, interpolate_refs, ExecutionPlan
from kyber.agent.workspace_index import WorkspaceIndex
from kyber.agent.narrator import LiveNarrator
from kyber.meta_messages import looks_like_prompt_leak
from kyber.providers.base import LLMProvider
from kyber.utils.helpers import redact_secrets


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,
        timezone: str | None = None,
        workspace_index: WorkspaceIndex | None = None,
        narrator: LiveNarrator | None = None,
    ):
        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.timezone = timezone
        self.workspace_index = workspace_index
        self.narrator = narrator

        self.tools = self._build_tools()

    @staticmethod
    def _result_is_unusable(text: str) -> bool:
        """Check if a final result is unsuitable for user delivery."""
        import re
        raw = text or ""
        t = " ".join(raw.split()).strip()
        if not t:
            return True
        if looks_like_prompt_leak(t):
            return True
        lower = t.lower()
        if len(t.split()) < 6:
            return True
        if any(tok in lower for tok in ["tool call", "tool_calls", "system prompt", "developer message"]):
            return True
        if lower.startswith("error calling llm:") or "litellm." in lower or "badrequest" in lower.replace(" ", ""):
            return True
        # Bot should never ask for permission — it should just do the work.
        if any(phrase in lower for phrase in [
            "would you like me to",
            "shall i ",
            "do you want me to",
            "want me to proceed",
            "like me to proceed",
            "like me to apply",
            "like me to fix",
        ]):
            return True
        # Markdown headers in chat messages = report-style, not conversational.
        # Check the raw text (before whitespace normalization) to preserve newlines.
        if re.search(r'^#{1,3}\s', raw, re.MULTILINE) and raw.count('\n') > 3:
            return True
        # Unfilled template placeholders — brackets containing instruction-like
        # text. Matches both "[Identify the error]" and "[specific issue here]".
        if re.search(r'\[[a-zA-Z][\w\s]{4,}\]', t):
            return True
        # Excessively long responses usually mean the LLM dumped raw file
        # contents or tool output instead of summarizing.
        if len(t) > 2000:
            return True
        return False

    async def _rewrite_result(self, raw_result: str) -> str:
        """Rewrite an unusable result into a concise user-facing message."""
        recent_actions = "\n".join(f"- {a}" for a in (self.task.actions_completed[-12:] or []))
        if not recent_actions:
            recent_actions = "- (none yet)"

        # Build a brief summary of what actually happened from the raw result.
        # Truncate the raw result so the rewrite LLM sees the key info without
        # being tempted to copy-paste it all.
        evidence = raw_result[:1500] if raw_result else "(no output)"

        # Pre-detect failure signals in the evidence so the LLM can't ignore them.
        failure_signals = []
        evidence_lower = evidence.lower()
        if "exit code:" in evidence_lower and "exit code: 0" not in evidence_lower:
            failure_signals.append("COMMAND FAILED (non-zero exit code)")
        if "traceback" in evidence_lower:
            failure_signals.append("PYTHON TRACEBACK DETECTED")
        if "error:" in evidence_lower or "error :" in evidence_lower:
            failure_signals.append("ERROR IN OUTPUT")
        if "not found" in evidence_lower:
            failure_signals.append("FILE OR COMMAND NOT FOUND")

        failure_context = ""
        if failure_signals:
            failure_context = (
                "\n\nFAILURE DETECTED: " + ", ".join(failure_signals) + "\n"
                "The task FAILED. Your message MUST say it failed and why. "
                "Do NOT say it succeeded. Do NOT invent positive results."
            )

        system = (
            f"{self.persona_prompt}\n\n"
            "You are rewriting a task result into a concise status update for the user.\n"
            "Stay in character.\n"
            "RULES:\n"
            "- If the output contains errors, tracebacks, or non-zero exit codes, the task FAILED. Say so.\n"
            "- Do NOT fabricate success. Do NOT say 'executed successfully' unless the output proves it.\n"
            "- Do NOT include API keys, tokens, passwords, or credentials.\n"
            "- Do NOT paste raw file contents, source code, tracebacks, or command output.\n"
            "- Do NOT use markdown headers (##) or structured formatting.\n"
            "- Do NOT ask the user for permission ('Would you like me to...').\n"
            "- Do NOT use placeholder brackets like [error here].\n"
            "- State the specific error in plain language (e.g., 'The script crashed because X').\n"
            "- Keep it under 800 characters. Talk naturally."
        )
        user = (
            f"Task: {self.task.label}\n"
            f"Request: {self.task.description}\n\n"
            f"Actions taken:\n{recent_actions}\n\n"
            f"Raw output (truncated):\n{evidence}\n"
            f"{failure_context}\n\n"
            "Rewrite this into a brief, natural status update."
        )

        try:
            r = await self.provider.chat(
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                tools=None,
                model=self.model,
                max_tokens=600,
            )
            out = (r.content or "").strip()
            if out and not self._result_is_unusable(out):
                return out
        except Exception as e:
            logger.warning(f"Worker [{self.task.id}] result rewrite failed: {e}")

        return raw_result
    
    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."""
        from kyber.utils.helpers import current_datetime_str

        # Workers run without the orchestrator's "intent" wrapper, so we include
        # user/workspace bootstrap + skills here (but NOT the orchestrator's tool-policy rules).
        bootstrap_files = ["AGENTS.md", "USER.md", "TOOLS.md", "IDENTITY.md"]
        bootstrap_parts: list[str] = []
        for fn in bootstrap_files:
            p = self.workspace / fn
            if p.exists():
                try:
                    bootstrap_parts.append(f"## {fn}\n\n{p.read_text(encoding='utf-8')}")
                except Exception:
                    continue
        bootstrap = "\n\n---\n\n".join(bootstrap_parts).strip()

        skills = SkillsLoader(self.workspace)
        always = skills.get_always_skills()
        always_content = skills.load_skills_for_context(always) if always else ""
        skills_summary = skills.build_skills_summary()
        skills_block = ""
        if always_content or skills_summary:
            skills_block = (
                "## Skills\n\n"
                "Skills are markdown playbooks. If your task mentions a skill by name, read its SKILL.md and follow it.\n"
                f"- Workspace skills: {self.workspace}/skills/<skill>/SKILL.md\n"
                "- Managed skills: ~/.kyber/skills/<skill>/SKILL.md\n\n"
            )
            if always_content:
                skills_block += f"### Active Skills\n\n{always_content}\n\n"
            if skills_summary:
                skills_block += f"### Available Skills\n\n{skills_summary}\n"

        return f"""{self.persona_prompt}

    ## Current Time
    {current_datetime_str(self.timezone)}

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

    ## Workspace Context
    {bootstrap if bootstrap else "(no workspace bootstrap files found)"}

    {skills_block if skills_block else ""}

    {self._get_workspace_index_block()}

    ## How You Work
    You are completing a task in the background. Your job is to DO the work, not describe it.

    ACTION-ORIENTED RULES:
    - If the user asks you to run something, RUN IT. Don't just read the file and report.
    - If a command fails, DIAGNOSE the error and FIX IT, then re-run. Don't just report the error.
    - If code has a bug, use edit_file to fix it, then exec to verify the fix works.
    - NEVER ask the user for permission to proceed. You already have it — that's why you exist.
    - NEVER say "Would you like me to..." or "Shall I..." — just do it.
    - Complete the ENTIRE task. Reading files and reporting findings is NOT completing the task.

    FINAL MESSAGE RULES:
    - Your final message is a brief status update (under 1200 characters), NOT a report.
    - Say what you DID, not what you FOUND. "Fixed the import error and ran the script — 15 replies posted." not a wall of diagnostics.
    - Do NOT paste file contents, source code, tracebacks, or raw command output into your message. Summarize in your own words.
    - Do NOT use markdown headers (##) or structured report formatting. Just talk naturally.
    - If you truly cannot fix something (missing credentials, external service down), explain the blocker concisely.

    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 (PROTECTED: cannot write to cron/jobs.json or .kyber/ config files)
    - edit_file: Edit existing files (PROTECTED: same restrictions as write_file)
    - list_dir: List directory contents
    - exec: Run shell commands
    - web_search: Search the web
    - web_fetch: Fetch web page content

    ## Critical Rules
    - NEVER write directly to cron/jobs.json or any file under ~/.kyber/cron/. Use `kyber cron add` via exec instead.
    - NEVER overwrite existing config files. Read them first, then append or edit only what's needed.
    - When adding cron jobs, ALWAYS use `kyber cron add` — never write the JSON file directly.
    - NEVER fabricate or invent results. If a command fails, report the actual error. If a file doesn't exist, say so. Do NOT claim success when the output shows failure.
    - NEVER include API keys, tokens, passwords, or credentials in messages to the user. If you encounter them in files or output, do NOT repeat them.
    - If a tool returns an error (non-zero exit code, "Error:", traceback), your final message MUST acknowledge the failure and include the actual error. Do NOT describe a successful outcome.

    ## Workspace
    {self.workspace}"""

    def _get_workspace_index_block(self) -> str:
        """Get workspace index for injection into system prompt."""
        if not self.workspace_index:
            return ""
        try:
            index = self.workspace_index.get()
            if index:
                return (
                    "## Workspace File Map\n"
                    "You already know the workspace layout. Skip list_dir unless you need "
                    "a subdirectory not shown here. Jump straight to the work.\n\n"
                    f"{index}"
                )
        except Exception as e:
            logger.debug(f"Failed to get workspace index: {e}")
        return ""
    
    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}]")
            
            result = await self._execute()
            
            self.registry.mark_completed(self.task.id, result)
            logger.info(f"Worker completed: {self.task.label} [{self.task.id}]")
            
        except asyncio.CancelledError:
            self.registry.mark_cancelled(self.task.id, "cancelled")
            logger.info(f"Worker cancelled: {self.task.label} [{self.task.id}]")
            raise
        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.

        Strategy:
        1. Try PLANNED execution first (1 LLM call → deterministic steps)
        2. Fall back to ADAPTIVE loop if plan fails or can't be generated

        The planned path is dramatically faster for predictable tasks because
        it replaces N sequential LLM round-trips with a single planning call
        followed by deterministic tool execution.
        """
        # Try planned execution for tasks that aren't too complex
        plan_context: str | None = None
        if self.task.max_iterations is None or (self.task.max_iterations and self.task.max_iterations <= 20):
            planned_result, plan_context = await self._execute_planned()
            if planned_result is not None:
                result = redact_secrets(planned_result)
                if self._result_is_unusable(result):
                    logger.info(f"Worker [{self.task.id}] planned result unusable, rewriting")
                    result = await self._rewrite_result(result)
                return redact_secrets(result)
            logger.info(f"Worker [{self.task.id}] plan failed or unavailable, falling back to adaptive loop")

        result = await self._execute_adaptive(plan_context=plan_context)
        # The adaptive loop has its own quality gate, but apply the top-level
        # one as well in case something slipped through.
        result = redact_secrets(result)
        if self._result_is_unusable(result):
            logger.info(f"Worker [{self.task.id}] adaptive result unusable, rewriting")
            result = await self._rewrite_result(result)
        return redact_secrets(result)

    async def _execute_planned(self) -> tuple[str | None, str | None]:
        """
        Execute via the planner: one LLM call to generate a plan, then
        deterministic execution of each step.

        Returns a tuple of (result, plan_context):
        - result: the final result string, or None if planning fails
        - plan_context: a summary of steps already completed (for adaptive fallback)
        """
        ws_index = ""
        if self.workspace_index:
            try:
                ws_index = self.workspace_index.get()
            except Exception:
                pass

        plan = await generate_plan(
            provider=self.provider,
            model=self.model,
            system_prompt=self._build_system_prompt(),
            task_description=self.task.description,
            workspace_index=ws_index,
        )

        if not plan or not plan.steps:
            return None, None

        self.registry.update_progress(
            self.task.id,
            current_action=f"executing plan ({plan.step_count} steps)",
        )

        # Execute each step deterministically
        step_results: dict[int, str] = {}
        failed = False

        for i, step in enumerate(plan.steps, 1):
            tool_name = step.get("tool", "")
            raw_args = step.get("args", {})
            purpose = step.get("purpose", "working...")

            # Interpolate references to previous step results
            args = interpolate_args(raw_args, step_results)

            # Narrate to user in real-time
            if self.narrator:
                self.narrator.narrate(self.task.id, purpose)

            self.registry.update_progress(
                self.task.id,
                iteration=i,
                current_action=purpose,
            )

            logger.debug(f"Worker [{self.task.id}] plan step {i}/{plan.step_count}: {tool_name}")

            if not self.tools.has(tool_name):
                logger.warning(f"Worker [{self.task.id}] plan references unknown tool: {tool_name}")
                failed = True
                break

            try:
                result = await self.tools.execute(tool_name, args)
            except Exception as e:
                logger.warning(f"Worker [{self.task.id}] plan step {i} failed: {e}")
                result = f"Error: {e}"
                failed = True

            # Redact secrets before storing in step_results (used in summary_template)
            result = redact_secrets(result)
            step_results[i] = result

            self.registry.update_progress(
                self.task.id,
                current_action="",
                action_completed=purpose,
            )

            # If a step produced an error result, flag it and fall back.
            # Check for: tool-level errors ("Error:"), command failures
            # (non-zero exit codes), and Python tracebacks.
            step_failed = False
            if tool_name == "exec":
                # Non-zero exit code or traceback = command failed
                if "Exit code:" in result and "Exit code: 0" not in result:
                    step_failed = True
                elif "Traceback" in result:
                    step_failed = True
                elif result.startswith("Error:"):
                    step_failed = True
            elif tool_name in ("edit_file", "write_file"):
                # Write/edit failures are always fatal
                if result.startswith("Error:"):
                    step_failed = True
            # Note: read_file "Error: File not found" is NOT fatal — the plan
            # might be checking if a file exists. But if it's the ONLY thing
            # the plan does, the summary_template will be garbage anyway and
            # the quality gate in _execute will catch it.
            if step_failed:
                failed = True
                break

        if failed:
            # Plan hit a snag — fall back to adaptive loop.
            # Build context about what was already accomplished so the adaptive
            # loop doesn't start from scratch.
            completed_steps = []
            for j, step in enumerate(plan.steps[:len(step_results)], 1):
                purpose = step.get("purpose", "")
                tool = step.get("tool", "")
                result_preview = step_results.get(j, "")[:300]
                completed_steps.append(f"Step {j} ({tool}): {purpose}\nResult: {result_preview}")
            
            plan_context = None
            if completed_steps:
                plan_context = (
                    "A planned execution was attempted but failed partway through. "
                    "Here is what was already completed — do NOT redo these steps:\n\n"
                    + "\n\n".join(completed_steps)
                    + f"\n\nThe plan failed at step {len(step_results) + 1}. "
                    "Continue from where it left off."
                )
            return None, plan_context

        # Build final message from the plan's summary template.
        # Truncate step results before interpolation so the summary doesn't
        # become a wall of raw file contents or command output.
        truncated_results = {
            k: (v[:500] + "..." if len(v) > 500 else v)
            for k, v in step_results.items()
        }
        final = interpolate_refs(plan.summary_template, truncated_results)

        # If the template is too generic, enhance it
        if len(final.split()) < 6:
            # Build a richer summary from what we actually did
            actions = [s.get("purpose", "") for s in plan.steps[:len(step_results)]]
            final = f"{final}\n\n" + "\n".join(f"• {a}" for a in actions if a)

        return redact_secrets(final), None

    async def _execute_adaptive(self, plan_context: str | None = None) -> 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},
        ]

        # If a planned execution partially completed, inject that context
        # so the adaptive loop knows what was already done.
        if plan_context:
            messages.append({"role": "assistant", "content": plan_context})
            messages.append({
                "role": "user",
                "content": (
                    "The automated plan failed partway through. Pick up where it "
                    "left off — do NOT redo steps that already succeeded. Fix the "
                    "issue that caused the failure and complete the task."
                ),
            })

        max_iterations = self.task.max_iterations  # None = unlimited
        warned = False
        warn_at = 0
        if 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))

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

        def _one_line(s: str, limit: int = 140) -> str:
            s = " ".join((s or "").split()).strip()
            if len(s) > limit:
                s = s[: limit - 1].rstrip() + "…"
            return s

        def _redact_secrets(s: str) -> str:
            """Redact strings that look like API keys, tokens, or passwords."""
            import re
            # Common API key patterns: sk-..., key-..., Bearer ..., long hex/base64 strings
            s = re.sub(
                r'(?i)(api[_-]?key|token|secret|password|bearer)\s*[=:]\s*\S+',
                r'\1=***',
                s,
            )
            # Standalone key-like strings (sk-..., xai-..., etc.)
            s = re.sub(r'\b(sk|key|xai|gsk|pk|rk)-[A-Za-z0-9_-]{20,}\b', '***', s)
            return s

        def _format_action(tool_name: str, args: dict[str, Any]) -> str:
            """Human-readable action detail for progress/status; avoid internal tool jargon."""
            name = (tool_name or "").strip()
            args = args or {}
            if name == "exec":
                cmd = _redact_secrets(_one_line(str(args.get("command", ""))))
                return f"running `{cmd}`" if cmd else "running a shell command"
            if name == "read_file":
                p = _one_line(str(args.get("path", "")))
                return f"reading `{p}`" if p else "reading a file"
            if name == "list_dir":
                p = _one_line(str(args.get("path", "")))
                return f"checking `{p}`" if p else "checking a folder"
            if name == "write_file":
                p = _one_line(str(args.get("path", "")))
                return f"writing `{p}`" if p else "writing a file"
            if name == "edit_file":
                p = _one_line(str(args.get("path", "")))
                return f"editing `{p}`" if p else "editing a file"
            if name == "web_search":
                q = _one_line(str(args.get("query", "")))
                return f"searching for “{q}”" if q else "searching the web"
            if name == "web_fetch":
                u = _one_line(str(args.get("url", "")))
                return f"opening {u}" if u else "opening a web page"
            if name == "message":
                return "sending a message"
            return "making progress"

        iteration = 0
        consecutive_errors = 0
        max_consecutive_errors = 3
        while True:
            iteration += 1
            self.registry.update_progress(self.task.id, iteration=iteration)

            # Nudge the LLM to start wrapping up (only if a cap exists)
            if max_iterations and 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 write a brief status "
                        "update (what you did, what worked, what didn't). No reports, "
                        "no file dumps, no asking permission. Just results."
                    ),
                })

            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,
                })

                # Execute tool calls in parallel when possible.
                # Tools that mutate state (write_file, edit_file, exec) run
                # sequentially to avoid races; read-only tools run concurrently.
                _MUTATING_TOOLS = {"write_file", "edit_file", "exec", "message"}

                read_only = [tc for tc in response.tool_calls if tc.name not in _MUTATING_TOOLS]
                mutating = [tc for tc in response.tool_calls if tc.name in _MUTATING_TOOLS]

                async def _run_tool(tc):
                    action_detail = _format_action(tc.name, tc.arguments)
                    # Live-narrate to user
                    if self.narrator:
                        self.narrator.narrate(self.task.id, action_detail)
                    self.registry.update_progress(self.task.id, current_action=action_detail)
                    logger.debug(f"Worker [{self.task.id}] executing: {tc.name}")
                    result = await self.tools.execute(tc.name, tc.arguments)
                    # Redact secrets from tool output BEFORE it enters the LLM
                    # conversation. The LLM never needs to see real API keys.
                    result = redact_secrets(result)
                    tools_used.append(tc.name)
                    preview = result[:200] if result else "(empty)"
                    tool_outputs.append(f"{tc.name}: {preview}")
                    self.registry.update_progress(self.task.id, current_action="", action_completed=action_detail)
                    return tc, result

                # Run read-only tools concurrently
                tool_results: dict[str, str] = {}
                if read_only:
                    completed = await asyncio.gather(*[_run_tool(tc) for tc in read_only])
                    for tc, result in completed:
                        tool_results[tc.id] = result

                # Run mutating tools sequentially
                for tc in mutating:
                    _, result = await _run_tool(tc)
                    tool_results[tc.id] = result

                # Append results in original order (LLM expects this)
                for tool_call in response.tool_calls:
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": tool_call.name,
                        "content": tool_results[tool_call.id],
                    })
            else:
                # No tool calls — either the LLM is done or it errored out.
                if response.finish_reason == "error":
                    consecutive_errors += 1
                    if consecutive_errors >= max_consecutive_errors:
                        logger.error(
                            f"Worker [{self.task.id}] hit {max_consecutive_errors} consecutive LLM errors, giving up"
                        )
                        final_result = last_content
                        break
                    # Transient LLM failure (e.g. provider 500). Don't treat
                    # the error string as the final answer — retry so the
                    # agent can finish its work (e.g. write a report file).
                    logger.warning(
                        f"Worker [{self.task.id}] LLM returned error ({consecutive_errors}/{max_consecutive_errors}), retrying: "
                        f"{(response.content or '')[:200]}"
                    )
                    await asyncio.sleep(3)
                    continue

                consecutive_errors = 0

                # If this is the first iteration and the LLM just talked without
                # using any tools, it hasn't done any work yet. Push it to
                # actually use tools instead of accepting a fabricated answer.
                if iteration == 1 and not tools_used:
                    logger.info(f"Worker [{self.task.id}] LLM skipped tools on first turn, pushing back")
                    messages.append({
                        "role": "assistant",
                        "content": response.content or "",
                    })
                    messages.append({
                        "role": "user",
                        "content": (
                            "You haven't used any tools yet. You MUST use tools to "
                            "complete this task — read files, run commands, make edits. "
                            "Do NOT just describe what you would do. Actually do it now."
                        ),
                    })
                    continue

                final_result = response.content
                break

            # If a cap exists and we've hit it, break and force summary below.
            if max_iterations and iteration >= max_iterations:
                break

            # Compact message history periodically to prevent unbounded context growth on long tasks.
            # Lower threshold = smaller context windows = faster LLM inference per turn.
            #
            # IMPORTANT: tool-role messages must not be "orphaned" (a tool message without its
            # preceding assistant tool_calls message can break tool-calling schemas). When we
            # keep a tail, ensure it does not start mid tool-response block.
            if len(messages) > 60:
                system = messages[0]

                def _safe_tail_start(msgs: list[dict[str, Any]], desired_tail: int) -> int:
                    start = max(1, len(msgs) - desired_tail)
                    # Never start on a tool message; include the preceding assistant tool_calls.
                    while start > 1 and (msgs[start].get("role") == "tool"):
                        start -= 1
                    return start

                start = _safe_tail_start(messages, desired_tail=30)
                tail = messages[start:]

                unique_tools = list(dict.fromkeys(tools_used))
                tool_summary = ", ".join(unique_tools[-20:]) if unique_tools else "none"
                recent_outputs = "\n".join(tool_outputs[-12:]) if tool_outputs else "No output captured."
                recent_actions = "\n".join(f"- {a}" for a in (self.task.actions_completed[-12:] or []))
                if not recent_actions:
                    recent_actions = "- (none yet)"

                summary = (
                    "Earlier steps summary (compressed to save context):\n"
                    f"- iterations so far: {iteration}\n"
                    f"- recent actions:\n{recent_actions}\n"
                    f"- tools used recently: {tool_summary}\n"
                    f"- recent tool output previews:\n{recent_outputs}\n"
                )
                messages = [system, {"role": "assistant", "content": summary}] + tail

        # --- Budget exhausted: force a summary with tools disabled ---
        if final_result is None:
            messages.append({
                "role": "user",
                "content": (
                    "You've used all available steps. Write a brief status update: "
                    "what you did, what worked, what didn't. Under 800 characters. "
                    "No file contents, no reports, no asking permission."
                ),
            })
            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 {iteration} steps. Tools used: {tool_summary}.\n\n"
                f"Last tool outputs:\n{recent_outputs}"
            )

        def _final_is_unusable(text: str) -> bool:
            return self._result_is_unusable(text)

        # If the model leaked prompt/instructions (or produced a uselessly short
        # "receipt"), try a quick in-context recovery before the top-level gate.
        if _final_is_unusable(final_result):
            # Try one in-context finalization with the conversation history
            messages.append({
                "role": "user",
                "content": (
                    "Your previous response was not suitable for the user. "
                    "Write a brief status update (under 800 chars): what you did, "
                    "what worked, what didn't. No file contents, no source code, "
                    "no markdown headers, no asking permission. Just talk naturally."
                ),
            })
            try:
                recovery = await self.provider.chat(
                    messages=messages, tools=None, model=self.model, max_tokens=600,
                )
                if recovery.content and not _final_is_unusable(recovery.content):
                    final_result = recovery.content
            except Exception:
                pass

        # Final safety net: redact any secrets that slipped through.
        return redact_secrets(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,
        timezone: str | None = None,
        exec_timeout: int = 60,
        narrator: LiveNarrator | None = None,
    ):
        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.timezone = timezone
        self.exec_timeout = exec_timeout
        self.narrator = narrator

        # Shared workspace index — built once, reused across all workers
        self.workspace_index = WorkspaceIndex(workspace)

        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."""
        # Register with narrator for live updates
        if self.narrator:
            self.narrator.register_task(
                task.id, task.origin_channel, task.origin_chat_id,
            )

        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,
            exec_timeout=self.exec_timeout,
            timezone=self.timezone,
            workspace_index=self.workspace_index,
            narrator=self.narrator,
        )

        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)
            if self.narrator:
                self.narrator.unregister_task(task.id)

        async_task.add_done_callback(_on_done)
        logger.info(f"Spawned worker for task: {task.label} [{task.id}]")

    def cancel(self, task_id: str) -> bool:
        """Cancel a running task by task id."""
        t = self._running.get(task_id)
        if not t:
            return False
        t.cancel()
        return True

    async def run_inline(self, task: Task) -> Task:
        """
        Run a task to completion in the current event loop.

        This is used for "direct"/one-shot invocations where background tasks
        would be cancelled when the process exits (e.g., `kyber agent -m ...`).
        """
        # Use a throwaway queue; Worker.run() always puts the task on completion.
        q: asyncio.Queue[Task] = asyncio.Queue()
        worker = Worker(
            task=task,
            provider=self.provider,
            workspace=self.workspace,
            registry=self.registry,
            completion_queue=q,
            persona_prompt=self.persona_prompt,
            model=self.model,
            brave_api_key=self.brave_api_key,
            exec_timeout=self.exec_timeout,
            timezone=self.timezone,
            workspace_index=self.workspace_index,
            narrator=self.narrator,
        )
        await worker.run()
        return task

    @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
