from __future__ import annotations

import os
import signal
import subprocess
import threading
import time
from pathlib import Path
from typing import Callable


class ToolSandbox:
    """Applies simple working directory and timeout restrictions."""

    def __init__(
        self,
        allowed_root: str | Path,
        timeout_ms: int,
        command_timeout_ms: int | None = None,
        idle_timeout_ms: int | None = None,
    ) -> None:
        self._allowed_root = Path(allowed_root).resolve()
        self._timeout_seconds = max(timeout_ms / 1000, 0.001)
        self._command_timeout_seconds = max((command_timeout_ms or timeout_ms) / 1000, 0.001)
        self._idle_timeout_seconds = max((idle_timeout_ms or timeout_ms) / 1000, 0.001)

    def ensure_path_allowed(self, path: str | Path) -> Path:
        resolved = Path(path).expanduser().resolve()
        # Permission-first model: path is allowed if it exists
        # Permission service will handle write/shell restrictions
        return resolved

    def run_shell(
        self,
        command: str,
        cwd: str | Path | None = None,
        stdin_data: str | None = None,
        output_callback: Callable[[str, str], None] | None = None,
    ) -> subprocess.CompletedProcess[str]:
        working_directory = self.ensure_path_allowed(cwd or self._allowed_root)
        env = {"PATH": os.environ.get("PATH", "")}
        if output_callback is None:
            return subprocess.run(
                command,
                shell=True,
                cwd=str(working_directory),
                env=env,
                text=True,
                capture_output=True,
                input=stdin_data,
                timeout=self._command_timeout_seconds,
                check=False,
            )

        process = subprocess.Popen(
            command,
            shell=True,
            cwd=str(working_directory),
            env=env,
            text=True,
            stdin=subprocess.PIPE if stdin_data is not None else None,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            start_new_session=True,
        )

        stdout_chunks: list[str] = []
        stderr_chunks: list[str] = []
        output_lock = threading.Lock()
        last_output_at = time.monotonic()

        if stdin_data is not None and process.stdin is not None:
            process.stdin.write(stdin_data)
            process.stdin.close()

        def read_stream(stream_name: str) -> None:
            stream = process.stdout if stream_name == "stdout" else process.stderr
            if stream is None:
                return
            chunks = stdout_chunks if stream_name == "stdout" else stderr_chunks
            try:
                for line in iter(stream.readline, ""):
                    if not line:
                        break
                    chunks.append(line)
                    nonlocal last_output_at
                    with output_lock:
                        last_output_at = time.monotonic()
                    output_callback(stream_name, line)
            finally:
                stream.close()

        stdout_thread = threading.Thread(target=read_stream, args=("stdout",), daemon=True)
        stderr_thread = threading.Thread(target=read_stream, args=("stderr",), daemon=True)
        stdout_thread.start()
        stderr_thread.start()

        timed_out = False
        timeout_reason: str | None = None
        started_at = time.monotonic()
        return_code: int | None = None
        while return_code is None:
            return_code = process.poll()
            if return_code is not None:
                break

            now = time.monotonic()
            with output_lock:
                idle_for = now - last_output_at
            if now - started_at > self._command_timeout_seconds:
                timed_out = True
                timeout_reason = f"Command timed out after {self._command_timeout_seconds:.0f}s"
                break
            if idle_for > self._idle_timeout_seconds:
                timed_out = True
                timeout_reason = f"Command produced no output for {self._idle_timeout_seconds:.0f}s"
                break
            time.sleep(0.1)

        if timed_out:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            except PermissionError:
                process.kill()
            return_code = process.wait()
            timeout_message = f"{timeout_reason}\n"
            stderr_chunks.append(timeout_message)
            output_callback("stderr", timeout_message)

        stdout_thread.join(timeout=1)
        stderr_thread.join(timeout=1)
        return subprocess.CompletedProcess(
            args=command,
            returncode=return_code if not timed_out else -9,
            stdout="".join(stdout_chunks),
            stderr="".join(stderr_chunks),
        )