from __future__ import annotations

import asyncio
import logging
import os
import platform
import signal
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Awaitable, Callable

logger = logging.getLogger(__name__)

# Resolve once at import time so child processes can never write here
_EPOCHLER_SOURCE_ROOT = str(Path(__file__).resolve().parents[2])


def validate_script_safety(code: str, allowed_workspace: Path) -> None:
    """Raise ValueError if the script contains obvious writes outside the workspace."""
    blocked = Path(_EPOCHLER_SOURCE_ROOT).resolve()
    ws = allowed_workspace.resolve()
    import ast

    try:
        tree = ast.parse(code)
    except SyntaxError:
        return

    for node in ast.walk(tree):
        if isinstance(node, ast.Constant) and isinstance(node.value, str):
            try:
                p = Path(node.value).resolve()
                if str(p).startswith(str(blocked)) and not str(p).startswith(str(ws)):
                    raise ValueError(
                        f"Script references Epochler source path: {node.value!r}"
                    )
            except (OSError, ValueError):
                pass


@dataclass
class RunResult:
    exit_code: int = -1
    stdout: str = ""
    stderr: str = ""
    elapsed_seconds: float = 0.0
    timed_out: bool = False
    peak_memory_mb: float | None = None


def _kill_process(proc: asyncio.subprocess.Process) -> None:
    """Kill a subprocess and its entire process group."""
    try:
        if os.name != "nt":
            os.killpg(proc.pid, signal.SIGKILL)
        else:
            proc.kill()
    except (ProcessLookupError, OSError):
        pass


async def run_script(
    script_path: Path,
    working_dir: Path,
    timeout_seconds: int = 600,
    env_override: dict[str, str] | None = None,
    python_executable: str = "python",
    memory_limit_mb: int | None = None,
    on_stdout: Callable[[str], Awaitable[None]] | None = None,
    on_stderr: Callable[[str], Awaitable[None]] | None = None,
) -> RunResult:
    """Execute a Python script in a subprocess with streaming and timeout."""
    env = {
        key: value
        for key, value in os.environ.items()
        if key in {"PATH", "HOME", "TMPDIR", "TEMP", "TMP", "LANG", "LC_ALL"}
    }
    env["EPOCHLER_BLOCKED_PATHS"] = _EPOCHLER_SOURCE_ROOT
    # Ensure PYTHONPATH never includes the Epochler source tree
    if "PYTHONPATH" in env:
        filtered = [
            p for p in env["PYTHONPATH"].split(os.pathsep)
            if not p.startswith(_EPOCHLER_SOURCE_ROOT)
        ]
        env["PYTHONPATH"] = os.pathsep.join(filtered) if filtered else ""
    if env_override:
        env.update(env_override)

    result = RunResult()
    stdout_lines: list[str] = []
    stderr_lines: list[str] = []
    start = time.monotonic()

    def _preexec() -> None:
        os.setsid()
        if memory_limit_mb is None:
            return
        try:
            import resource

            limit_bytes = int(memory_limit_mb) * 1024 * 1024
            resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes))
        except Exception:
            pass

    async def _read_stream(stream, lines_buf, callback):
        while True:
            line = await stream.readline()
            if not line:
                break
            decoded = line.decode("utf-8", errors="replace")
            lines_buf.append(decoded)
            if callback:
                try:
                    await callback(decoded)
                except Exception:
                    pass

    stdout_task: asyncio.Task | None = None
    stderr_task: asyncio.Task | None = None

    try:
        proc = await asyncio.create_subprocess_exec(
            python_executable, str(script_path),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(working_dir),
            env=env,
            preexec_fn=_preexec if os.name != "nt" else None,
        )

        stdout_task = asyncio.create_task(_read_stream(proc.stdout, stdout_lines, on_stdout))
        stderr_task = asyncio.create_task(_read_stream(proc.stderr, stderr_lines, on_stderr))

        try:
            await asyncio.wait_for(
                asyncio.gather(proc.wait(), stdout_task, stderr_task),
                timeout=timeout_seconds,
            )
        except asyncio.TimeoutError:
            result.timed_out = True
            _kill_process(proc)
            await asyncio.gather(stdout_task, stderr_task, return_exceptions=True)
            try:
                await asyncio.wait_for(proc.wait(), timeout=5)
            except asyncio.TimeoutError:
                pass

        result.exit_code = proc.returncode if proc.returncode is not None else -1
        result.stdout = "".join(stdout_lines)
        result.stderr = "".join(stderr_lines)

        try:
            import resource
            rusage = resource.getrusage(resource.RUSAGE_CHILDREN)
            ru_maxrss = rusage.ru_maxrss
            if platform.system() == "Darwin":
                result.peak_memory_mb = ru_maxrss / (1024 * 1024)
            else:
                result.peak_memory_mb = ru_maxrss / 1024
        except Exception:
            pass

    except asyncio.CancelledError:
        if stdout_task:
            stdout_task.cancel()
        if stderr_task:
            stderr_task.cancel()
        try:
            _kill_process(proc)  # type: ignore[possibly-undefined]
            await proc.wait()  # type: ignore[possibly-undefined]
        except Exception:
            pass
        result.exit_code = -1
        result.stdout = "".join(stdout_lines)
        result.stderr = "".join(stderr_lines)
        raise

    except Exception as e:
        logger.exception("Subprocess execution error")
        result.stderr = str(e)
        result.exit_code = -1

    result.elapsed_seconds = time.monotonic() - start
    return result
