"""Manages tmux sessions for agent sub-processes."""

import asyncio
import fnmatch
import hashlib
import logging
import os
import shlex
import subprocess
import time
from pathlib import Path
from typing import Any, Coroutine, Dict, List, Optional, Tuple

from .file_attacher import FileAttacher
from .models import AgentSession, AgentTask

logger = logging.getLogger(__name__)


class AgentOrchestrator:
    """Manages tmux sessions for agent sub-processes."""

    def __init__(
        self,
        project_name: str = None,
        session_init_delay: float = 3.0,
        kollab_init_delay: float = 4.0,
        message_delay: float = 2.0,
        ready_timeout: float = 30.0,
        ready_poll_interval: float = 0.5,
        ready_stable_threshold: int = 2,
        ready_initial_wait: float = 1.0,
    ):
        """Initialize orchestrator.

        Args:
            project_name: Project name for session naming. Defaults to cwd.
                          Uses KOLLABOR_ROOT_SOCKET env var for nested agents.
            session_init_delay: Seconds to wait after creating tmux session.
            kollab_init_delay: Fallback delay if ready detection fails.
            message_delay: Seconds to wait after sending a message.
            ready_timeout: Max seconds to wait for kollab to be ready.
            ready_poll_interval: Seconds between ready detection polls.
            ready_stable_threshold: Consecutive stable polls required.
            ready_initial_wait: Seconds to wait before starting ready detection.
        """
        # Check for inherited ROOT socket (set by parent kollab instance)
        # This allows nested agents to register on the parent's socket for visibility
        self.project_name = (
            os.environ.get("KOLLABOR_ROOT_SOCKET")  # Inherit from parent
            or project_name  # Use provided name
            or Path.cwd().name  # Default to current folder
        )
        self.agents: Dict[str, AgentSession] = {}
        self.file_attacher = FileAttacher()

        # Track background tasks for cleanup
        self._background_tasks: set = set()

        # Timing configuration
        self.session_init_delay = session_init_delay
        self.kollab_init_delay = kollab_init_delay
        self.message_delay = message_delay

        # Ready detection configuration
        self.ready_timeout = ready_timeout
        self.ready_poll_interval = ready_poll_interval
        self.ready_stable_threshold = ready_stable_threshold
        self.ready_initial_wait = ready_initial_wait

    # -------------------------------------------------------------------------
    # DRY Helpers
    # -------------------------------------------------------------------------

    def _full_name(self, name: str) -> str:
        """Build full tmux session name."""
        return f"{self.project_name}-{name}"

    async def _setup_session(
        self, name: str, clean_stale: bool = True
    ) -> Optional[str]:
        """Create session and track agent.

        Args:
            name: Agent name
            clean_stale: If True, clean up stale sessions with same name

        Returns:
            Full session name on success, None on failure
        """
        full_name = self._full_name(name)

        if self._session_exists(full_name):
            if name not in self.agents and clean_stale:
                logger.warning(f"Found stale session {full_name}, cleaning up...")
                self._kill_session(full_name)
            else:
                logger.warning(f"Agent {name} is already tracked")
                return None

        if not await self._create_session(full_name, agent_name=name):
            logger.error(f"Failed to create session: {full_name}")
            return None

        self.agents[name] = AgentSession(
            name=name,
            full_name=full_name,
            status="initializing",
            start_time=time.time(),
        )
        return full_name

    def _schedule_background(self, coro: Coroutine[Any, Any, None]) -> None:
        """Schedule coroutine as tracked background task."""
        task = asyncio.create_task(coro)
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)

    async def _run_or_schedule(
        self, coro: Coroutine[Any, Any, None], wait: bool, name: str, desc: str
    ) -> None:
        """Run coroutine immediately or schedule in background."""
        if wait:
            await coro
            logger.info(f"Spawn complete for {desc}: {name} (waited)")
        else:
            self._schedule_background(coro)
            logger.info(f"Spawn initiated for {desc}: {name} (background)")

    def _update_status(self, name: str, status: str) -> None:
        """Update agent status if tracked."""
        if name in self.agents:
            self.agents[name].status = status

    def _build_task_message(self, task: str, files: List[str] = None) -> str:
        """Build message with optional file attachments."""
        message = ""
        if files:
            message = self.file_attacher.attach(files) + "\n\n"
        message += task
        return message

    def _build_team_lead_message(self, max_workers: int, task: AgentTask) -> str:
        """Build team lead message with instructions and task."""
        prompt = (
            "You are a team lead agent.\n"
            f"You can spawn up to {max_workers} subagents using <agent> tags.\n"
            "Coordinate their work and integrate results.\n"
            "Use <status /> to check on workers.\n"
            "Use <capture>worker-name 100</capture> to see their progress.\n\n"
        )
        return prompt + self._build_task_message(task.task, task.files)

    async def _wait_for_ready(self, full_name: str) -> bool:
        """Wait for kollab to be ready by detecting stable output.

        Polls the tmux pane and waits for output to stabilize, indicating
        the app has finished initializing and is ready for input.

        Uses instance configuration:
            - ready_timeout: Max seconds to wait
            - ready_poll_interval: Seconds between polls
            - ready_stable_threshold: Consecutive stable polls required
            - ready_initial_wait: Initial wait before polling

        Args:
            full_name: Full tmux session name.

        Returns:
            True if ready detected, False if timeout.
        """
        last_hash = ""
        stable_count = 0
        start = time.time()

        # Wait a minimum time for app to start producing output
        await asyncio.sleep(self.ready_initial_wait)

        while time.time() - start < self.ready_timeout:
            content = self.capture_output(full_name, 50)
            current_hash = hashlib.md5(content.encode()).hexdigest()

            if current_hash == last_hash and content.strip():
                # Content unchanged and not empty
                stable_count += 1
                logger.debug(
                    f"[ready] {full_name} stable count: {stable_count}/{self.ready_stable_threshold}"
                )
                if stable_count >= self.ready_stable_threshold:
                    logger.info(
                        f"[ready] {full_name} ready after {time.time() - start:.1f}s"
                    )
                    return True
            else:
                # Content changed or empty, reset
                if stable_count > 0:
                    logger.debug(f"[ready] {full_name} content changed, resetting")
                stable_count = 0
                last_hash = current_hash

            await asyncio.sleep(self.ready_poll_interval)

        logger.warning(f"[ready] {full_name} timeout after {self.ready_timeout}s")
        return False

    async def _init_and_send_task(
        self,
        name: str,
        full_name: str,
        kollab_cmd: str,
        task_message: str,
    ) -> None:
        """Common initialization: start kollab, send task, update status."""
        try:
            logger.info(f"[init] Sending '{kollab_cmd}' to {full_name}")
            await self._send_keys(full_name, kollab_cmd)

            # Wait for kollab to be ready (polls for stable output)
            logger.info(f"[init] Waiting for {full_name} to be ready...")
            ready = await self._wait_for_ready(full_name)

            if not ready:
                # Fallback to fixed delay if detection fails
                logger.warning("[init] Ready detection failed, using fallback delay")
                await asyncio.sleep(self.kollab_init_delay)

            logger.info(f"[init] Sending task to {full_name}: {task_message[:100]}...")
            result = await self._send_keys(full_name, task_message)
            logger.info(f"[init] Task send result: {result}")

            self._update_status(name, "running")
            logger.info(f"Background initialization complete for agent: {name}")
        except Exception as e:
            logger.error(f"Error during background initialization of {name}: {e}")
            self._update_status(name, "error")

    # -------------------------------------------------------------------------
    # Spawn
    # -------------------------------------------------------------------------

    def _build_kollab_cmd(
        self, agent_type: str = "", skills: List[str] = None, extra_args: str = ""
    ) -> str:
        """Build kollab command with optional flags."""
        cmd = "kollab --simple --permissions trust "
        if agent_type:
            cmd += f" --agent {shlex.quote(agent_type)}"
        for skill in skills or []:
            cmd += f" --skill {shlex.quote(skill)}"
        if extra_args:
            cmd += f" {extra_args}"
        return cmd

    async def spawn(
        self,
        name: str,
        task: str,
        files: List[str] = None,
        wait: bool = False,
        agent_type: str = "",
        skills: List[str] = None,
    ) -> bool:
        """Spawn a new agent session.

        By default, this is NON-BLOCKING - it creates the session and kicks off
        background initialization, returning immediately so multiple agents
        can be spawned in parallel.

        If wait=True, waits for full initialization before returning.

        Args:
            name: Agent name
            task: Task description
            files: Optional list of files to attach
            wait: If True, wait for initialization to complete
            agent_type: Optional agent type to use (e.g., "coder", "research")
            skills: Optional list of skills to load

        Returns:
            True if spawn was initiated successfully
        """
        full_name = await self._setup_session(name)
        if not full_name:
            return False

        kollab_cmd = self._build_kollab_cmd(agent_type, skills)
        message = self._build_task_message(task, files)

        coro = self._init_and_send_task(name, full_name, kollab_cmd, message)
        await self._run_or_schedule(coro, wait, name, "agent")

        return True

    async def spawn_clone(
        self,
        name: str,
        task: str,
        files: List[str],
        conversation_file: str,
        wait: bool = False,
    ) -> bool:
        """Spawn agent with conversation context.

        By default NON-BLOCKING. If wait=True, waits for initialization.

        Args:
            name: Agent name
            task: Task description
            files: Files to attach
            conversation_file: Path to exported conversation JSON
            wait: If True, wait for initialization to complete

        Returns:
            True if spawn was initiated successfully
        """
        full_name = await self._setup_session(name, clean_stale=False)
        if not full_name:
            return False

        kollab_cmd = f"kollab --resume {shlex.quote(conversation_file)}"
        message = self._build_task_message(task, files)

        coro = self._init_and_send_task(name, full_name, kollab_cmd, message)
        await self._run_or_schedule(coro, wait, name, "clone agent")

        return True

    async def spawn_team_lead(
        self,
        lead_name: str,
        max_workers: int,
        task: AgentTask,
        wait: bool = False,
    ) -> bool:
        """Spawn a team lead agent that can spawn workers.

        By default NON-BLOCKING. If wait=True, waits for initialization.

        Args:
            lead_name: Lead agent name
            max_workers: Maximum number of workers the lead can spawn
            task: Task for the lead
            wait: If True, wait for initialization to complete

        Returns:
            True if spawn was initiated successfully
        """
        full_name = await self._setup_session(lead_name, clean_stale=False)
        if not full_name:
            return False

        message = self._build_team_lead_message(max_workers, task)
        coro = self._init_and_send_task(lead_name, full_name, "kollab", message)
        await self._run_or_schedule(coro, wait, lead_name, "team lead")

        return True

    # -------------------------------------------------------------------------
    # Message
    # -------------------------------------------------------------------------

    async def message(self, name: str, content: str) -> bool:
        """Send message to agent.

        Args:
            name: Agent name
            content: Message content

        Returns:
            True if successful
        """
        full_name = self._full_name(name)

        if not self._session_exists(full_name):
            logger.warning(f"Session does not exist: {full_name}")
            return False

        await self._send_keys(full_name, content)
        logger.info(f"Sent message to agent: {name}")
        return True

    # -------------------------------------------------------------------------
    # Stop
    # -------------------------------------------------------------------------

    async def stop(self, name: str) -> Tuple[str, str]:
        """Stop agent and return final output + duration.

        Args:
            name: Agent name

        Returns:
            Tuple of (output, duration)
        """
        full_name = self._full_name(name)

        # Wait for agent to finish initializing if needed
        agent = self.agents.get(name)
        max_wait = 10  # Maximum seconds to wait for initialization
        wait_start = time.time()

        while (
            agent
            and agent.status == "initializing"
            and (time.time() - wait_start) < max_wait
        ):
            logger.debug(f"Waiting for {name} to finish initializing...")
            await asyncio.sleep(0.5)
            agent = self.agents.get(name)

        # Capture final output
        output = self.capture_output(name, 100)

        # Get duration
        duration = agent.duration if agent else "?"

        # Kill session
        self._kill_session(full_name)

        # Remove from tracking
        if name in self.agents:
            del self.agents[name]

        logger.info(f"Stopped agent: {name} @ {duration}")
        return output, duration

    # -------------------------------------------------------------------------
    # Status / Capture
    # -------------------------------------------------------------------------

    def list_agents(self) -> List[AgentSession]:
        """List all active agents.

        Returns:
            List of agent sessions
        """
        # Refresh status from tmux
        self._refresh_agents()
        return list(self.agents.values())

    def get_agent(self, name: str) -> Optional[AgentSession]:
        """Get specific agent.

        Args:
            name: Agent name

        Returns:
            Agent session or None
        """
        return self.agents.get(name)

    def get_all_agents(self) -> Dict[str, AgentSession]:
        """Get all tracked agents.

        Returns:
            Dictionary of agent name to AgentSession.
        """
        # Refresh to ensure we have latest state
        self._refresh_agents()
        return self.agents.copy()

    def find_agents(self, pattern: str) -> List[str]:
        """Find agents matching glob pattern.

        Args:
            pattern: Glob pattern (e.g., "lint-*")

        Returns:
            List of matching agent names
        """
        return [name for name in self.agents.keys() if fnmatch.fnmatch(name, pattern)]

    def capture_output(self, name: str, lines: int = 50) -> str:
        """Capture last N lines from agent.

        Args:
            name: Agent name
            lines: Number of lines to capture

        Returns:
            Captured output
        """
        full_name = self._full_name(name)

        result = subprocess.run(
            self._tmux_cmd("capture-pane", "-t", full_name, "-p", "-S", f"-{lines}"),
            capture_output=True,
            text=True,
        )
        return result.stdout

    # -------------------------------------------------------------------------
    # Private Helpers
    # -------------------------------------------------------------------------

    def _tmux_cmd(self, *args) -> List[str]:
        """Build tmux command with project socket isolation.

        Args:
            *args: Tmux command arguments

        Returns:
            List of command parts with socket flag
        """
        return ["tmux", "-L", self.project_name] + list(args)

    def _session_exists(self, full_name: str) -> bool:
        """Check if tmux session exists."""
        result = subprocess.run(
            self._tmux_cmd("has-session", "-t", full_name), capture_output=True
        )
        return result.returncode == 0

    async def _create_session(self, full_name: str, agent_name: str = "") -> bool:
        """Create new tmux session in current working directory.

        Args:
            full_name: Full tmux session name (project-agent)
            agent_name: Agent name for KOLLABOR_AGENT_NAME env var
        """
        cwd = str(Path.cwd())
        root_socket = os.environ.get("KOLLABOR_ROOT_SOCKET", self.project_name)

        cmd = self._tmux_cmd(
            "new-session",
            "-d",
            "-s",
            full_name,
            "-e",
            f"KOLLABOR_ROOT_SOCKET={root_socket}",
            "-e",
            f"KOLLABOR_AGENT_NAME={agent_name}",
            "-c",
            cwd,
        )

        result = subprocess.run(cmd, capture_output=True)
        if result.returncode != 0:
            logger.error(f"Failed to create tmux session: {result.stderr}")
            return False

        await asyncio.sleep(self.session_init_delay)
        return True

    async def _send_keys(self, full_name: str, content: str) -> bool:
        """Send keys to tmux session."""
        logger.debug(f"[send_keys] Sending to {full_name}: {len(content)} chars")

        # Send content first
        result = subprocess.run(
            self._tmux_cmd("send-keys", "-t", full_name, content),
            capture_output=True,
        )
        if result.returncode != 0:
            logger.error(f"[send_keys] Content send failed: {result.stderr}")
            return False

        logger.debug("[send_keys] Content sent, waiting 1s before Enter")
        # Wait for paste to be processed (non-blocking)
        await asyncio.sleep(1)

        # Then send Enter to submit
        logger.debug("[send_keys] Sending Enter")
        result = subprocess.run(
            self._tmux_cmd("send-keys", "-t", full_name, "Enter"),
            capture_output=True,
        )
        if result.returncode != 0:
            logger.error(f"[send_keys] Enter send failed: {result.stderr}")

        logger.debug(f"[send_keys] Waiting {self.message_delay}s message delay")
        await asyncio.sleep(self.message_delay)
        return result.returncode == 0

    def _kill_session(self, full_name: str) -> bool:
        """Kill tmux session."""
        result = subprocess.run(
            self._tmux_cmd("kill-session", "-t", full_name), capture_output=True
        )
        if result.returncode == 0:
            logger.info(f"Killed session: {full_name}")
        else:
            logger.warning(f"Failed to kill session {full_name}: {result.stderr}")
        return result.returncode == 0

    def cleanup_stale_sessions(self) -> int:
        """Clean up stale tmux sessions that are no longer tracked.

        Returns:
            Number of sessions cleaned up.
        """
        # Get all sessions for this project
        result = subprocess.run(
            self._tmux_cmd("list-sessions", "-F", "#{session_name}"),
            capture_output=True,
            text=True,
        )

        if result.returncode != 0:
            return 0

        stdout = result.stdout.strip()
        if not stdout:
            return 0

        all_sessions = stdout.split("\n")
        cleaned = 0

        for session in all_sessions:
            if not session.startswith(f"{self.project_name}-"):
                continue

            # Extract agent name from session name
            agent_name = session[len(f"{self.project_name}-") :]

            # Skip if we're tracking this agent (even if initializing)
            # If we're tracking it, it's not stale
            if agent_name in self.agents:
                # But if it's in error state, it might be worth cleaning up
                if self.agents[agent_name].status != "error":
                    continue

            # If we're not tracking this agent, it's stale
            logger.info(f"Cleaning up stale session: {session}")
            if self._kill_session(session):
                cleaned += 1

        if cleaned > 0:
            logger.info(f"Cleaned up {cleaned} stale session(s)")

        return cleaned

    def _refresh_agents(self) -> None:
        """Refresh agent status from tmux."""
        result = subprocess.run(
            self._tmux_cmd("list-sessions", "-F", "#{session_name}"),
            capture_output=True,
            text=True,
        )

        if result.returncode != 0:
            return

        stdout = result.stdout.strip()
        if not stdout:
            active_sessions = set()
        else:
            active_sessions = set(stdout.split("\n"))

        prefix = f"{self.project_name}-"

        # Remove dead agents (but skip initializing agents to avoid race condition)
        dead = [
            name
            for name in self.agents
            if f"{prefix}{name}" not in active_sessions
            and self.agents[name].status != "initializing"
        ]
        for name in dead:
            del self.agents[name]

        # Discover new agents (created externally)
        for session in active_sessions:
            if session.startswith(prefix):
                agent_name = session[len(prefix) :]
                if agent_name and agent_name not in self.agents:
                    self.agents[agent_name] = AgentSession(
                        name=agent_name,
                        full_name=session,
                        status="running",
                        start_time=time.time(),  # approximate
                    )

    async def shutdown(self) -> None:
        """Cancel all pending background tasks."""
        if self._background_tasks:
            count = len(self._background_tasks)
            logger.info(f"Cancelling {count} pending background tasks")
            for task in self._background_tasks:
                if not task.done():
                    task.cancel()
            # Wait for tasks to be cancelled
            await asyncio.gather(*self._background_tasks, return_exceptions=True)
            self._background_tasks.clear()
            logger.info("All background tasks cancelled")
