#!/usr/bin/env python3

import hashlib
import logging
import os
import subprocess
import sys
import time
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stderr)
logger = logging.getLogger(__name__)

USAGE = """\
Usage:
  agent-container opencode [args...]   Run OpenCode in the container
  agent-container claude [args...]     Run Claude Code in the container
  agent-container update               Rebuild image with latest versions
  agent-container force-cleanup        Remove all containers, image, and data
"""


class AgentContainer:
    IMAGE = "agent-container:latest"
    CONTAINER_PREFIX = "ac-"
    DATA_DIR_NAME = "agent-container"
    PROJECT_ID_LEN = 12
    START_TIMEOUT = 3.0

    def __init__(self):
        self.project_path = Path.cwd().resolve()
        self.project_id = self._project_id(self.project_path)

        self.host_username = os.environ.get("USER", "dev")
        self.host_uid = os.getuid()
        self.host_gid = os.getgid()

        self.container_name = (
            f"{self.CONTAINER_PREFIX}{self.project_path.name}-{self.project_id}"
        )
        self.docker_context_dir = Path(__file__).resolve().parent

    def _get_xdg_data_home(self) -> Path:
        """Get XDG_DATA_HOME with fallback to ~/.local/share"""
        xdg_data = os.environ.get("XDG_DATA_HOME")
        if xdg_data:
            return Path(xdg_data)
        return Path.home() / ".local" / "share"

    @property
    def container_home_path(self) -> Path:
        """Container home directory in XDG_DATA_HOME"""
        return self._get_xdg_data_home() / self.DATA_DIR_NAME / "container-home"

    # =========================
    # Subcommand dispatch
    # =========================

    def dispatch(self, argv: list[str]) -> None:
        if not argv:
            print(USAGE, file=sys.stderr)
            sys.exit(1)

        command = argv[0]
        args = argv[1:]

        if command == "opencode":
            self._run_tool("opencode", args, env_vars=[])
        elif command == "claude":
            self._run_tool(
                "claude",
                args,
                env_vars=[
                    "ANTHROPIC_API_KEY",
                    "ANTHROPIC_BASE_URL",
                ],
                extra_env={"DISABLE_AUTOUPDATER": "1"},
            )
        elif command == "update":
            self.update()
        elif command == "force-cleanup":
            self.force_cleanup()
        else:
            logger.error(f"Unknown command: {command}")
            print(USAGE, file=sys.stderr)
            sys.exit(1)

    # =========================
    # Run a tool inside the container
    # =========================

    def _run_tool(
        self,
        tool: str,
        args: list[str],
        env_vars: list[str],
        extra_env: dict[str, str] | None = None,
    ) -> None:
        self.container_home_path.mkdir(parents=True, exist_ok=True)

        if self.container_exists() and self.container_running():
            logger.error(
                f"Project '{self.project_path.name}' already has a running agent container."
            )
            logger.error(f"Container name: {self.container_name}")
            logger.error(
                "Wait for the current instance to finish or manually stop it with:"
            )
            logger.error(f"  docker stop {self.container_name}")
            sys.exit(1)

        # Pre-create project directory structure to prevent root-owned directories
        try:
            relative_path = self.project_path.relative_to(Path.home())
            (self.container_home_path / relative_path).mkdir(
                parents=True, exist_ok=True
            )
        except ValueError:
            pass

        if not self.image_exists():
            self.build_image()

        if not self.container_exists():
            self.create_container()

        if not self.start_container():
            logger.warning("Recreating container due to failed start")
            self.remove_container()
            self.create_container()
            if not self.start_container():
                raise RuntimeError("Container failed to start after recreation")

        try:
            self._exec_tool(tool, args, env_vars, extra_env or {})
        finally:
            self.stop_container()

    def _exec_tool(
        self, tool: str, args: list[str], env_vars: list[str], extra_env: dict[str, str]
    ) -> None:
        env_args: list[str] = []

        # Pass through host environment variables if set
        for var in env_vars:
            val = os.environ.get(var)
            if val:
                env_args += ["-e", f"{var}={val}"]

        # Set additional environment variables
        for var, val in extra_env.items():
            env_args += ["-e", f"{var}={val}"]

        subprocess.run(
            [
                "docker",
                "exec",
                "-it",
                *env_args,
                "-w",
                str(self.project_path),
                self.container_name,
                tool,
                *args,
            ],
            check=True,
        )

    # =========================
    # Image handling
    # =========================

    def image_exists(self) -> bool:
        return (
            subprocess.run(
                ["docker", "image", "inspect", self.IMAGE],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            ).returncode
            == 0
        )

    def build_image(self, no_cache: bool = False) -> None:
        logger.info(
            f"Building image '{self.IMAGE}' with user {self.host_username} ({self.host_uid}:{self.host_gid})"
        )
        cmd = [
            "docker",
            "build",
            "--build-arg",
            f"USERNAME={self.host_username}",
            "--build-arg",
            f"UID={self.host_uid}",
            "--build-arg",
            f"GID={self.host_gid}",
            "-t",
            self.IMAGE,
        ]
        if no_cache:
            cmd.append("--no-cache")
        cmd.append(str(self.docker_context_dir))
        self._run_cmd(cmd)

    # =========================
    # Container lifecycle
    # =========================

    def create_container(self) -> None:
        logger.info(f"Creating container '{self.container_name}'")
        self._run_cmd(
            [
                "docker",
                "create",
                "--name",
                self.container_name,
                "--user",
                f"{self.host_uid}:{self.host_gid}",
                "--volume",
                f"{self.project_path}:{self.project_path}",
                "--volume",
                f"{self.container_home_path}:/home/{self.host_username}",
                self.IMAGE,
                "sh",
                "-c",
                "sleep infinity",
            ]
        )

    def start_container(self) -> bool:
        subprocess.run(
            ["docker", "start", self.container_name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        return self.wait_until_running()

    def stop_container(self) -> None:
        if self.container_running():
            self._run_cmd(["docker", "stop", "-t", "0", self.container_name])

    def remove_container(self) -> None:
        subprocess.run(
            ["docker", "rm", "-f", self.container_name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

    # =========================
    # Update & cleanup
    # =========================

    def update(self) -> None:
        logger.info("Updating agent-container...")
        self._remove_all_containers()
        if self.image_exists():
            logger.info(f"Removing image '{self.IMAGE}'...")
            subprocess.run(
                ["docker", "rmi", self.IMAGE],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        logger.info("Rebuilding image with latest versions...")
        self.build_image(no_cache=True)
        logger.info("Update complete. Containers will be recreated on next run.")

    def force_cleanup(self) -> None:
        logger.info("Running force cleanup...")
        self._remove_all_containers()

        if self.image_exists():
            logger.info(f"Removing image '{self.IMAGE}'...")
            subprocess.run(
                ["docker", "rmi", self.IMAGE],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )

        data_dir = self._get_xdg_data_home() / self.DATA_DIR_NAME
        if data_dir.exists():
            logger.info(f"Removing {data_dir}...")
            import shutil

            shutil.rmtree(data_dir)

        logger.info("Force cleanup complete.")

    def _remove_all_containers(self) -> None:
        logger.info("Removing existing containers...")
        result = subprocess.run(
            [
                "docker",
                "ps",
                "-a",
                "--filter",
                f"name=^{self.CONTAINER_PREFIX}",
                "--format",
                "{{.Names}}",
            ],
            capture_output=True,
            text=True,
        )
        for name in result.stdout.strip().splitlines():
            if name:
                logger.info(f"  Removing container '{name}'")
                subprocess.run(
                    ["docker", "rm", "-f", name],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )

    # =========================
    # State checks
    # =========================

    def container_exists(self) -> bool:
        return (
            subprocess.run(
                ["docker", "inspect", self.container_name],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            ).returncode
            == 0
        )

    def container_running(self) -> bool:
        result = subprocess.run(
            ["docker", "inspect", "-f", "{{.State.Running}}", self.container_name],
            capture_output=True,
            text=True,
        )
        return result.returncode == 0 and result.stdout.strip() == "true"

    def wait_until_running(self) -> bool:
        deadline = time.time() + self.START_TIMEOUT
        while time.time() < deadline:
            if self.container_running():
                return True
            time.sleep(0.1)
        return False

    # =========================
    # Helpers
    # =========================

    @staticmethod
    def _project_id(path: Path) -> str:
        return hashlib.sha256(str(path).encode()).hexdigest()[:12]

    @staticmethod
    def _run_cmd(cmd: list[str]) -> None:
        subprocess.run(cmd, check=True)


def main() -> None:
    AgentContainer().dispatch(sys.argv[1:])


if __name__ == "__main__":
    main()