from pydantic import BaseModel
from typing import Dict, List, Optional, Union, Any

# =============================================================================
# REQUEST SCHEMA
# =============================================================================

class PredictRequest(BaseModel):
    """
    Forest segmentation prediction request.

    This schema is intentionally flexible to support:
    - Supabase Edge Functions
    - Hugging Face remote inference
    - Local inference scripts

    Required:
    - bands: Dict[str, Union[str, List[float]]]

    Allowed extra fields (sent by Supabase):
    - width, height
    - bbox
    - band_names
    - preprocessing
    - model_name, model_version
    """

    model_name: str = "forest_segmentation"
    model_version: str = "landsat8_v1"

    # Band data:
    # - base64-encoded float32 (remote calls)
    # - list/array of floats (local calls)
    bands: Dict[str, Union[str, List[float]]]

    class Config:
        # 🔑 CRITICAL FIX
        # Allows Supabase to send extra metadata without 422 errors
        extra = "allow"


# =============================================================================
# RESPONSE SCHEMA
# =============================================================================

class PredictResponse(BaseModel):
    """
    Forest segmentation prediction response.

    Mask values are CONTINUOUS (0–255) — NOT binary.
    """

    # Flattened mask (length = width * height)
    mask: List[int]

    # Inverted mask (optional utility)
    inverted_mask: List[int]

    # Statistics
    forest_percentage: float
    forest_confidence: float
    mean_prediction: float

    # Class mapping
    classes: Dict[str, int]

    # Model metadata
    model_info: Dict[str, Any]

    # Optional debug block
    debug: Optional[Dict[str, Any]] = None