"""Container image builder - Modal-style image definition."""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class Image:
    """Defines a container image for deployment.

    Usage:
        image = Image.from_registry("pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime")
        image = image.pip_install("transformers", "accelerate")
        image = image.run_commands("apt-get update && apt-get install -y ffmpeg")

        # Or from a Dockerfile
        image = Image.from_dockerfile("./Dockerfile")
    """

    base: str = "python:3.11-slim"
    pip_packages: list[str] = field(default_factory=list)
    commands: list[str] = field(default_factory=list)
    env_vars: dict[str, str] = field(default_factory=dict)
    dockerfile_path: str | None = None

    @classmethod
    def from_registry(cls, image: str) -> Image:
        """Create an image from a Docker registry image."""
        return cls(base=image)

    @classmethod
    def from_dockerfile(cls, path: str) -> Image:
        """Create an image from a Dockerfile."""
        return cls(dockerfile_path=path)

    def pip_install(self, *packages: str) -> Image:
        """Add pip packages to install."""
        return Image(
            base=self.base,
            pip_packages=[*self.pip_packages, *packages],
            commands=self.commands.copy(),
            env_vars=self.env_vars.copy(),
            dockerfile_path=self.dockerfile_path,
        )

    def run_commands(self, *cmds: str) -> Image:
        """Add shell commands to run during build."""
        return Image(
            base=self.base,
            pip_packages=self.pip_packages.copy(),
            commands=[*self.commands, *cmds],
            env_vars=self.env_vars.copy(),
            dockerfile_path=self.dockerfile_path,
        )

    def env(self, **kwargs: str) -> Image:
        """Set environment variables."""
        new_env = {**self.env_vars, **kwargs}
        return Image(
            base=self.base,
            pip_packages=self.pip_packages.copy(),
            commands=self.commands.copy(),
            env_vars=new_env,
            dockerfile_path=self.dockerfile_path,
        )

    def to_dockerfile(self) -> str:
        """Generate a Dockerfile from this image spec."""
        if self.dockerfile_path:
            with open(self.dockerfile_path) as f:
                return f.read()

        lines = [f"FROM {self.base}"]

        for key, val in self.env_vars.items():
            lines.append(f"ENV {key}={val}")

        if self.pip_packages:
            pkgs = " ".join(self.pip_packages)
            lines.append(f"RUN pip install --no-cache-dir {pkgs}")

        for cmd in self.commands:
            lines.append(f"RUN {cmd}")

        lines.append("WORKDIR /app")
        lines.append("COPY . /app")

        return "\n".join(lines)
