# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import asyncio
import os
import shutil
import sys
from collections.abc import Sequence
from typing import Any, Callable, ClassVar

from mirage.runtime.config import HomeConfig, RuntimeConfig
from mirage.runtime.python.base import PythonRuntime
from mirage.runtime.python.bootstrap import bootstrap
from mirage.runtime.python.flags import init_argv
from mirage.runtime.types import RunArgs, RunResult, RuntimeReach, ScriptSource

LOCAL_HOME_ENV = "MIRAGE_LOCAL_HOME"


class LocalRuntime(PythonRuntime):
    """Run Python code on a host interpreter as a subprocess.

    Each run spawns `<interpreter> -c <code>`; the code sees the host
    filesystem and environment, not the workspace mounts. Cancelling the
    run kills the subprocess, so a limit timeout reclaims it.

    The interpreter defaults to the one running mirage; point the
    config `home` (the yaml entry's ``config`` block ends up here) or
    the MIRAGE_LOCAL_HOME environment variable at another binary, e.g.
    a project venv whose packages the code needs.
    """

    name = "local"
    # Spawns the host interpreter: a real process with the user's own
    # filesystem and network, doors the workspace gate never sees.
    # This is the base default; declared here so the claim is explicit
    # at the one builtin runtime that voids a world's sandbox claim.
    reach: RuntimeReach = "process"

    config_cls: ClassVar[type[RuntimeConfig]] = HomeConfig
    config: HomeConfig

    def __init__(
            self,
            captures: Sequence[str] | None = None,
            config: HomeConfig | dict[str, Any] | None = None,
            script: Callable[..., Any] | ScriptSource | None = None) -> None:
        super().__init__(captures, config, script)
        chosen = self.config.home or os.environ.get(LOCAL_HOME_ENV)
        if chosen:
            resolved = shutil.which(chosen)
            if resolved is None:
                raise FileNotFoundError(
                    f"local python interpreter not found: {chosen!r} "
                    "(from the runtime entry's config `home` or "
                    f"{LOCAL_HOME_ENV})")
            self._python = resolved
        else:
            self._python = sys.executable

    async def run(self, args: RunArgs) -> RunResult:
        # Honoring the init switches is just handing them back to the
        # real interpreter, which is why this tier gets them exactly
        # right (sys.flags included) where an in-process engine cannot.
        proc = await asyncio.create_subprocess_exec(
            self._python,
            *init_argv(args.flags),
            "-c",
            bootstrap(args.code, args.prog),
            *args.args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env={
                **os.environ,
                **args.env
            },
        )
        try:
            stdout, stderr = await proc.communicate(input=args.stdin)
        except asyncio.CancelledError:
            proc.kill()
            await proc.wait()
            raise
        return RunResult(
            stdout=stdout,
            stderr=stderr or None,
            exit_code=proc.returncode if proc.returncode is not None else 1,
        )
