"""Types for the cloud CUDA compiler."""
from dataclasses import dataclass, field
from enum import Enum
from typing import Literal
class OutputFormat(str, Enum):
    """Supported output formats."""
    PTX = "ptx"
    SASS = "sass"
# Valid SM architectures (CUDA compute capabilities)
VALID_ARCHITECTURES = frozenset(
    {
        # Ampere
        "sm_80",
        "sm_86",
        "sm_87",
        # Ada Lovelace
        "sm_89",
        # Hopper
        "sm_90",
        "sm_90a",
        # Blackwell
        "sm_100",
        "sm_100a",
        "sm_101",
        "sm_101a",
        "sm_120",
        "sm_120a",
    }
)
@dataclass(frozen=True)
class CompileRequest:
    """Request to compile CUDA code.
    Attributes:
        files: Mapping of filename to file content. At least one .cu file required.
        arch: Target GPU architecture (e.g., "sm_90a" for Hopper).
        flags: Additional nvcc flags (e.g., ["-O3", "--maxrregcount=64"]).
        output: Requested output formats (ptx and/or sass).
    """
    files: dict[str, str]
    arch: str = "sm_90a"
    flags: tuple[str, ...] = field(default_factory=tuple)
    output: tuple[OutputFormat, ...] = field(
        default_factory=lambda: (OutputFormat.PTX, OutputFormat.SASS)
    )
    def __post_init__(self) -> None:
        """Validate the request."""
        # Validate files
        if not self.files:
            raise ValueError("At least one file is required")

        cu_files = [f for f in self.files if f.endswith(".cu")]
        if not cu_files:
            raise ValueError("At least one .cu file is required")
        # Validate architecture
        if self.arch not in VALID_ARCHITECTURES:
            raise ValueError(
                f"Invalid architecture '{self.arch}'. "
                f"Valid architectures: {sorted(VALID_ARCHITECTURES)}"
            )
        # Validate output formats
        if not self.output:
            raise ValueError("At least one output format is required")
    @property
    def main_cu_file(self) -> str:
        """Get the main .cu file (first one found)."""
        for filename in self.files:
            if filename.endswith(".cu"):
                return filename
        # This should never happen due to __post_init__ validation
        raise ValueError("No .cu file found")  # pragma: no cover
@dataclass(frozen=True)
class CompileResponse:
    """Response from CUDA compilation.
    Attributes:
        success: Whether compilation succeeded.
        ptx: Generated PTX code (if requested and successful).
        sass: Generated SASS code (if requested and successful).
        stderr: Compiler warnings or errors.
        compilation_time_ms: Time taken to compile in milliseconds.
    """
    success: bool
    ptx: str | None = None
    sass: str | None = None
    stderr: str = ""
    compilation_time_ms: int = 0
    @classmethod
    def error(cls, message: str, compilation_time_ms: int = 0) -> "CompileResponse":
        """Create an error response."""
        return cls(
            success=False,
            stderr=message,
            compilation_time_ms=compilation_time_ms,
        )
# Type alias for the output format literal
OutputFormatLiteral = Literal["ptx", "sass"]
