from pydantic import BaseModel
from typing import Optional, List
from enum import Enum


class PoolHealth(str, Enum):
    ONLINE = "ONLINE"
    DEGRADED = "DEGRADED"
    FAULTED = "FAULTED"
    OFFLINE = "OFFLINE"
    UNAVAIL = "UNAVAIL"


class Vdev(BaseModel):
    name: str
    state: str
    read: int = 0  # Read error count
    write: int = 0  # Write error count
    cksum: int = 0  # Checksum error count
    children: List["Vdev"] = []

    class Config:
        json_schema_extra = {
            "example": {
                "name": "mirror-0",
                "state": "ONLINE",
                "read": 0,
                "write": 0,
                "cksum": 0,
                "children": [
                    {"name": "sda", "state": "ONLINE", "read": 0, "write": 0, "cksum": 0},
                    {"name": "sdb", "state": "ONLINE", "read": 0, "write": 0, "cksum": 0}
                ]
            }
        }


# Update forward reference for recursive type
Vdev.model_rebuild()


class Pool(BaseModel):
    name: str
    size: int  # bytes
    alloc: int  # bytes
    free: int  # bytes
    fragmentation: str  # percentage
    capacity: str  # percentage
    health: PoolHealth

    class Config:
        json_schema_extra = {
            "example": {
                "name": "tank",
                "size": 3865470976,
                "alloc": 2040109465,
                "free": 1825361511,
                "fragmentation": "0%",
                "capacity": "52%",
                "health": "ONLINE"
            }
        }


class PoolStatus(BaseModel):
    name: str
    state: Optional[str] = None
    health: PoolHealth
    scan: Optional[str] = None
    errors: Optional[str] = None
    last_scrub: Optional[str] = None
    vdevs: List[Vdev] = []

    class Config:
        json_schema_extra = {
            "example": {
                "name": "tank",
                "state": "ONLINE",
                "health": "ONLINE",
                "scan": "scrub in progress since Sat Apr 14 10:30:00 2026",
                "errors": "No known data errors",
                "vdevs": [
                    {
                        "name": "mirror-0",
                        "state": "ONLINE",
                        "read": 0,
                        "write": 0,
                        "cksum": 0,
                        "children": [
                            {"name": "sda", "state": "ONLINE", "read": 0, "write": 0, "cksum": 0},
                            {"name": "sdb", "state": "ONLINE", "read": 0, "write": 0, "cksum": 0}
                        ]
                    }
                ]
            }
        }