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

import subprocess
import time
import asyncio
import fnmatch
import logging
from pathlib import Path
from typing import List, Optional, Tuple, Dict

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

logger = logging.getLogger(__name__)


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

    def __init__(self, project_name: str = None):
        """Initialize orchestrator.

        Args:
            project_name: Project name for session naming. Defaults to cwd name.
        """
        self.project_name = project_name or Path.cwd().name
        self.agents: Dict[str, AgentSession] = {}
        self.file_attacher = FileAttacher()

        # Timing configuration
        self.session_init_delay = 2
        self.kollab_init_delay = 7  # Allow time for kollab to fully initialize
        self.message_delay = 2

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

    async def spawn(self, name: str, task: str, files: List[str] = None, wait: bool = False) -> 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

        Returns:
            True if spawn was initiated successfully
        """
        full_name = f"{self.project_name}-{name}"

        # Check if already exists
        if self._session_exists(full_name):
            logger.warning(f"Session already exists: {full_name}")
            return False

        # Create session (has small delay, but necessary)
        if not await self._create_session(full_name):
            logger.error(f"Failed to create session: {full_name}")
            return False

        # Track agent immediately (before background init)
        self.agents[name] = AgentSession(
            name=name,
            full_name=full_name,
            status="initializing",
            start_time=time.time(),
        )

        if wait:
            # Wait for full initialization (blocking)
            await self._initialize_agent_background(
                name=name,
                full_name=full_name,
                task=task,
                files=files or []
            )
            logger.info(f"Spawn complete for agent: {name} (waited for initialization)")
        else:
            # Kick off background initialization (NON-BLOCKING)
            # This allows multiple agents to be spawned in parallel
            asyncio.create_task(self._initialize_agent_background(
                name=name,
                full_name=full_name,
                task=task,
                files=files or []
            ))
            logger.info(f"Spawn initiated for agent: {name} (initializing in background)")

        return True

    async def _initialize_agent_background(
        self,
        name: str,
        full_name: str,
        task: str,
        files: List[str],
    ) -> None:
        """Initialize agent in the background (non-blocking).

        This runs after spawn() returns, allowing parallel agent spawning.

        Args:
            name: Agent name
            full_name: Full tmux session name
            task: Task description
            files: Files to attach
        """
        try:
            # Start kollab
            logger.info(f"[init] Sending 'kollab' command to {full_name}")
            await self._send_keys(full_name, "kollab")
            logger.info(f"[init] Waiting {self.kollab_init_delay}s for kollab to initialize")
            await asyncio.sleep(self.kollab_init_delay)

            # Prepare message with attached files
            message = ""
            if files:
                message = self.file_attacher.attach(files)
                message += "\n\n"
            message += task

            logger.info(f"[init] Sending task to {full_name}: {message[:100]}...")

            # Send task
            result = await self._send_keys(full_name, message)
            logger.info(f"[init] Task send result: {result}")

            # Update agent status to running
            if name in self.agents:
                self.agents[name].status = "running"

            logger.info(f"Background initialization complete for agent: {name}")
        except Exception as e:
            logger.error(f"Error during background initialization of {name}: {e}")
            # Update agent status to error
            if name in self.agents:
                self.agents[name].status = "error"

    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 = f"{self.project_name}-{name}"

        if self._session_exists(full_name):
            logger.warning(f"Session already exists: {full_name}")
            return False

        if not await self._create_session(full_name):
            logger.error(f"Failed to create session: {full_name}")
            return False

        # Track agent immediately
        self.agents[name] = AgentSession(
            name=name,
            full_name=full_name,
            status="initializing",
            start_time=time.time(),
        )

        if wait:
            # Wait for full initialization (blocking)
            await self._initialize_clone_background(
                name=name,
                full_name=full_name,
                task=task,
                files=files or [],
                conversation_file=conversation_file,
            )
            logger.info(f"Spawn complete for clone agent: {name} (waited for initialization)")
        else:
            # Kick off background initialization
            asyncio.create_task(self._initialize_clone_background(
                name=name,
                full_name=full_name,
                task=task,
                files=files or [],
                conversation_file=conversation_file,
            ))
            logger.info(f"Spawn initiated for clone agent: {name} (initializing in background)")

        return True

    async def _initialize_clone_background(
        self,
        name: str,
        full_name: str,
        task: str,
        files: List[str],
        conversation_file: str,
    ) -> None:
        """Initialize clone agent in the background (non-blocking).

        Args:
            name: Agent name
            full_name: Full tmux session name
            task: Task description
            files: Files to attach
            conversation_file: Path to exported conversation JSON
        """
        try:
            # Start kollab with resume
            await self._send_keys(full_name, f"kollab --resume {conversation_file}")
            await asyncio.sleep(self.kollab_init_delay)

            # Send task with files
            message = ""
            if files:
                message = self.file_attacher.attach(files) + "\n\n"
            message += task

            await self._send_keys(full_name, message)

            # Update agent status to running
            if name in self.agents:
                self.agents[name].status = "running"

            logger.info(f"Background initialization complete for clone agent: {name}")
        except Exception as e:
            logger.error(f"Error during background initialization of clone {name}: {e}")
            if name in self.agents:
                self.agents[name].status = "error"

    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 = f"{self.project_name}-{lead_name}"

        if self._session_exists(full_name):
            logger.warning(f"Session already exists: {full_name}")
            return False

        if not await self._create_session(full_name):
            logger.error(f"Failed to create session: {full_name}")
            return False

        # Track agent immediately
        self.agents[lead_name] = AgentSession(
            name=lead_name,
            full_name=full_name,
            status="initializing",
            start_time=time.time(),
        )

        if wait:
            # Wait for full initialization (blocking)
            await self._initialize_team_lead_background(
                lead_name=lead_name,
                full_name=full_name,
                max_workers=max_workers,
                task=task,
            )
            logger.info(f"Spawn complete for team lead: {lead_name} (waited for initialization)")
        else:
            # Kick off background initialization
            asyncio.create_task(self._initialize_team_lead_background(
                lead_name=lead_name,
                full_name=full_name,
                max_workers=max_workers,
                task=task,
            ))
            logger.info(f"Spawn initiated for team lead: {lead_name} (initializing in background)")

        return True

    async def _initialize_team_lead_background(
        self,
        lead_name: str,
        full_name: str,
        max_workers: int,
        task: AgentTask,
    ) -> None:
        """Initialize team lead agent in the background (non-blocking).

        Args:
            lead_name: Lead agent name
            full_name: Full tmux session name
            max_workers: Maximum number of workers the lead can spawn
            task: Task for the lead
        """
        try:
            await self._send_keys(full_name, "kollab")
            await asyncio.sleep(self.kollab_init_delay)

            # Inject team lead prompt
            lead_prompt = f"""You are a team lead agent.
You can spawn up to {max_workers} worker agents using <agent> tags.
Coordinate their work and integrate results.
Use <status /> to check on workers.
Use <capture>worker-name 100</capture> to see their progress.

"""
            message = lead_prompt
            if task.files:
                message += self.file_attacher.attach(task.files) + "\n\n"
            message += task.task

            await self._send_keys(full_name, message)

            # Update agent status to running
            if lead_name in self.agents:
                self.agents[lead_name].status = "running"

            logger.info(f"Background initialization complete for team lead: {lead_name}")
        except Exception as e:
            logger.error(f"Error during background initialization of team lead {lead_name}: {e}")
            if lead_name in self.agents:
                self.agents[lead_name].status = "error"

    # -------------------------------------------------------------------------
    # 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 = f"{self.project_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 = f"{self.project_name}-{name}"

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

        # Get duration
        agent = self.agents.get(name)
        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 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 = f"{self.project_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) -> bool:
        """Create new tmux session in current working directory."""
        cwd = str(Path.cwd())
        result = subprocess.run(
            self._tmux_cmd("new-session", "-d", "-s", full_name, "-c", cwd), 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(f"[send_keys] Content sent, waiting 2s before Enter")
        # Wait for paste to be processed (non-blocking)
        await asyncio.sleep(2)

        # Then send Enter to submit
        logger.debug(f"[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
        )
        return result.returncode == 0

    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

        active_sessions = set(result.stdout.strip().split("\n"))
        prefix = f"{self.project_name}-"

        # Remove dead agents
        dead = [
            name
            for name in self.agents
            if f"{prefix}{name}" not in active_sessions
        ]
        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
                    )
