from __future__ import annotations
from typing import Optional

import uuid
from datetime import date, datetime

from sqlalchemy import Boolean, Date, DateTime, ForeignKey, Integer, JSON, String, Text, func
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import Mapped, mapped_column

from app.db.base_class import Base


class Study(Base):
    __tablename__ = "studies"

    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    code: Mapped[str] = mapped_column(String(50), unique=True, nullable=False, index=True)
    name: Mapped[str] = mapped_column(String(200), nullable=False)
    project_full_name: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
    sponsor: Mapped[Optional[str]] = mapped_column(String(200), nullable=True)
    protocol_no: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    lead_unit: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    principal_investigator: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    main_pm: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    research_analysis: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    research_product: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    control_product: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    indication: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    research_population: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    research_design: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    plan_start_date: Mapped[Optional[date]] = mapped_column(Date, nullable=True)
    plan_end_date: Mapped[Optional[date]] = mapped_column(Date, nullable=True)
    planned_site_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    planned_enrollment_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    enrollment_monthly_goal_note: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    enrollment_stage_breakdown: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    phase: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
    status: Mapped[str] = mapped_column(String(20), nullable=False, default="DRAFT")
    is_locked: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    visit_schedule: Mapped[list[dict]] = mapped_column(JSON, nullable=False, default=list)
    created_by: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now())