models.py

from datetime import datetime
from typing import Optional

import numpy as np
from pydantic import BaseModel, ConfigDict


class FrameMeta:
    """A video frame with associated metadata. Not a Pydantic model because it holds numpy arrays."""

    def __init__(
        self,
        frame_image: np.ndarray,
        timestamp: datetime,
        gps_lat: float,
        gps_lon: float,
        frame_index: int = 0,
        source_path: Optional[str] = None,
    ):
        self.frame_image = frame_image
        self.timestamp = timestamp
        self.gps_lat = gps_lat
        self.gps_lon = gps_lon
        self.frame_index = frame_index
        self.source_path = source_path


class HazardDetection(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    hazard_type: str
    description: str
    confidence: float
    severity: str  # low | medium | high | critical
    timestamp: datetime
    gps_lat: float
    gps_lon: float
    frame_index: int = 0


class VideoChunk:
    """A clip-sized segment of video with metadata."""

    def __init__(
        self,
        frames: list[FrameMeta],
        start_time: datetime,
        end_time: datetime,
        gps_lat: float,
        gps_lon: float,
        chunk_index: int,
        temp_clip_path: str,
        is_stationary: bool = False,
    ):
        self.frames = frames
        self.start_time = start_time
        self.end_time = end_time
        self.gps_lat = gps_lat
        self.gps_lon = gps_lon
        self.chunk_index = chunk_index
        self.temp_clip_path = temp_clip_path
        self.is_stationary = is_stationary


# ---------------------------------------------------------------------------
# CivicAurAI Spanner models (V2 — hackathon schema)
# ---------------------------------------------------------------------------


class User(BaseModel):
    user_id: str
    name: Optional[str] = None
    role: str  # CITIZEN | CITY_WORKER | AI_SYSTEM | PARTNER
    is_anonymous: bool = False


class Organization(BaseModel):
    org_id: str
    name: str
    org_type: str  # CITY_AGENCY | COMMUNITY_GROUP | CONTRACTOR
    capabilities: list[str] = []


class District(BaseModel):
    district_id: str
    name: str


class IssueCategory(BaseModel):
    category_id: str
    display_name: str


class Issue(BaseModel):
    issue_id: str
    category_id: str
    district_id: Optional[str] = None
    severity: Optional[str] = None
    status: str = "NEW"  # NEW | IN_PROGRESS | RESOLVED | VERIFIED
    assigned_org_id: Optional[str] = None


class Report(BaseModel):
    report_id: str
    issue_id: str
    reporter_id: Optional[str] = None
    segment_id: Optional[str] = None
    source_type: str  # AI_VISION | CITIZEN_APP | CITY_WORKER
    description: Optional[str] = None
    ai_metadata: Optional[dict] = None
    embedding: Optional[list[float]] = None


class Video(BaseModel):
    video_id: str
    source_device: Optional[str] = None
    gcs_uri: Optional[str] = None
    capture_start_time: Optional[datetime] = None
    capture_end_time: Optional[datetime] = None


class VideoSegment(BaseModel):
    segment_id: str
    video_id: str
    start_time_offset: float
    end_time_offset: float
    ai_summary: Optional[str] = None
    gcs_uri: Optional[str] = None
    embedding: Optional[list[float]] = None


class IssueEpisode(BaseModel):
    issue_id: str
    episode_id: str
    actor_id: Optional[str] = None
    previous_status: Optional[str] = None
    new_status: str
    comment: Optional[str] = None


class MediaBlob(BaseModel):
    media_id: str
    report_id: Optional[str] = None
    episode_id: Optional[str] = None
    gcs_uri: str
    media_type: str  # VIDEO_CLIP | PHOTO | THUMBNAIL
    embedding: Optional[list[float]] = None


# ---------------------------------------------------------------------------
# Category mapping — AI free-form hazard_type → controlled CategoryId
# ---------------------------------------------------------------------------

HAZARD_TYPE_TO_CATEGORY: dict[str, str] = {
    # Potholes / road damage → STREET_LIGHT_POTHOLE
    "pothole": "STREET_LIGHT_POTHOLE",
    "potholes": "STREET_LIGHT_POTHOLE",
    "road surface damage": "STREET_LIGHT_POTHOLE",
    "pothole or road surface damage": "STREET_LIGHT_POTHOLE",
    "potholes or road surface damage": "STREET_LIGHT_POTHOLE",
    "street light": "STREET_LIGHT_POTHOLE",
    "broken street light": "STREET_LIGHT_POTHOLE",
    # Garbage / debris
    "garbage": "GARBAGE_WASTE",
    "trash": "GARBAGE_WASTE",
    "waste": "GARBAGE_WASTE",
    "debris": "GARBAGE_WASTE",
    "debris, fallen objects, or obstructions on the road": "GARBAGE_WASTE",
    "overflowing trash": "GARBAGE_WASTE",
    # Biohazard
    "biohazard": "BIOHAZARD",
    "needles": "BIOHAZARD",
    "needle": "BIOHAZARD",
    "human feces": "BIOHAZARD",
    # Homeless outreach
    "homeless": "HOMELESS_OUTREACH",
    "encampment": "HOMELESS_OUTREACH",
    "tent": "HOMELESS_OUTREACH",
    # Person in distress
    "person in distress": "PERSON_IN_DISTRESS",
    "person lying": "PERSON_IN_DISTRESS",
    # Animal waste
    "dog poop": "ANIMAL_WASTE",
    "animal waste": "ANIMAL_WASTE",
    # Legacy mappings that still go to STREET_LIGHT_POTHOLE
    "damaged sidewalks or curbs": "STREET_LIGHT_POTHOLE",
    "damaged sidewalk or curb": "STREET_LIGHT_POTHOLE",
    "damaged sidewalk": "STREET_LIGHT_POTHOLE",
    "exposed utility covers": "STREET_LIGHT_POTHOLE",
    "exposed utility cover": "STREET_LIGHT_POTHOLE",
    "missing manhole cover": "STREET_LIGHT_POTHOLE",
    "broken traffic lights": "STREET_LIGHT_POTHOLE",
    "damaged, missing, or obscured road signs": "STREET_LIGHT_POTHOLE",
    "flooding": "GARBAGE_WASTE",
    "fallen trees": "GARBAGE_WASTE",
}

# Keywords used as fallback when exact match fails
_CATEGORY_KEYWORDS: dict[str, str] = {
    "pothole": "STREET_LIGHT_POTHOLE",
    "street light": "STREET_LIGHT_POTHOLE",
    "sidewalk": "STREET_LIGHT_POTHOLE",
    "curb": "STREET_LIGHT_POTHOLE",
    "manhole": "STREET_LIGHT_POTHOLE",
    "utility": "STREET_LIGHT_POTHOLE",
    "sign": "STREET_LIGHT_POTHOLE",
    "traffic light": "STREET_LIGHT_POTHOLE",
    "garbage": "GARBAGE_WASTE",
    "trash": "GARBAGE_WASTE",
    "debris": "GARBAGE_WASTE",
    "flood": "GARBAGE_WASTE",
    "tree": "GARBAGE_WASTE",
    "obstruction": "GARBAGE_WASTE",
    "scooter": "GARBAGE_WASTE",
    "biohazard": "BIOHAZARD",
    "needle": "BIOHAZARD",
    "feces": "BIOHAZARD",
    "homeless": "HOMELESS_OUTREACH",
    "encampment": "HOMELESS_OUTREACH",
    "tent": "HOMELESS_OUTREACH",
    "distress": "PERSON_IN_DISTRESS",
    "poop": "ANIMAL_WASTE",
    "animal": "ANIMAL_WASTE",
}

DEFAULT_ISSUE_CATEGORIES: list[IssueCategory] = [
    IssueCategory(category_id="GARBAGE_WASTE", display_name="Garbage & Waste"),
    IssueCategory(category_id="BIOHAZARD", display_name="Biohazard / Needles / Human Feces"),
    IssueCategory(category_id="HOMELESS_OUTREACH", display_name="Homeless Outreach / Encampment"),
    IssueCategory(category_id="PERSON_IN_DISTRESS", display_name="Person in Distress"),
    IssueCategory(category_id="ANIMAL_WASTE", display_name="Dog Poop / Animal Waste"),
    IssueCategory(category_id="STREET_LIGHT_POTHOLE", display_name="Street Light / Pothole"),
]


# ---------------------------------------------------------------------------
# Resolution policy — infrastructure vs transient
# ---------------------------------------------------------------------------

# Infrastructure: requires civic ticket, manual resolution, camera verification
CIVIC_TICKET_CATEGORIES: set[str] = {"STREET_LIGHT_POTHOLE"}

# Transient: auto-resolvable when not detected on follow-up pass
AUTO_RESOLVE_CATEGORIES: set[str] = {
    "GARBAGE_WASTE", "BIOHAZARD", "ANIMAL_WASTE", "PERSON_IN_DISTRESS",
}


def is_auto_resolvable(category_id: str) -> bool:
    """Return True if the category can be auto-resolved by the pipeline."""
    return category_id not in CIVIC_TICKET_CATEGORIES


def normalize_category(hazard_type: str) -> str:
    """Map AI hazard_type string to a CategoryId. Falls back to GARBAGE_WASTE."""
    key = hazard_type.lower().strip()
    # Exact match first
    if key in HAZARD_TYPE_TO_CATEGORY:
        return HAZARD_TYPE_TO_CATEGORY[key]
    # Keyword fallback
    for keyword, cat in _CATEGORY_KEYWORDS.items():
        if keyword in key:
            return cat
    return "GARBAGE_WASTE"