"""Spoke agent — lightweight daemon that connects to a remote Dobby hub.

Polls for tasks, executes them locally, and streams logs back to the hub.
"""
from __future__ import annotations

import asyncio
import logging
import os
import platform
import shutil
from datetime import UTC

import aiohttp

from dobby.core.collector import Collector
from dobby.core.dispatcher import PIPE, STDOUT

logger = logging.getLogger(__name__)


class SpokeAgent:
    """Connects to a Dobby hub and executes tasks on this machine."""

    def __init__(
        self,
        hub_url: str,
        spoke_id: str,
        token: str | None = None,
        poll_interval: float = 5.0,
        results_dir: str = "results",
    ) -> None:
        self.hub_url = hub_url.rstrip("/")
        self.spoke_id = spoke_id
        self.token = token
        self.poll_interval = poll_interval
        self.results_dir = results_dir
        self._running = True

    def _headers(self) -> dict[str, str]:
        h: dict[str, str] = {"Content-Type": "application/json"}
        if self.token:
            h["Authorization"] = f"Bearer {self.token}"
        return h

    def _system_info(self) -> dict:
        return {
            "platform": platform.system(),
            "arch": platform.machine(),
            "hostname": platform.node(),
            "python": platform.python_version(),
            "claude_available": shutil.which("claude") is not None,
            "codex_available": shutil.which("codex") is not None,
        }

    async def _register(self, session: aiohttp.ClientSession) -> None:
        """Register this spoke with the hub."""
        capabilities = []
        if shutil.which("claude"):
            capabilities.append("claude")
        if shutil.which("codex"):
            capabilities.append("codex")

        body = {
            "id": self.spoke_id,
            "capabilities": capabilities,
            "hostname": platform.node(),
            "system_info": self._system_info(),
        }
        async with session.post(
            f"{self.hub_url}/spokes/register",
            json=body,
            headers=self._headers(),
        ) as resp:
            if resp.status == 201:
                logger.info("Registered with hub as '%s'", self.spoke_id)
            else:
                text = await resp.text()
                logger.warning("Registration failed (%d): %s", resp.status, text[:200])

    async def _heartbeat(self, session: aiohttp.ClientSession, running_tasks: int = 0) -> None:
        """Send heartbeat to hub."""
        body = {
            "status": "busy" if running_tasks > 0 else "online",
            "running_tasks": running_tasks,
            "system_info": self._system_info(),
        }
        try:
            async with session.post(
                f"{self.hub_url}/spokes/{self.spoke_id}/heartbeat",
                json=body,
                headers=self._headers(),
            ) as resp:
                if resp.status != 200:
                    logger.warning("Heartbeat failed: %d", resp.status)
        except Exception as exc:
            logger.warning("Heartbeat error: %s", exc)

    async def _claim_task(self, session: aiohttp.ClientSession) -> dict | None:
        """Try to claim a pending task from the hub."""
        try:
            async with session.post(
                f"{self.hub_url}/spokes/{self.spoke_id}/claim",
                headers=self._headers(),
            ) as resp:
                if resp.status != 200:
                    return None
                data = await resp.json()
                return data if data.get("task") else None
        except Exception as exc:
            logger.warning("Claim error: %s", exc)
            return None

    async def _append_log(self, session: aiohttp.ClientSession, task_id: str, lines: str) -> None:
        """Send log lines to the hub."""
        try:
            async with session.post(
                f"{self.hub_url}/tasks/{task_id}/log-append",
                json={"lines": lines},
                headers=self._headers(),
            ) as resp:
                if resp.status != 200:
                    logger.warning("Log append failed: %d", resp.status)
        except Exception:
            pass

    async def _submit_result(self, session: aiohttp.ClientSession, result_dict: dict) -> None:
        """Submit task result to the hub."""
        try:
            async with session.post(
                f"{self.hub_url}/spokes/{self.spoke_id}/result",
                json=result_dict,
                headers=self._headers(),
            ) as resp:
                if resp.status == 201:
                    return
                text = await resp.text()
                # Task may have been cancelled while running — not an error
                if "Cannot transition" in text:
                    logger.info("Task %s: status transition rejected (likely cancelled): %s",
                                result_dict.get("task_id", "?"), text[:100])
                else:
                    logger.warning("Result submit failed (%d): %s", resp.status, text[:200])
        except Exception as exc:
            logger.warning("Result submit error: %s", exc)

    async def _execute_task(
        self, session: aiohttp.ClientSession, task: dict, project: dict | None
    ) -> None:
        """Execute a task locally and stream logs back to hub."""
        task_id = task["id"]
        backend = task.get("backend", "claude")
        description = task.get("description", "")
        exit_criteria = task.get("exit_criteria", "")

        # Build prompt (same as dispatcher)
        prompt_parts = [f"# Project: {task.get('project', 'unknown')}"]
        if project:
            prompt_parts.append(f"Repository: {project.get('repo', '.')}")
            if project.get("context"):
                prompt_parts.append(f"Context: {project['context']}")
        prompt_parts.append(f"\n# Task\nDescription: {description}")
        if exit_criteria:
            prompt_parts.append(f"Exit criteria: {exit_criteria}")
        prompt_parts.append(
            "\nWhen you are done, output a summary in this exact format:\n"
            "---DOBBY_RESULT---\n"
            '{\n  "status": "completed" or "failed",\n'
            '  "summary": "one-line summary of what was done",\n'
            '  "pr_url": "url or null",\n'
            '  "files_changed": ["list", "of", "files"],\n'
            '  "metrics": { ... any numeric results ... },\n'
            '  "error": "error message if failed, null otherwise"\n'
            "}\n---END_DOBBY_RESULT---"
        )
        prompt = "\n".join(prompt_parts)

        # Create worktree for parallel execution (skip when resuming a session)
        repo_path = project.get("repo", ".") if project else "."
        session_id = task.get("session_id")
        from dobby.core.worktree import cleanup_worktree, create_worktree, merge_worktree
        if session_id:
            worktree_path = None  # don't use worktree when resuming
        else:
            worktree_path = (
                await create_worktree(repo_path, task_id) if os.path.isdir(repo_path) else None
            )
        cwd = worktree_path or repo_path

        if backend == "codex":
            args = ["codex", "exec", "--json", "-s", "danger-full-access",
                    "-c", 'model_reasoning_effort="high"', "-C", cwd, prompt]
        else:
            args = ["claude", "-p", prompt, "--output-format", "stream-json",
                    "--verbose", "--max-turns", "200", "--dangerously-skip-permissions"]
            if session_id:
                args.extend(["--resume", session_id])

        logger.info("Executing task %s (%s) in %s", task_id, backend, cwd)

        from datetime import datetime
        started_at = datetime.now(UTC)

        proc = await asyncio.create_subprocess_exec(
            *args,
            stdout=PIPE,
            stderr=STDOUT,
            cwd=cwd,
            limit=10 * 1024 * 1024,
        )

        # Stream output to hub in batches
        output_lines: list[str] = []
        log_buffer: list[str] = []
        flush_interval = 2.0  # seconds
        last_flush = asyncio.get_event_loop().time()

        try:
            while True:
                try:
                    line_bytes = await asyncio.wait_for(
                        proc.stdout.readline(), timeout=flush_interval
                    )
                except TimeoutError:
                    # Flush buffered logs
                    if log_buffer:
                        await self._append_log(session, task_id, "".join(log_buffer))
                        log_buffer.clear()
                    continue

                if not line_bytes:
                    break

                decoded = line_bytes.decode(errors="replace")
                output_lines.append(decoded)
                log_buffer.append(decoded)

                # Flush every flush_interval seconds
                now = asyncio.get_event_loop().time()
                if now - last_flush >= flush_interval:
                    await self._append_log(session, task_id, "".join(log_buffer))
                    log_buffer.clear()
                    last_flush = now

        except Exception as exc:
            logger.error("Task %s execution error: %s", task_id, exc)

        # Flush remaining logs
        if log_buffer:
            await self._append_log(session, task_id, "".join(log_buffer))

        await proc.wait()
        exit_code = proc.returncode or 0

        # Parse result
        full_output = "".join(output_lines)
        result = Collector.parse_output(full_output, exit_code, task_id)
        result_dict = result.to_dict()
        result_dict["started_at"] = started_at.isoformat()

        # Merge worktree back to main if task succeeded
        if worktree_path and result.status == "completed":
            merge_ok, merge_msg = await merge_worktree(repo_path, task_id)
            logger.info(
                "Task %s merge: %s — %s", task_id, "OK" if merge_ok else "FAILED", merge_msg
            )
            result_dict["metrics"] = result_dict.get("metrics", {})
            result_dict["metrics"]["merge"] = merge_msg
        elif worktree_path:
            await cleanup_worktree(repo_path, task_id)

        await self._submit_result(session, result_dict)
        logger.info("Task %s completed: %s", task_id, result.status)

    async def run(self) -> None:
        """Main loop: register, poll for tasks, execute concurrently."""
        logger.info("Spoke agent '%s' connecting to %s", self.spoke_id, self.hub_url)

        async with aiohttp.ClientSession() as session:
            await self._register(session)

            heartbeat_counter = 0
            self._running_count = 0

            async def _run_task(task: dict, project: dict | None) -> None:
                task_id = task.get("id", "?")
                self._running_count += 1
                await self._heartbeat(session, self._running_count)
                try:
                    await self._execute_task(session, task, project)
                except Exception as exc:
                    logger.error("Task %s failed: %s", task_id, exc)
                finally:
                    self._running_count -= 1

            while self._running:
                heartbeat_counter += 1
                if heartbeat_counter >= 6:
                    await self._heartbeat(session, self._running_count)
                    heartbeat_counter = 0

                # Claim and spawn concurrently — don't block
                claimed = await self._claim_task(session)
                if claimed:
                    task = claimed["task"]
                    project = claimed.get("project")
                    asyncio.create_task(_run_task(task, project))

                await asyncio.sleep(self.poll_interval)
