from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field, HttpUrl, field_validator


class Info(BaseModel):
    year: int
    version: str
    description: str
    contributor: str
    url: str
    # date_created is expected in ISO format; if parsing fails it remains a string.
    date_created: Union[datetime, str]

    @field_validator("date_created", mode="before")
    def parse_date_created(cls, v):
        if isinstance(v, str):
            try:
                return datetime.fromisoformat(v)
            except ValueError:
                # Return as-is if parsing fails
                return v
        return v


class License(BaseModel):
    id: int
    name: str
    url: Optional[str] = None


class Image(BaseModel):
    id: int
    width: int
    height: int
    file_name: str
    # Optional additional fields.
    license: Optional[int] = None
    flickr_url: Optional[str] = None
    coco_url: Optional[str] = None
    date_captured: Optional[Union[datetime, str]] = None

    @field_validator("date_captured", mode="before")
    def parse_date_captured(cls, v):
        if v is None:
            return v
        if isinstance(v, str):
            try:
                return datetime.fromisoformat(v)
            except ValueError:
                return v
        return v


class Annotation(BaseModel):
    id: int
    image_id: int
    category_id: int
    # segmentation in COCO can be a list of lists (for polygons) or a dict (for RLE)
    segmentation: Optional[Union[List[float], List[List[float]], dict]] = None
    area: float
    bbox: List[float]
    iscrowd: int

    @field_validator("bbox")
    def validate_bbox_length(cls, v):
        if len(v) != 4:
            raise ValueError("bbox must be a list of 4 numbers [x, y, width, height]")
        return v


class Category(BaseModel):
    id: int
    name: str
    supercategory: Optional[str] = None


# --- Define the top-level COCO dataset model ---
#
# Optional fields (like "info" and "licenses") are allowed to be missing.
# An extra field, e.g. "bbox_format", can also be supported if needed.


class CocoDataset(BaseModel):
    info: Optional[Info] = None
    licenses: Optional[List[License]] = None
    images: List[Image]
    annotations: List[Annotation]
    categories: List[Category]
    bbox_format: Optional[str] = None
    categories: List[Category]
    bbox_format: Optional[str] = None
