Anton_wireframe/app/schemas/router_schemas.py

from datetime import datetime
from enum import Enum
from typing import Any, Generic, List, Optional, TypeVar

from pydantic import BaseModel

# Generic type for pagination
T = TypeVar("T")


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 InvestmentStageSchema(BaseModel):
    id: int
    name: str

    class Config:
        from_attributes = True


class InvestorMemberSchema(BaseModel):
    id: int
    name: str
    role: str | None
    email: str | None
    linkedin: str | None

    class Config:
        from_attributes = True


class FundSchema(BaseModel):
    id: int
    fund_name: str | None
    fund_size: int | None  # Changed to int for numerical filtering
    fund_size_source_url: str | None
    check_size_lower: int | None  # NEW: Lower bound of check size range
    check_size_upper: int | None  # NEW: Upper bound of check size range
    source_url: str | None
    source_provider: str | None
    geographic_focus: str | None  # Changed from List[str] to string
    investment_stages: List[InvestmentStageSchema] | None  # Changed to relationship
    sectors: List[SectorSchema] | None  # Changed to relationship
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        from_attributes = True


class CompanyMemberSchema(BaseModel):
    id: int
    name: Optional[str]
    linkedin: Optional[str]
    role: Optional[str]
    company_id: int

    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 InvestorSchema(BaseModel):
    id: int
    name: str
    description: Optional[str]
    website: Optional[str] = None
    headquarters: Optional[str] = None
    aum: int | None
    aum_as_of_date: str | None = None
    aum_source_url: str | None = None
    check_size_lower: int | None
    check_size_upper: int | None
    geographic_focus: str | None
    investment_thesis: Any = (
        None  # Flexible JSON field - can be list, dict, or list of dicts
    )
    portfolio_highlights: Any = (
        None  # Flexible JSON field - can be list, dict, or list of dicts
    )
    number_of_investments: int | None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        from_attributes = True


class InvestorData(BaseModel):
    """Comprehensive investor data schema - used for individual investor requests"""

    investor: InvestorSchema
    portfolio_companies: List[CompanySchema]
    team_members: List[InvestorMemberSchema]
    sectors: List[SectorSchema]
    funds: List[FundSchema]

    class Config:
        from_attributes = True


class InvestorFundData(BaseModel):
    """Investor-Fund combined data - used for list/filter requests

    Each row represents one investor-fund combination.
    An investor with 3 funds will appear as 3 separate entries.
    """

    # Investor fields
    investor_id: int
    investor_name: str
    investor_description: Optional[str]
    investor_website: Optional[str]
    investor_headquarters: Optional[str]
    aum: int | None
    aum_as_of_date: str | None
    aum_source_url: str | None
    investment_thesis: Any = None  # Flexible JSON field
    portfolio_highlights: Any = None  # Flexible JSON field
    number_of_investments: int | None

    # Fund fields
    fund_id: int | None
    fund_name: str | None
    fund_size: int | None  # Changed to int for numerical filtering
    fund_size_source_url: str | None
    check_size_lower: int | None  # NEW: Lower bound of check size range
    check_size_upper: int | None  # NEW: Upper bound of check size range
    geographic_focus: str | None  # Changed from List[str] to string
    fund_investment_stages: (
        List[InvestmentStageSchema] | None
    )  # Changed to relationship
    fund_sectors: List[SectorSchema] | None  # Changed to relationship

    # Related data
    portfolio_companies: List[CompanySchema]
    team_members: List[InvestorMemberSchema]
    sectors: List[SectorSchema]

    class Config:
        from_attributes = True


class InvestorMinimal(BaseModel):
    """Minimal investor info with just id and name"""

    id: int
    name: str

    class Config:
        from_attributes = True


class CompanySchemaMinimal(BaseModel):
    id: int
    name: str
    industry: str | None
    location: str | None
    founded_year: Optional[int]
    website: Optional[str]

    class Config:
        from_attributes = True


class CompanyData(BaseModel):  # Renamed from CompaniesData for consistency
    company: CompanySchemaMinimal
    investors: List[InvestorMinimal]
    members: List[CompanyMemberSchema] = []
    sectors: List[SectorSchema] = []

    class Config:
        from_attributes = True


class InvestorList(BaseModel):
    investors: List[InvestorData]


class InvestorFundList(BaseModel):
    """List of investor-fund combinations"""

    investor_funds: List[InvestorFundData]


class CompanyMinimal(BaseModel):
    """Minimal company info with just id and name"""

    id: int
    name: str

    class Config:
        from_attributes = True


class SectorMinimal(BaseModel):
    """Minimal sector info with just id and name"""

    id: int
    name: str

    class Config:
        from_attributes = True


class InvestmentResponse(BaseModel):
    """Simplified investment response schema

    One row per investor-fund combination with streamlined data
    """

    id: int  # Investor ID
    name: (
        str  # Combination of investor name and fund name (e.g., "Investor A - Fund A")
    )
    aum: int | None  # From investor
    check_size_lower: int | None  # From fund
    check_size_upper: int | None  # From fund
    geographic_focus: str | None  # From fund
    stage_focus: str | None  # Comma-separated stages from fund
    portfolio_companies: List[CompanyMinimal]  # Top 3 companies from investor
    sectors: List[SectorMinimal]  # Top 3 sectors from fund
    compatibility_score: float  # 0 to 1 (default 1 for now)

    class Config:
        from_attributes = True


class PaginatedResponse(BaseModel, Generic[T]):
    """Generic paginated response schema"""

    items: List[T]
    total: int
    page: int
    page_size: int
    total_pages: int

    class Config:
        from_attributes = True