#!/usr/bin/env python3

import sys
import logging
import subprocess
import time
from pathlib import Path
from typing import Dict, Any, Optional, List
import yaml
import toml
import argparse
from dataclasses import dataclass
from enum import Enum, auto
import random
import string

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class ShellType(Enum):
    BASH = auto()
    SH = auto()
    UNKNOWN = auto()


@dataclass
class ContainerConfig:
    path: Path
    service: str
    image: str
    shell: ShellType = ShellType.UNKNOWN


def validate_path(path: Path) -> None:
    """Validate the provided path."""
    if not path:
        raise ValueError("Path cannot be empty")

    if not path.is_absolute():
        path = path.absolute()

    if not path.parent.exists():
        raise FileNotFoundError(f"Parent directory {path.parent} does not exist")


def validate_image(image: str) -> None:
    """Validate the Docker image name."""
    if not image:
        raise ValueError("Docker image cannot be empty")

    # Basic validation for image format
    if not any(char in image for char in [":", "@"]):
        logger.warning("No tag or digest specified for the image. Using 'latest' tag")


def check_image_exists(image: str) -> bool:
    """Check if the Docker image exists locally."""
    try:
        subprocess.run(
            ["docker", "image", "inspect", image], check=True, capture_output=True
        )
        return True
    except subprocess.CalledProcessError:
        return False


def pull_image(image: str) -> None:
    """Pull the Docker image if it doesn't exist locally."""
    if not check_image_exists(image):
        logger.info(f"Pulling image {image}...")
        try:
            subprocess.run(["docker", "pull", image], check=True)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to pull image {image}: {e}")
            raise


def detect_shell(image: str) -> ShellType:
    """Detect available shell in the image."""
    try:
        # Try bash first
        subprocess.run(
            ["docker", "run", "--rm", image, "which", "bash"],
            check=True,
            capture_output=True,
        )
        return ShellType.BASH
    except subprocess.CalledProcessError:
        try:
            # Fallback to sh
            subprocess.run(
                ["docker", "run", "--rm", image, "which", "sh"],
                check=True,
                capture_output=True,
            )
            return ShellType.SH
        except subprocess.CalledProcessError:
            return ShellType.UNKNOWN


def create_docker_compose_config(
    workspace_path: Path, image: str, service_name: str
) -> Dict[str, Any]:
    """Create the Docker Compose configuration."""
    return {
        "services": {
            service_name: {
                "image": image,
                "command": "tail -f /dev/null",  # Keep container running with a shell
                "volumes": [f"{workspace_path.absolute()}:/workspace"],
                "working_dir": "/workspace",
                "stdin_open": True,
                "tty": True,
                "environment": {
                    "TERM": "xterm-256color",
                },
                "security_opt": ["seccomp=unconfined"],
            }
        }
    }


def create_toml_config(
    container_path: Path, service_name: str, image: str, shell: ShellType
) -> Dict[str, Any]:
    """Create the TOML configuration."""
    return {
        "container": {
            "path": str(container_path.absolute()),
            "service": service_name,
            "image": image,
            "shell": shell.name,
        }
    }


def write_config(path: Path, config: Dict[str, Any], config_type: str) -> None:
    """Write configuration to a file."""
    try:
        if config_type == "yaml":
            path.write_text(yaml.dump(config, sort_keys=False))
        elif config_type == "toml":
            path.write_text(toml.dumps(config))
        logger.info(f"Created {config_type.upper()} file at: {path}")
    except Exception as e:
        logger.error(f"Failed to write {config_type.upper()} file: {e}")
        raise


def run_docker_compose_command(
    container_path: Path, command: str, detach: bool = False
) -> None:
    """Run a Docker Compose command with retries."""
    max_retries = 3
    retry_delay = 1

    for attempt in range(max_retries):
        try:
            cmd = [
                "docker",
                "compose",
                "-f",
                str(container_path / "docker-compose.yml"),
            ]
            cmd.extend(command.split())
            if detach:
                cmd.append("-d")

            subprocess.run(cmd, check=True)
            return
        except subprocess.CalledProcessError as e:
            if attempt == max_retries - 1:
                logger.error(
                    f"Failed to execute Docker Compose command after {max_retries} attempts: {e}"
                )
                raise
            logger.warning(f"Command failed, retrying in {retry_delay} seconds...")
            time.sleep(retry_delay)
            retry_delay *= 2


def get_container_config() -> ContainerConfig:
    """Get the container configuration from deu.toml."""
    try:
        config = toml.loads((Path.cwd() / "deu.toml").read_text())
        container_config = config["container"]
        return ContainerConfig(
            path=Path(container_config["path"]),
            service=container_config["service"],
            image=container_config["image"],
            shell=ShellType[container_config["shell"]],
        )
    except Exception as e:
        logger.error(f"Failed to read container configuration: {e}")
        sys.exit(1)


def handle_activate() -> None:
    """Activate the container shell."""
    config = get_container_config()
    print(f"Trying to activate container shell for service '{config.service}'...")

    # Check if container is running, start it if not
    try:
        ps_output = subprocess.run(
            ["docker", "compose", "-f", str(config.path / "docker-compose.yml"), "ps"],
            check=True,
            capture_output=True,
            text=True,
        )
        if config.service not in ps_output.stdout:
            logger.info(f"Container '{config.service}' is not running. Starting it...")
            run_docker_compose_command(config.path, "up -d")
    except subprocess.CalledProcessError:
        logger.info(f"Container '{config.service}' is not running. Starting it...")
        run_docker_compose_command(config.path, "up -d")

    # Use appropriate shell based on configuration
    shell = "bash" if config.shell == ShellType.BASH else "sh"
    run_docker_compose_command(config.path, f"exec {config.service} {shell}")


def handle_background() -> None:
    """Start the container in background mode."""
    container_path = get_container_config().path
    run_docker_compose_command(container_path, "up", detach=True)


def handle_logs() -> None:
    """Show container logs."""
    container_path = get_container_config().path
    run_docker_compose_command(container_path, "logs -f")


def handle_stop() -> None:
    """Stop the container."""
    container_path = get_container_config().path
    run_docker_compose_command(container_path, "stop")


def handle_rm() -> None:
    """Remove the container."""
    container_path = get_container_config().path
    run_docker_compose_command(container_path, "down")


def create_dev_container(path: str, image: str, service_name: str) -> None:
    """Create a development container configuration."""
    try:
        container_path = Path(path)
        validate_path(container_path)
        validate_image(image)

        # Ensure image exists and detect shell
        pull_image(image)
        shell_type = detect_shell(image)
        if shell_type == ShellType.UNKNOWN:
            logger.error(
                "No shell (bash/sh) found in the image. Please use an image with a shell."
            )
            sys.exit(1)

        # Create the container directory if it doesn't exist
        container_path.mkdir(parents=True, exist_ok=True)

        # Create configurations
        workspace_path = container_path.parent
        compose_config = create_docker_compose_config(
            workspace_path, image, service_name
        )
        toml_config = create_toml_config(
            container_path, service_name, image, shell_type
        )

        # Write configurations
        write_config(container_path / "docker-compose.yml", compose_config, "yaml")
        write_config(Path.cwd() / "deu.toml", toml_config, "toml")

    except Exception as e:
        logger.error(f"Failed to create development container: {e}")
        sys.exit(1)


def generate_service_name() -> str:
    """Generate a random, friendly service name."""
    adjectives = [
        "happy",
        "sleepy",
        "daisy",
        "sunny",
        "lucky",
        "fuzzy",
        "cozy",
        "merry",
    ]
    nouns = ["dog", "cat", "bear", "fox", "bird", "fish", "tree", "star"]
    random_number = random.randint(1000, 9999)
    return f"{random.choice(adjectives)}_{random.choice(nouns)}_{random_number}"


def main() -> None:
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(
        description="Development container utility",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # Init command
    init_parser = subparsers.add_parser("init", help="Initialize a new container")
    init_parser.add_argument("path", help="Path to the .container directory")
    init_parser.add_argument("--image", required=True, help="Docker image to use")
    init_parser.add_argument(
        "--service",
        default=generate_service_name(),
        help="Name of the container service",
    )

    # Other commands
    subparsers.add_parser("activate", help="Activate container shell")
    subparsers.add_parser("background", help="Start container in background")
    subparsers.add_parser("logs", help="Show container logs")
    subparsers.add_parser("stop", help="Stop container")
    subparsers.add_parser("rm", help="Remove container")

    args = parser.parse_args()

    if args.command == "init":
        create_dev_container(args.path, args.image, args.service)
    elif args.command == "activate":
        handle_activate()
    elif args.command == "background":
        handle_background()
    elif args.command == "logs":
        handle_logs()
    elif args.command == "stop":
        handle_stop()
    elif args.command == "rm":
        handle_rm()
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()