"""Background command sessions for long-running local processes."""

from __future__ import annotations

import asyncio
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
import shlex
import time

from pydantic import BaseModel, Field

from yier_agents.src.tool import Tool, ToolContext, ToolOutput
from yier_agents.src.tools.workspace import (
    SHELL_EVENT_EXIT,
    SHELL_EVENT_STARTED,
    SHELL_EVENT_STATE_CHANGED,
    SHELL_EVENT_STDERR,
    SHELL_EVENT_STDIN,
    SHELL_EVENT_STDOUT,
    WorkspaceAccess,
    _build_shell_event,
    _build_shell_process_payload,
    _build_shell_raw_payload,
    _default_shell_program,
    _normalize_allowed_roots,
    _normalize_default_root,
    _validate_command,
)


@dataclass(slots=True)
class OutputRingBuffer:
    """Keep only the latest output up to a character limit."""

    max_chars: int
    chunks: deque[str] = field(default_factory=deque)
    total_chars: int = 0
    truncated: bool = False

    def append(self, text: str) -> None:
        if not text:
            return

        self.chunks.append(text)
        self.total_chars += len(text)
        self._trim()

    def render(self, max_chars: int | None = None) -> str:
        text = "".join(self.chunks)
        if max_chars is None or len(text) <= max_chars:
            return text
        return text[-max_chars:]

    def _trim(self) -> None:
        while self.total_chars > self.max_chars and self.chunks:
            overflow = self.total_chars - self.max_chars
            head = self.chunks[0]
            if len(head) <= overflow:
                self.chunks.popleft()
                self.total_chars -= len(head)
                self.truncated = True
                continue

            self.chunks[0] = head[overflow:]
            self.total_chars -= overflow
            self.truncated = True


@dataclass(slots=True)
class EventRingBuffer:
    """Keep only the latest raw shell events up to an event limit."""

    max_events: int
    events: deque[dict[str, object]] = field(default_factory=deque)
    next_index: int = 0
    dropped_count: int = 0

    def append(
        self,
        event_type: str,
        *,
        timestamp: float | None = None,
        **payload: object,
    ) -> dict[str, object]:
        event = _build_shell_event(
            self.next_index,
            event_type,
            time.time() if timestamp is None else timestamp,
            **payload,
        )
        self.events.append(event)
        self.next_index += 1
        while len(self.events) > self.max_events:
            self.events.popleft()
            self.dropped_count += 1
        return event

    def latest_index(self) -> int | None:
        if self.next_index == 0:
            return None
        return self.next_index - 1

    def render(
        self,
        *,
        after_index: int | None = None,
    ) -> tuple[list[dict[str, object]], bool]:
        if after_index is None:
            return list(self.events), self.dropped_count > 0

        first_retained_index = self.dropped_count
        was_truncated = after_index < first_retained_index - 1
        filtered = [
            event
            for event in self.events
            if isinstance(event.get("index"), int) and event["index"] > after_index
        ]
        return filtered, was_truncated


@dataclass(slots=True)
class BackgroundCommandSession:
    """One long-running background command."""

    session_id: str
    command: str
    cwd: Path
    process: asyncio.subprocess.Process
    allow_shell: bool
    shell_program: str | None
    started_at_monotonic: float
    started_at_timestamp: float
    stdout_buffer: OutputRingBuffer
    stderr_buffer: OutputRingBuffer
    event_buffer: EventRingBuffer
    state: str = "running"
    exit_code: int | None = None
    finished_at_monotonic: float | None = None
    finished_at_timestamp: float | None = None
    completion_notified: bool = False
    completion_event: asyncio.Event = field(default_factory=asyncio.Event)
    tasks: list[asyncio.Task] = field(default_factory=list)

    def runtime_seconds(self) -> int:
        finished_at = self.finished_at_monotonic or time.monotonic()
        return max(0, int(finished_at - self.started_at_monotonic))

    def is_running(self) -> bool:
        return self.state == "running"


class StartBackgroundCommandParams(BaseModel):
    """Parameters for starting one background command."""

    command: str = Field(description="Command to execute in the background.")
    cwd: str | None = Field(default=None, description="Optional working directory.")


class ListBackgroundCommandsParams(BaseModel):
    """Parameters for listing background commands."""

    include_completed: bool = Field(
        default=True,
        description="Whether to include finished commands.",
    )


class ReadBackgroundCommandParams(BaseModel):
    """Parameters for reading command output."""

    session_id: str = Field(description="Background command session id.")
    max_output_chars: int = Field(
        default=6000,
        ge=1,
        le=20000,
        description="Maximum characters of output to return.",
    )
    after_event_index: int | None = Field(
        default=None,
        ge=0,
        description="Optional event index cursor to fetch only newer raw shell events.",
    )


class WaitBackgroundCommandParams(BaseModel):
    """Parameters for waiting on one background command."""

    session_id: str = Field(description="Background command session id.")
    timeout_seconds: int = Field(
        default=30,
        ge=1,
        le=300,
        description="Maximum time to wait for completion.",
    )
    max_output_chars: int = Field(
        default=6000,
        ge=1,
        le=20000,
        description="Maximum characters of output to return.",
    )
    after_event_index: int | None = Field(
        default=None,
        ge=0,
        description="Optional event index cursor to fetch only newer raw shell events.",
    )


class StopBackgroundCommandParams(BaseModel):
    """Parameters for stopping one background command."""

    session_id: str = Field(description="Background command session id.")


class SendBackgroundCommandInputParams(BaseModel):
    """Parameters for sending stdin to a background command."""

    session_id: str = Field(description="Background command session id.")
    text: str = Field(description="Text to send to stdin.")
    append_newline: bool = Field(
        default=False,
        description="Whether to append a newline after the text.",
    )


class BackgroundCommandManager:
    """Manage long-running subprocesses for the CLI."""

    def __init__(
        self,
        allowed_roots: list[Path] | None = None,
        default_root: Path | None = None,
        *,
        allow_shell: bool = False,
        shell_program: str | None = None,
        max_buffer_chars: int = 40000,
        max_event_buffer_size: int = 1000,
    ) -> None:
        self.access = WorkspaceAccess(
            allowed_roots=_normalize_allowed_roots(allowed_roots),
            default_root=_normalize_default_root(default_root),
        )
        self.allow_shell = allow_shell
        self.shell_program = shell_program or _default_shell_program()
        self.max_buffer_chars = max_buffer_chars
        self.max_event_buffer_size = max_event_buffer_size
        self._sessions: dict[str, BackgroundCommandSession] = {}
        self._session_counter = 0

    async def start(self, command: str, cwd: str | None = None) -> BackgroundCommandSession:
        """Start one subprocess and return its session."""
        _validate_command(command, allow_shell=self.allow_shell)

        default_cwd = self.access.default_root or Path.cwd()
        working_directory = self.access.resolve_file(cwd or str(default_cwd))
        if not working_directory.exists():
            raise FileNotFoundError(f"Working directory not found: {working_directory}")
        if not working_directory.is_dir():
            raise NotADirectoryError(
                f"Working directory is not a directory: {working_directory}"
            )

        process = await self._spawn_process(command, working_directory)
        self._session_counter += 1
        session_id = f"bg-{self._session_counter}"
        started_at_timestamp = time.time()
        session = BackgroundCommandSession(
            session_id=session_id,
            command=command,
            cwd=working_directory,
            process=process,
            allow_shell=self.allow_shell,
            shell_program=self.shell_program if self.allow_shell else None,
            started_at_monotonic=time.monotonic(),
            started_at_timestamp=started_at_timestamp,
            stdout_buffer=OutputRingBuffer(self.max_buffer_chars),
            stderr_buffer=OutputRingBuffer(self.max_buffer_chars),
            event_buffer=EventRingBuffer(self.max_event_buffer_size),
        )
        session.event_buffer.append(
            SHELL_EVENT_STARTED,
            timestamp=started_at_timestamp,
            command=command,
            cwd=str(working_directory),
            allow_shell=self.allow_shell,
            shell_program=session.shell_program,
        )
        self._sessions[session_id] = session

        if process.stdout is not None:
            session.tasks.append(
                asyncio.create_task(
                    self._read_stream(
                        process.stdout,
                        session.stdout_buffer,
                        session,
                        SHELL_EVENT_STDOUT,
                    )
                )
            )
        if process.stderr is not None:
            session.tasks.append(
                asyncio.create_task(
                    self._read_stream(
                        process.stderr,
                        session.stderr_buffer,
                        session,
                        SHELL_EVENT_STDERR,
                    )
                )
            )
        session.tasks.append(asyncio.create_task(self._watch_process(session)))
        return session

    def list_sessions(
        self,
        *,
        include_completed: bool = True,
    ) -> list[BackgroundCommandSession]:
        """List known sessions in creation order."""
        sessions = sorted(
            self._sessions.values(),
            key=lambda item: int(item.session_id.removeprefix("bg-")),
        )
        if include_completed:
            return sessions
        return [session for session in sessions if session.is_running()]

    def count_running(self) -> int:
        """Return the number of active background commands."""
        return sum(1 for session in self._sessions.values() if session.is_running())

    def active_session_ids(self) -> tuple[str, ...]:
        """Return running session ids for CLI completion."""
        return tuple(session.session_id for session in self.list_sessions())

    def drain_notifications(self) -> list[str]:
        """Return newly finished sessions that have not been surfaced yet."""
        notifications: list[str] = []
        for session in self.drain_completed_sessions():
            if session.state == "completed":
                message = (
                    f"Background command {session.session_id} finished successfully: "
                    f"{session.command}"
                )
            elif session.state == "stopped":
                message = (
                    f"Background command {session.session_id} was stopped: "
                    f"{session.command}"
                )
            else:
                message = (
                    f"Background command {session.session_id} exited with code "
                    f"{session.exit_code}: {session.command}"
                )
            notifications.append(message)
        return notifications

    def drain_completed_sessions(self) -> list[BackgroundCommandSession]:
        """Return newly completed sessions that have not been surfaced yet."""
        completed_sessions: list[BackgroundCommandSession] = []
        for session in self.list_sessions():
            if session.is_running() or session.completion_notified:
                continue
            session.completion_notified = True
            completed_sessions.append(session)
        return completed_sessions

    async def wait(
        self,
        session_id: str,
        timeout_seconds: int,
    ) -> BackgroundCommandSession:
        """Wait for one session to complete or time out."""
        session = self.require_session(session_id)
        if not session.is_running():
            return session

        try:
            await asyncio.wait_for(
                session.completion_event.wait(),
                timeout=timeout_seconds,
            )
        except TimeoutError:
            return session
        return session

    async def stop(self, session_id: str) -> BackgroundCommandSession:
        """Stop one background command."""
        session = self.require_session(session_id)
        if not session.is_running():
            return session

        session.state = "stopping"
        session.event_buffer.append(
            SHELL_EVENT_STATE_CHANGED,
            state=session.state,
        )
        session.process.terminate()
        try:
            await asyncio.wait_for(session.process.wait(), timeout=3)
        except TimeoutError:
            session.process.kill()
            await session.process.wait()
        if not session.completion_event.is_set():
            try:
                await asyncio.wait_for(session.completion_event.wait(), timeout=1)
            except TimeoutError:
                session.exit_code = session.process.returncode
                session.finished_at_monotonic = time.monotonic()
                session.finished_at_timestamp = time.time()
                session.state = "stopped"
                session.event_buffer.append(
                    SHELL_EVENT_STATE_CHANGED,
                    timestamp=session.finished_at_timestamp,
                    state=session.state,
                )
                session.event_buffer.append(
                    SHELL_EVENT_EXIT,
                    timestamp=session.finished_at_timestamp,
                    exit_code=session.exit_code,
                    state=session.state,
                    timed_out=False,
                )
                session.completion_event.set()
        return session

    async def write_stdin(
        self,
        session_id: str,
        text: str,
        *,
        append_newline: bool = False,
    ) -> BackgroundCommandSession:
        """Send text to a running background command."""
        session = self.require_session(session_id)
        if not session.is_running():
            raise ValueError(f"Background command is not running: {session.session_id}")
        if session.process.stdin is None:
            raise ValueError(f"Background command does not accept stdin: {session.session_id}")

        payload = f"{text}\n" if append_newline else text
        session.process.stdin.write(payload.encode("utf-8"))
        await session.process.stdin.drain()
        session.event_buffer.append(
            SHELL_EVENT_STDIN,
            text=payload,
            append_newline=append_newline,
        )
        return session

    def require_session(self, session_id: str) -> BackgroundCommandSession:
        """Resolve one session id or unique prefix."""
        normalized = session_id.strip().lower()
        if normalized in self._sessions:
            return self._sessions[normalized]

        matches = [
            session
            for session in self._sessions.values()
            if session.session_id.startswith(normalized)
        ]
        if len(matches) == 1:
            return matches[0]
        if len(matches) > 1:
            raise ValueError(f"Session id prefix is ambiguous: {session_id}")
        raise KeyError(f"Unknown background command session: {session_id}")

    def format_session_list(self, *, include_completed: bool = True) -> str:
        """Render a compact list of known sessions."""
        sessions = self.list_sessions(include_completed=include_completed)
        if not sessions:
            return "No background commands."

        lines = ["Background commands:"]
        for session in sessions:
            exit_text = (
                f"exit={session.exit_code}"
                if session.exit_code is not None
                else "running"
            )
            lines.append(
                f"- {session.session_id} [{session.state}] {exit_text} "
                f"({session.runtime_seconds()}s) {session.command}"
            )
        return "\n".join(lines)

    def format_session_output(
        self,
        session_id: str,
        *,
        max_output_chars: int,
    ) -> str:
        """Render one session with recent stdout/stderr."""
        session = self.require_session(session_id)
        header = [
            f"Session: {session.session_id}",
            f"Command: {session.command}",
            f"Working directory: {session.cwd}",
            f"State: {session.state}",
            f"Runtime: {session.runtime_seconds()}s",
        ]
        if session.exit_code is not None:
            header.append(f"Exit code: {session.exit_code}")

        stdout_text = session.stdout_buffer.render(max_output_chars)
        stderr_text = session.stderr_buffer.render(max_output_chars)

        sections = header
        if stdout_text:
            sections.extend(["", "[stdout]", stdout_text])
        if stderr_text:
            sections.extend(["", "[stderr]", stderr_text])
        if not stdout_text and not stderr_text:
            sections.extend(["", "No output captured yet."])
        return "\n".join(sections)

    def get_session_raw_payload(
        self,
        session_id: str,
        *,
        after_event_index: int | None = None,
        append_newline: bool | None = None,
    ) -> dict[str, object]:
        """Build the structured shell payload for one background session."""
        session = self.require_session(session_id)
        events, events_truncated = session.event_buffer.render(
            after_index=after_event_index,
        )

        request: dict[str, object] = {
            "command": session.command,
            "cwd": str(session.cwd),
            "allow_shell": session.allow_shell,
            "shell_program": session.shell_program,
        }
        if append_newline is not None:
            request["append_newline"] = append_newline

        return _build_shell_raw_payload(
            kind="background_command",
            request=request,
            process=_build_shell_process_payload(
                session_id=session.session_id,
                state=session.state,
                exit_code=session.exit_code,
                started_at=session.started_at_timestamp,
                finished_at=session.finished_at_timestamp,
                runtime_seconds=session.runtime_seconds(),
                timed_out=False,
            ),
            events=events,
            latest_event_index=session.event_buffer.latest_index(),
            stdout_text=session.stdout_buffer.render(),
            stderr_text=session.stderr_buffer.render(),
            stdout_truncated=session.stdout_buffer.truncated,
            stderr_truncated=session.stderr_buffer.truncated,
            events_truncated=events_truncated,
            dropped_event_count=session.event_buffer.dropped_count,
        )

    async def close(self) -> None:
        """Stop all running background commands."""
        for session in self.list_sessions(include_completed=False):
            await self.stop(session.session_id)

    async def _spawn_process(
        self,
        command: str,
        working_directory: Path,
    ) -> asyncio.subprocess.Process:
        if self.allow_shell:
            return await asyncio.create_subprocess_shell(
                command,
                cwd=str(working_directory),
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                executable=self.shell_program,
            )

        argv = shlex.split(command)
        if not argv:
            raise ValueError("Command must not be empty")

        return await asyncio.create_subprocess_exec(
            *argv,
            cwd=str(working_directory),
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

    async def _read_stream(
        self,
        stream: asyncio.StreamReader,
        buffer: OutputRingBuffer,
        session: BackgroundCommandSession,
        event_type: str,
    ) -> None:
        while True:
            chunk = await stream.read(4096)
            if not chunk:
                break
            text = chunk.decode("utf-8", errors="replace")
            buffer.append(text)
            session.event_buffer.append(
                event_type,
                text=text,
                stream="stdout" if event_type == SHELL_EVENT_STDOUT else "stderr",
            )

    async def _watch_process(self, session: BackgroundCommandSession) -> None:
        exit_code = await session.process.wait()
        session.exit_code = exit_code
        session.finished_at_monotonic = time.monotonic()
        session.finished_at_timestamp = time.time()
        if session.state in {"stopping", "stopped"}:
            session.state = "stopped"
        elif exit_code == 0:
            session.state = "completed"
        else:
            session.state = "failed"
        session.event_buffer.append(
            SHELL_EVENT_STATE_CHANGED,
            timestamp=session.finished_at_timestamp,
            state=session.state,
        )
        session.event_buffer.append(
            SHELL_EVENT_EXIT,
            timestamp=session.finished_at_timestamp,
            exit_code=exit_code,
            state=session.state,
            timed_out=False,
        )
        session.completion_event.set()


def create_start_background_command_tool(
    manager: BackgroundCommandManager,
) -> Tool[StartBackgroundCommandParams]:
    async def execute(
        params: StartBackgroundCommandParams,
        ctx: ToolContext,
    ) -> ToolOutput:
        session = await manager.start(params.command, params.cwd)
        return ToolOutput(
            content=(
                f"Started background command {session.session_id}\n"
                f"Command: {session.command}\n"
                f"Working directory: {session.cwd}\n"
                "Use list_background_commands, read_background_command, wait_background_command, "
                "or stop_background_command to manage it."
            ),
            metadata={
                "session_id": session.session_id,
                "command": session.command,
                "cwd": str(session.cwd),
                "state": session.state,
            },
            raw=manager.get_session_raw_payload(session.session_id),
        )

    return Tool(
        name="start_background_command",
        description=START_BACKGROUND_COMMAND_DESCRIPTION,
        parameters=StartBackgroundCommandParams,
        execute=execute,
    )


def create_list_background_commands_tool(
    manager: BackgroundCommandManager,
) -> Tool[ListBackgroundCommandsParams]:
    async def execute(
        params: ListBackgroundCommandsParams,
        ctx: ToolContext,
    ) -> ToolOutput:
        sessions = manager.list_sessions(include_completed=params.include_completed)
        return ToolOutput(
            content=manager.format_session_list(
                include_completed=params.include_completed,
            ),
            metadata={
                "count": len(sessions),
                "running_count": manager.count_running(),
            },
            raw={
                "kind": "background_command",
                "sessions": [
                    {
                        "session_id": session.session_id,
                        "command": session.command,
                        "cwd": str(session.cwd),
                        "state": session.state,
                        "exit_code": session.exit_code,
                        "runtime_seconds": session.runtime_seconds(),
                    }
                    for session in sessions
                ],
            },
        )

    return Tool(
        name="list_background_commands",
        description=LIST_BACKGROUND_COMMANDS_DESCRIPTION,
        parameters=ListBackgroundCommandsParams,
        execute=execute,
    )


def create_read_background_command_tool(
    manager: BackgroundCommandManager,
) -> Tool[ReadBackgroundCommandParams]:
    async def execute(
        params: ReadBackgroundCommandParams,
        ctx: ToolContext,
    ) -> ToolOutput:
        session = manager.require_session(params.session_id)
        return ToolOutput(
            content=manager.format_session_output(
                session.session_id,
                max_output_chars=params.max_output_chars,
            ),
            metadata={
                "session_id": session.session_id,
                "state": session.state,
                "exit_code": session.exit_code,
            },
            raw=manager.get_session_raw_payload(
                session.session_id,
                after_event_index=params.after_event_index,
            ),
        )

    return Tool(
        name="read_background_command",
        description=READ_BACKGROUND_COMMAND_DESCRIPTION,
        parameters=ReadBackgroundCommandParams,
        execute=execute,
    )


def create_wait_background_command_tool(
    manager: BackgroundCommandManager,
) -> Tool[WaitBackgroundCommandParams]:
    async def execute(
        params: WaitBackgroundCommandParams,
        ctx: ToolContext,
    ) -> ToolOutput:
        session = await manager.wait(
            params.session_id,
            timeout_seconds=params.timeout_seconds,
        )
        return ToolOutput(
            content=manager.format_session_output(
                session.session_id,
                max_output_chars=params.max_output_chars,
            ),
            metadata={
                "session_id": session.session_id,
                "state": session.state,
                "exit_code": session.exit_code,
                "running": session.is_running(),
            },
            raw=manager.get_session_raw_payload(
                session.session_id,
                after_event_index=params.after_event_index,
            ),
        )

    return Tool(
        name="wait_background_command",
        description=WAIT_BACKGROUND_COMMAND_DESCRIPTION,
        parameters=WaitBackgroundCommandParams,
        execute=execute,
    )


def create_stop_background_command_tool(
    manager: BackgroundCommandManager,
) -> Tool[StopBackgroundCommandParams]:
    async def execute(
        params: StopBackgroundCommandParams,
        ctx: ToolContext,
    ) -> ToolOutput:
        session = await manager.stop(params.session_id)
        return ToolOutput(
            content=(
                f"Stopped background command {session.session_id}\n"
                f"Command: {session.command}\n"
                f"Exit code: {session.exit_code}"
            ),
            metadata={
                "session_id": session.session_id,
                "state": session.state,
                "exit_code": session.exit_code,
            },
            raw=manager.get_session_raw_payload(session.session_id),
        )

    return Tool(
        name="stop_background_command",
        description=STOP_BACKGROUND_COMMAND_DESCRIPTION,
        parameters=StopBackgroundCommandParams,
        execute=execute,
    )


def create_send_background_command_input_tool(
    manager: BackgroundCommandManager,
) -> Tool[SendBackgroundCommandInputParams]:
    async def execute(
        params: SendBackgroundCommandInputParams,
        ctx: ToolContext,
    ) -> ToolOutput:
        session = await manager.write_stdin(
            params.session_id,
            params.text,
            append_newline=params.append_newline,
        )
        return ToolOutput(
            content=(
                f"Sent input to background command {session.session_id}\n"
                f"Command: {session.command}"
            ),
            metadata={
                "session_id": session.session_id,
                "state": session.state,
            },
            raw=manager.get_session_raw_payload(
                session.session_id,
                append_newline=params.append_newline,
            ),
        )

    return Tool(
        name="send_background_command_input",
        description=SEND_BACKGROUND_COMMAND_INPUT_DESCRIPTION,
        parameters=SendBackgroundCommandInputParams,
        execute=execute,
    )


START_BACKGROUND_COMMAND_DESCRIPTION = """Start a long-running local command in the background.

Use this tool when:
- A build, dev server, watcher, or test run may take a while
- You want to continue helping the user without waiting for the command to finish
- The user asks you to keep an eye on a task while doing something else

After starting it, use the other background command tools to inspect output, wait, or stop it."""


LIST_BACKGROUND_COMMANDS_DESCRIPTION = """List background commands started in this CLI session.

Use this tool to:
- See which long-running tasks are still active
- Check session ids before reading output or stopping a task
- Confirm whether a background command already finished"""


READ_BACKGROUND_COMMAND_DESCRIPTION = """Read the latest output from one background command.

Use this tool to:
- Inspect recent stdout and stderr from a running task
- Check whether a background job is producing errors
- Review logs without blocking on completion"""


WAIT_BACKGROUND_COMMAND_DESCRIPTION = """Wait for one background command to finish or for a timeout window to expire.

Use this tool to:
- Pause briefly while monitoring a background task
- Check whether a build or test run has completed
- Fetch final output once the task exits"""


STOP_BACKGROUND_COMMAND_DESCRIPTION = """Stop one running background command.

Use this tool to:
- Cancel a stale build, watcher, or dev server
- Clean up a background process before starting another one
- End work the user no longer needs"""


SEND_BACKGROUND_COMMAND_INPUT_DESCRIPTION = """Send stdin text to a running background command.

Use this tool to:
- Interact with a process that prompts for input
- Send a newline or simple command to a background session
- Continue a process started with start_background_command"""
