import uuid
from datetime import date, datetime
from typing import Literal, Optional

from pydantic import BaseModel, ConfigDict, Field, model_validator

StudyStatus = Literal["DRAFT", "ACTIVE", "CLOSED"]


class VisitScheduleItem(BaseModel):
    visit_code: str = Field(min_length=1, max_length=50)
    baseline_offset_days: int = Field(ge=0, le=3650)
    window_before_days: int = Field(ge=0, le=365)
    window_after_days: int = Field(ge=0, le=365)


class VisitScheduleMixin(BaseModel):
    visit_schedule: list[VisitScheduleItem] = Field(default_factory=list)

    @model_validator(mode="after")
    def validate_visit_schedule(self):
        codes: set[str] = set()
        for index, item in enumerate(self.visit_schedule):
            code = item.visit_code.strip()
            if not code:
                raise ValueError(f"第 {index + 1} 行访视编号不能为空")
            if code in codes:
                raise ValueError(f"访视编号重复：{code}")
            codes.add(code)
            item.visit_code = code
        return self


class StudyCreate(VisitScheduleMixin):
    name: str = Field(min_length=1)
    code: str = Field(min_length=1)
    project_full_name: Optional[str] = None
    sponsor: Optional[str] = None
    protocol_no: Optional[str] = None
    lead_unit: Optional[str] = None
    principal_investigator: Optional[str] = None
    main_pm: Optional[str] = None
    research_analysis: Optional[str] = None
    research_product: Optional[str] = None
    control_product: Optional[str] = None
    indication: Optional[str] = None
    research_population: Optional[str] = None
    research_design: Optional[str] = None
    plan_start_date: Optional[date] = None
    plan_end_date: Optional[date] = None
    planned_site_count: Optional[int] = None
    planned_enrollment_count: Optional[int] = None
    enrollment_monthly_goal_note: Optional[str] = None
    enrollment_stage_breakdown: Optional[str] = None
    phase: Optional[str] = None
    status: StudyStatus = "DRAFT"


class StudyUpdate(BaseModel):
    code: Optional[str] = None
    name: Optional[str] = None
    project_full_name: Optional[str] = None
    sponsor: Optional[str] = None
    protocol_no: Optional[str] = None
    lead_unit: Optional[str] = None
    principal_investigator: Optional[str] = None
    main_pm: Optional[str] = None
    research_analysis: Optional[str] = None
    research_product: Optional[str] = None
    control_product: Optional[str] = None
    indication: Optional[str] = None
    research_population: Optional[str] = None
    research_design: Optional[str] = None
    plan_start_date: Optional[date] = None
    plan_end_date: Optional[date] = None
    planned_site_count: Optional[int] = None
    planned_enrollment_count: Optional[int] = None
    enrollment_monthly_goal_note: Optional[str] = None
    enrollment_stage_breakdown: Optional[str] = None
    phase: Optional[str] = None
    status: Optional[StudyStatus] = None
    is_locked: Optional[bool] = None
    visit_schedule: Optional[list[VisitScheduleItem]] = None

    @model_validator(mode="after")
    def validate_visit_schedule(self):
        if self.visit_schedule is None:
            return self
        codes: set[str] = set()
        for index, item in enumerate(self.visit_schedule):
            code = item.visit_code.strip()
            if not code:
                raise ValueError(f"第 {index + 1} 行访视编号不能为空")
            if code in codes:
                raise ValueError(f"访视编号重复：{code}")
            codes.add(code)
            item.visit_code = code
        return self


class StudyRead(BaseModel):
    id: uuid.UUID
    code: str
    name: str
    project_full_name: Optional[str]
    sponsor: Optional[str]
    protocol_no: Optional[str]
    lead_unit: Optional[str]
    principal_investigator: Optional[str]
    main_pm: Optional[str]
    research_analysis: Optional[str]
    research_product: Optional[str]
    control_product: Optional[str]
    indication: Optional[str]
    research_population: Optional[str]
    research_design: Optional[str]
    plan_start_date: Optional[date]
    plan_end_date: Optional[date]
    planned_site_count: Optional[int]
    planned_enrollment_count: Optional[int]
    enrollment_monthly_goal_note: Optional[str]
    enrollment_stage_breakdown: Optional[str]
    phase: Optional[str]
    status: StudyStatus
    is_locked: bool
    visit_schedule: list[VisitScheduleItem]
    created_by: Optional[uuid.UUID]
    created_at: datetime

    model_config = ConfigDict(from_attributes=True)