# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 Substratum Labs

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from roche_sandbox.daemon import detect_daemon, _find_bundled_binary, _spawn_daemon, _wait_for_daemon_ready
from roche_sandbox.sandbox import AsyncSandbox, Sandbox
from roche_sandbox.transport.cli import CliTransport
from roche_sandbox.transport.grpc import GrpcTransport
from roche_sandbox.intent import CodeIntent
from roche_sandbox.types import (
    Budget, DynamicPermissions, ExecOutput, SandboxConfig, SandboxInfo, SessionInfo,
)

if TYPE_CHECKING:
    from roche_sandbox.transport import Transport


class AsyncRoche:
    def __init__(
        self,
        *,
        mode: str = "auto",
        daemon_port: int | None = None,
        provider: str = "docker",
        binary: str = "roche",
        transport: Transport | None = None,
    ):
        self._provider = provider
        if transport is not None:
            self._transport = transport
        elif mode == "direct":
            self._transport = CliTransport(binary=binary)
        elif daemon_port is not None:
            self._transport = GrpcTransport(port=daemon_port)
        else:
            daemon = detect_daemon()
            if daemon is not None:
                self._transport = GrpcTransport(port=daemon["port"])
            else:
                roched_path = _find_bundled_binary("roched")
                if roched_path:
                    _spawn_daemon(roched_path)
                    if _wait_for_daemon_ready(timeout=3.0):
                        daemon = detect_daemon()
                        if daemon is not None:
                            self._transport = GrpcTransport(port=daemon["port"])
                        else:
                            self._transport = CliTransport(binary=binary)
                    else:
                        self._transport = CliTransport(binary=binary)
                else:
                    self._transport = CliTransport(binary=binary)

    @property
    def transport(self) -> Transport:
        return self._transport

    async def create(
        self,
        *,
        provider: str | None = None,
        image: str = "python:3.12-slim",
        memory: str | None = None,
        cpus: float | None = None,
        timeout_secs: int = 300,
        network: bool = False,
        writable: bool = False,
        env: dict[str, str] | None = None,
        mounts: list | None = None,
        kernel: str | None = None,
        rootfs: str | None = None,
        network_allowlist: list[str] | None = None,
        fs_paths: list[str] | None = None,
    ) -> AsyncSandbox:
        effective_provider = provider or self._provider
        config = SandboxConfig(
            provider=effective_provider,
            image=image,
            memory=memory,
            cpus=cpus,
            timeout_secs=timeout_secs,
            network=network,
            writable=writable,
            env=env or {},
            mounts=mounts or [],
            kernel=kernel,
            rootfs=rootfs,
            network_allowlist=network_allowlist or [],
            fs_paths=fs_paths or [],
        )
        sandbox_id = await self._transport.create(config, effective_provider)
        return AsyncSandbox(sandbox_id, effective_provider, self._transport)

    async def create_id(self, **kwargs) -> str:
        sb = await self.create(**kwargs)
        return sb.id

    async def exec(
        self, sandbox_id: str, command: list[str], timeout_secs: int | None = None, trace_level: str | None = None
    ) -> ExecOutput:
        return await self._transport.exec(sandbox_id, command, self._provider, timeout_secs, trace_level=trace_level)

    async def destroy(self, sandbox_id: str) -> None:
        await self._transport.destroy([sandbox_id], self._provider)

    async def list(self) -> list[SandboxInfo]:
        return await self._transport.list(self._provider)

    async def gc(self, dry_run: bool = False, all: bool = False) -> list[str]:
        return await self._transport.gc(self._provider, dry_run, all)

    async def create_session(
        self,
        sandbox_id: str,
        *,
        provider: str | None = None,
        permissions: DynamicPermissions | None = None,
        budget: Budget | None = None,
    ) -> str:
        return await self._transport.create_session(
            sandbox_id, provider or self._provider, permissions, budget
        )

    async def destroy_session(self, session_id: str) -> SessionInfo:
        return await self._transport.destroy_session(session_id)

    async def list_sessions(self) -> list[SessionInfo]:
        return await self._transport.list_sessions()

    async def update_permissions(self, session_id: str, **change) -> DynamicPermissions:
        return await self._transport.update_permissions(session_id, change)

    async def analyze_intent(self, code: str, language: str = "python") -> CodeIntent:
        return await self._transport.analyze_intent(code, language)


class Roche:
    def __init__(self, **kwargs):
        self._async = AsyncRoche(**kwargs)

    def create(self, **kwargs) -> Sandbox:
        sb = asyncio.run(self._async.create(**kwargs))
        return Sandbox(sb.id, sb.provider, self._async.transport)

    def create_id(self, **kwargs) -> str:
        return asyncio.run(self._async.create_id(**kwargs))

    def exec(
        self, sandbox_id: str, command: list[str], timeout_secs: int | None = None, trace_level: str | None = None
    ) -> ExecOutput:
        return asyncio.run(self._async.exec(sandbox_id, command, timeout_secs, trace_level=trace_level))

    def destroy(self, sandbox_id: str) -> None:
        asyncio.run(self._async.destroy(sandbox_id))

    def list(self) -> list[SandboxInfo]:
        return asyncio.run(self._async.list())

    def gc(self, dry_run: bool = False, all: bool = False) -> list[str]:
        return asyncio.run(self._async.gc(dry_run, all))

    def create_session(self, sandbox_id: str, **kwargs) -> str:
        return asyncio.run(self._async.create_session(sandbox_id, **kwargs))

    def destroy_session(self, session_id: str) -> SessionInfo:
        return asyncio.run(self._async.destroy_session(session_id))

    def list_sessions(self) -> list[SessionInfo]:
        return asyncio.run(self._async.list_sessions())

    def update_permissions(self, session_id: str, **change) -> DynamicPermissions:
        return asyncio.run(self._async.update_permissions(session_id, **change))

    def analyze_intent(self, code: str, language: str = "python") -> CodeIntent:
        return asyncio.run(self._async.analyze_intent(code, language))
