from __future__ import annotations

from dataclasses import dataclass, field


@dataclass(slots=True)
class Prediction:
    class_name: str = ""
    confidence: float = 0.0
    bbox: list[float] = field(default_factory=list)
    color: str | None = None
    keypoints: list[list[float]] | None = None
    mask: list[list[float]] | None = None
    obb: dict | None = None
    task: str | None = None
    probs: list[dict] | None = None

    def __repr__(self) -> str:
        return f"<Prediction {self.class_name} conf={self.confidence:.3f}>"


@dataclass(slots=True)
class PredictResult:
    predictions: list[Prediction] = field(default_factory=list)
    inference_ms: float = 0.0
    image_width: int | None = None
    image_height: int | None = None
    model_version_id: str | None = None

    @property
    def count(self) -> int:
        return len(self.predictions)


@dataclass(slots=True)
class BatchResult:
    results: list[PredictResult] = field(default_factory=list)
    total_ms: float = 0.0
    count: int = 0

    def __iter__(self):
        return iter(self.results)

    def __len__(self) -> int:
        return self.count


@dataclass(slots=True, frozen=True)
class ClassInfo:
    name: str
    color: str


@dataclass(slots=True)
class ModelClasses:
    model_version_id: str
    classes: list[ClassInfo] = field(default_factory=list)
    model_name: str | None = None
    architecture: str | None = None
    total: int = 0
    imgsz: int = 640

    def __contains__(self, name: str) -> bool:
        return any(c.name == name for c in self.classes)

    def names(self) -> list[str]:
        return [c.name for c in self.classes]


@dataclass(slots=True)
class ModelInfo:
    id: str
    name: str
    slug: str
    architecture: str | None = None
    owner_username: str | None = None

    @property
    def full_slug(self) -> str:
        if self.owner_username:
            return f"{self.owner_username}/{self.slug}"
        return self.slug


@dataclass(slots=True)
class ModelVersion:
    id: str
    version_tag: str
    weights_url: str | None = None
    framework: str = "pytorch"
    metrics: dict = field(default_factory=dict)
