import asyncio
import time
from collections.abc import Awaitable, Callable

import structlog

from app.config import settings
from app.models import ToolResult

log = structlog.get_logger(__name__)

# Track running processes for cancellation: scan_id -> list[asyncio.subprocess.Process]
_running: dict[str, list[asyncio.subprocess.Process]] = {}


def register_process(scan_id: str, proc: asyncio.subprocess.Process) -> None:
    _running.setdefault(scan_id, []).append(proc)


def unregister_process(scan_id: str, proc: asyncio.subprocess.Process) -> None:
    procs = _running.get(scan_id, [])
    if proc in procs:
        procs.remove(proc)


async def cancel_scan_processes(scan_id: str) -> None:
    for proc in _running.get(scan_id, []):
        try:
            proc.kill()
        except ProcessLookupError:
            pass
    _running[scan_id] = []


async def run_tool(
    command: list[str],
    scan_id: str | None = None,
    cwd: str | None = None,
    timeout: int | None = None,
    on_line: Callable[[str], Awaitable[None]] | None = None,
) -> ToolResult:
    """Spawn process without shell; stream stdout line-by-line."""
    timeout = timeout or settings.tool_timeout_seconds
    start = time.monotonic()
    stdout_parts: list[str] = []
    stderr_parts: list[str] = []

    proc = await asyncio.create_subprocess_exec(
        *command,
        cwd=cwd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    if scan_id:
        register_process(scan_id, proc)

    async def read_stream(stream: asyncio.StreamReader, is_stderr: bool) -> None:
        while True:
            line = await stream.readline()
            if not line:
                break
            text = line.decode(errors="replace").rstrip("\n")
            if is_stderr:
                stderr_parts.append(text)
            else:
                stdout_parts.append(text)
                if on_line:
                    await on_line(text)

    try:
        await asyncio.wait_for(
            asyncio.gather(
                read_stream(proc.stdout, False),  # type: ignore[arg-type]
                read_stream(proc.stderr, True),  # type: ignore[arg-type]
            ),
            timeout=timeout,
        )
        exit_code = await proc.wait()
    except asyncio.TimeoutError:
        proc.kill()
        await proc.wait()
        exit_code = -1
        stderr_parts.append(f"timeout after {timeout}s")
        log.warning("tool_timeout", command=command, timeout=timeout)
    finally:
        if scan_id:
            unregister_process(scan_id, proc)

    duration = time.monotonic() - start
    return ToolResult(
        stdout="\n".join(stdout_parts),
        stderr="\n".join(stderr_parts),
        exit_code=exit_code,
        duration=duration,
    )