from enum import Enum
from typing import List, Optional

from pydantic import BaseModel, field_validator


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 InvestorMemberSchema(BaseModel):
    id: int
    name: str
    role: str
    email: str
    investor_id: int

    class Config:
        from_attributes = True


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

    class Config:
        from_attributes = True


class CompanySchema(BaseModel):
    id: int
    name: str
    industry: str
    location: str
    description: Optional[str] = None  # Fixed typo from 'nullabel'
    founded_year: Optional[int] = None  # Changed from str to int to match model
    website: Optional[str] = None

    @field_validator("founded_year", mode="before")
    @classmethod
    def validate_founded_year(cls, v):
        if v is None or v == "Not Available" or v == "":
            return None
        if isinstance(v, str):
            try:
                return int(v)
            except ValueError:
                return None
        return v

    class Config:
        from_attributes = True


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

    class Config:
        from_attributes = True


class InvestorData(BaseModel):
    """Comprehensive investor data schema for LLM processing"""

    investor: InvestorSchema
    portfolio_companies: List[CompanySchema] = []
    team_members: List[InvestorMemberSchema] = []  # Changed from TeamMember
    sectors: List[SectorSchema] = []

    class Config:
        from_attributes = True


class CompanyData(BaseModel):  # Renamed from CompaniesData for consistency
    company: CompanySchema
    sectors: List[SectorSchema] = []
    members: List[CompanyMemberSchema] = []  # Changed to match model relationship name
    investors: List[InvestorSchema] = []

    class Config:
        from_attributes = True


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