from typing import Any, Optional

from pydantic import BaseModel, Field


class AutogenContext(BaseModel):
    # Init params
    project_key: str
    dataset_name: str
    llm_id: Optional[str] = None
    user_instructions: Optional[str] = None

    schema_cache: Optional[dict[str, Any]] = Field(default=None, description="Cached dataset schema")
    profile_cache: dict[str, dict[str, Any]] = Field(
        default_factory=dict, description="Column profiles keyed by column name"
    )

    entity_groups: list[dict[str, Any]] = Field(default_factory=list, description="Proposed entity groups with columns")
    working_model: Optional[dict[str, Any]] = Field(default=None, description="In-progress semantic model spec")

    validation_attempts: int = Field(default=0, description="Number of validation retry attempts")

    # Result
    warnings: list[str] = Field(default_factory=list, description="Non-fatal warnings accumulated during generation")
    status: str = Field(default="running", description="Status: running | completed | failed")

    class Config:
        extra = "ignore"
