"""hdx sandbox — run MCP tools inside containerized sandboxes."""

from __future__ import annotations

import asyncio
import json

import click

from hatchdx.utils import console
from hatchdx.utils.config import ConfigError, load_config


def _resolve_runtime(runtime_flag: str | None):
    """Resolve the container runtime from CLI flag, config, or auto-detect.

    Resolution order: --runtime CLI flag > container_runtime in config > auto-detect.
    Returns a CLIContainerRuntime instance.
    """
    from hatchdx.sandbox.runtime import detect_runtime

    if runtime_flag is not None:
        return detect_runtime(override=runtime_flag)

    try:
        config = load_config()
        if config.container_runtime:
            return detect_runtime(override=config.container_runtime)
    except ConfigError:
        pass

    return detect_runtime()


_runtime_option = click.option(
    "--runtime",
    "-r",
    "runtime_override",
    default=None,
    help="Container runtime binary (docker/podman/nerdctl or absolute path). Auto-detected if not set.",
)


@click.group("sandbox")
def sandbox_cmd():
    """Run MCP tools in a sandboxed container environment."""


@sandbox_cmd.command("run")
@click.argument("tool_name")
@click.argument("args_json", default="{}")
@click.option(
    "--policy",
    "-p",
    default=None,
    help="Policy preset name (strict/standard/permissive) or path to a YAML file",
)
@click.option(
    "--build",
    "force_build",
    is_flag=True,
    help="Force rebuild the container image",
)
@click.option(
    "--keep",
    is_flag=True,
    help="Keep the container after execution (don't remove it)",
)
@click.option(
    "--timeout",
    "-t",
    type=int,
    default=None,
    help="Override the policy timeout (seconds)",
)
@_runtime_option
def sandbox_run(tool_name, args_json, policy, force_build, keep, timeout, runtime_override):
    """Execute a single MCP tool call inside the sandbox.

    TOOL_NAME is the name of the tool to call.
    ARGS_JSON is the tool arguments as a JSON string (default: {}).
    """
    # Parse tool arguments
    try:
        arguments = json.loads(args_json)
    except json.JSONDecodeError as exc:
        console.error(f"Invalid JSON arguments: {exc}")
        raise SystemExit(1)

    if not isinstance(arguments, dict):
        console.error("Tool arguments must be a JSON object (dict).")
        raise SystemExit(1)

    # Load project config
    try:
        config = load_config()
    except ConfigError as exc:
        console.error(str(exc))
        raise SystemExit(1)

    console.info(f"Running {tool_name} in sandbox for {config.server_name}")

    # Resolve policy
    policy_path = None
    preset = None
    if policy is not None:
        if policy in ("strict", "standard", "permissive"):
            preset = policy
        else:
            policy_path = policy

    from hatchdx.sandbox.engine import SandboxEngine
    from hatchdx.sandbox.reporter import (
        report_cleanup,
        report_container_status,
        report_policy_summary,
        report_resource_stats,
        report_tool_result,
    )
    from hatchdx.sandbox.runtime import SandboxError

    # Resolve runtime: CLI flag > config > auto-detect
    try:
        runtime = _resolve_runtime(runtime_override)
    except SandboxError as exc:
        console.error(str(exc))
        raise SystemExit(1)

    try:
        engine = SandboxEngine(
            project_path=".",
            server_command=config.server_command,
            preset=preset,
            policy_path=policy_path,
            force_build=force_build,
            keep_container=keep,
            timeout_override=timeout,
            runtime=runtime,
        )
    except (SandboxError, Exception) as exc:
        console.error(str(exc))
        raise SystemExit(1)

    async def _run() -> int:
        try:
            await engine.start()
            # Report policy after start so enforcement status is known
            report_policy_summary(
                engine.policy, network_enforced=engine.network_enforced
            )
            if engine.container_id:
                report_container_status(
                    container_id=engine.container_id, started=True
                )

            result = await engine.run_tool(tool_name, arguments)
            report_tool_result(result)
            report_resource_stats(result.stats)

            return 1 if result.is_error else 0
        except SandboxError as exc:
            console.error(str(exc))
            return 1
        finally:
            container_id = engine.container_id or ""
            engine.cleanup()
            if container_id:
                report_cleanup(container_id=container_id, removed=not keep)

    try:
        exit_code = asyncio.run(_run())
    except KeyboardInterrupt:
        console.warning("Interrupted — cleaning up sandbox...")
        engine.cleanup()
        exit_code = 130

    raise SystemExit(exit_code)


@sandbox_cmd.command("shell")
@click.option(
    "--policy",
    "-p",
    default=None,
    help="Policy preset name or path to a YAML file",
)
@click.option(
    "--build",
    "force_build",
    is_flag=True,
    help="Force rebuild the container image",
)
@click.option(
    "--keep",
    is_flag=True,
    help="Keep the container after the shell exits",
)
@_runtime_option
def sandbox_shell(policy, force_build, keep, runtime_override):
    """Start an interactive shell inside the sandbox container."""
    # Load project config
    try:
        config = load_config()
    except ConfigError as exc:
        console.error(str(exc))
        raise SystemExit(1)

    console.info(f"Opening sandbox shell for {config.server_name}")

    # Resolve policy
    policy_path = None
    preset = None
    if policy is not None:
        if policy in ("strict", "standard", "permissive"):
            preset = policy
        else:
            policy_path = policy

    from hatchdx.sandbox.engine import SandboxEngine
    from hatchdx.sandbox.reporter import (
        report_cleanup,
        report_container_status,
        report_policy_summary,
    )
    from hatchdx.sandbox.runtime import SandboxError

    # Resolve runtime
    try:
        runtime = _resolve_runtime(runtime_override)
    except SandboxError as exc:
        console.error(str(exc))
        raise SystemExit(1)

    try:
        engine = SandboxEngine(
            project_path=".",
            server_command=config.server_command,
            preset=preset,
            policy_path=policy_path,
            force_build=force_build,
            keep_container=keep,
            runtime=runtime,
        )
    except (SandboxError, Exception) as exc:
        console.error(str(exc))
        raise SystemExit(1)

    async def _shell() -> int:
        try:
            await engine.start()
            report_policy_summary(
                engine.policy, network_enforced=engine.network_enforced
            )
            if engine.container_id:
                report_container_status(
                    container_id=engine.container_id, started=True
                )

            console.step("Starting shell (type 'exit' to quit)...")
            click.echo()

            exit_code, output = await engine.exec_shell()
            if output.strip():
                click.echo(output)
            return exit_code
        except SandboxError as exc:
            console.error(str(exc))
            return 1
        finally:
            container_id = engine.container_id or ""
            engine.cleanup()
            if container_id:
                report_cleanup(container_id=container_id, removed=not keep)

    try:
        exit_code = asyncio.run(_shell())
    except KeyboardInterrupt:
        console.warning("Interrupted — cleaning up sandbox...")
        engine.cleanup()
        exit_code = 130

    raise SystemExit(exit_code)


@sandbox_cmd.command("build")
@click.option(
    "--policy",
    "-p",
    default=None,
    help="Policy preset name or path to a YAML file",
)
@_runtime_option
def sandbox_build(policy, runtime_override):
    """Build the sandbox container image without running it."""
    # Load project config
    try:
        config = load_config()
    except ConfigError as exc:
        console.error(str(exc))
        raise SystemExit(1)

    console.info(f"Building sandbox image for {config.server_name}")

    # Resolve policy
    policy_path = None
    preset = None
    if policy is not None:
        if policy in ("strict", "standard", "permissive"):
            preset = policy
        else:
            policy_path = policy

    from hatchdx.sandbox.engine import SandboxEngine
    from hatchdx.sandbox.reporter import report_build_status, report_policy_summary
    from hatchdx.sandbox.runtime import SandboxError

    # Resolve runtime
    try:
        runtime = _resolve_runtime(runtime_override)
    except SandboxError as exc:
        console.error(str(exc))
        raise SystemExit(1)

    try:
        engine = SandboxEngine(
            project_path=".",
            preset=preset,
            policy_path=policy_path,
            force_build=True,
            runtime=runtime,
        )
    except (SandboxError, Exception) as exc:
        console.error(str(exc))
        raise SystemExit(1)

    # Build-only: enforcement happens at runtime, not at build time
    report_policy_summary(engine.policy, network_enforced=False)
    report_build_status(image_tag=engine.image_tag, building=True)

    async def _build() -> int:
        try:
            await engine.build_image()
            report_build_status(image_tag=engine.image_tag, built=True)
            return 0
        except SandboxError as exc:
            console.error(str(exc))
            return 1

    try:
        exit_code = asyncio.run(_build())
    except KeyboardInterrupt:
        console.warning("Build interrupted.")
        exit_code = 130

    raise SystemExit(exit_code)
