Anton_wireframe/app/schemas/project_schemas.py

from datetime import datetime
from enum import Enum
from typing import List, Optional

from pydantic import BaseModel


class InvestmentStage(str, Enum):
    SEED = "SEED"
    SERIES_A = "SERIES_A"
    SERIES_B = "SERIES_B"
    SERIES_C = "SERIES_C"
    GROWTH = "GROWTH"
    LATE_STAGE = "LATE_STAGE"


class SectorSchema(BaseModel):
    id: int
    name: str

    class Config:
        from_attributes = True


class InvestorSchema(BaseModel):
    id: int
    name: str
    description: Optional[str]
    aum: int | None
    check_size_lower: int | None
    check_size_upper: int | None
    geographic_focus: str | None
    stage_focus: InvestmentStage
    number_of_investments: int | None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        from_attributes = True


class CompanySchema(BaseModel):
    id: int
    name: str
    industry: str | None
    location: str | None
    description: Optional[str]
    founded_year: Optional[int]
    website: Optional[str]
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        from_attributes = True


class ProjectSchema(BaseModel):
    id: int
    name: str
    valuation: int | None
    stage: InvestmentStage | None
    location: str | None
    description: Optional[str]
    start_date: Optional[datetime]
    end_date: Optional[datetime]
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        from_attributes = True


class ProjectCreate(BaseModel):
    name: str
    valuation: Optional[int] = None
    stage: Optional[InvestmentStage] = None
    location: Optional[str] = None
    description: Optional[str] = None
    start_date: Optional[datetime] = None
    end_date: Optional[datetime] = None


class ProjectUpdate(BaseModel):
    name: Optional[str] = None
    valuation: Optional[int] = None
    stage: Optional[InvestmentStage] = None
    location: Optional[str] = None
    description: Optional[str] = None
    start_date: Optional[datetime] = None
    end_date: Optional[datetime] = None


class ProjectData(BaseModel):
    """Comprehensive project data schema"""

    project: ProjectSchema
    sector: List[SectorSchema]
    investors: List[InvestorSchema]
    companies: List[CompanySchema]

    class Config:
        from_attributes = True


class ProjectInvestorAssociation(BaseModel):
    project_id: int
    investor_id: int


class ProjectCompanyAssociation(BaseModel):
    project_id: int
    company_id: int


class ProjectSectorAssociation(BaseModel):
    project_id: int
    sector_id: int