from datetime import datetime
from enum import Enum
from typing import Any, Literal, Optional

from pydantic import BaseModel, ConfigDict, Field

from semantic_models.glossary_types import GlossaryTermSource

class AttributeType(str, Enum):
    COLUMN = "COLUMN"
    EXPRESSION = "EXPRESSION"


class DistinctValuesHandlingMode(str, Enum):
    NONE = "NONE"
    MANUAL = "MANUAL"
    AUTO_INDEX = "AUTO_INDEX"


class Attribute(BaseModel):
    """DSS attribute definition."""

    name: str = Field(..., min_length=1, description="Attribute name")
    description: Optional[str] = Field(None, description="Human-readable description")
    dssType: Optional[str] = Field(None, description="DSS type for the attribute")
    type: AttributeType = Field(AttributeType.COLUMN, description="Attribute type (COLUMN or EXPRESSION)")
    column: Optional[str] = Field(None, description="Source column name (for COLUMN type)")
    pseudoSQLExpression: Optional[str] = Field(None, description="Pseudo-SQL expression (for EXPRESSION type)")
    distinctValuesHandlingMode: DistinctValuesHandlingMode = Field(
        DistinctValuesHandlingMode.NONE, description="How distinct values are handled"
    )
    manualValues: list[str] = Field(default_factory=list, description="Manual distinct values")
    indexDistinctValues: bool = Field(False, description="Whether to index distinct values")
    resolveInUserRequests: bool = Field(
        False, description="Include in term resolution for fuzzy matching user query terms"
    )
    llmInstructions: Optional[str] = Field(None, description="LLM instruction for this attribute")
    sqlGenerationConfig: dict[str, Any] = Field(
        default_factory=dict, description="SQL generation config for this attribute"
    )


class Metric(BaseModel):
    """DSS metric definition with aggregation."""

    name: str = Field(..., min_length=1, description="Metric name")
    description: Optional[str] = Field(None, description="Human-readable description")
    pseudoSQLExpression: str = Field(..., description="Aggregation expression (e.g., SUM(amount))")
    llmInstructions: Optional[str] = Field(None, description="LLM instruction for this metric")
    created: Optional["DSSNiceVersionTag"] = None


class SemanticModelFilter(BaseModel):
    """DSS filter definition (pseudo-SQL WHERE clause fragment)."""

    name: str = Field(..., min_length=1, description="Filter name")
    description: Optional[str] = Field(None, description="Human-readable description")
    pseudoSQLExpression: str = Field(..., description="SQL WHERE clause fragment")
    llmInstructions: Optional[str] = Field(None, description="LLM instruction for this filter")
    created: Optional["DSSNiceVersionTag"] = None


class Key(BaseModel):
    """DSS key definition (list of attribute names)."""

    attributes: list[str] = Field(default_factory=list)


class ForeignKey(BaseModel):
    referencedEntity: str
    key: Key


class EntityClass(BaseModel):
    """DSS entity class representing a table or dataset."""

    model_config = ConfigDict(protected_namespaces=())

    name: str = Field(..., min_length=1, description="Entity name")
    description: Optional[str] = Field(None, description="Human-readable description")
    tags: list[str] = Field(default_factory=list)

    type: Literal["TABLE", "DATASET"] = Field("DATASET", description="Entity source type")
    datasetRef: Optional[str] = Field(None, description="Dataset reference (PROJECT_KEY.DATASET_NAME)")

    connection: Optional[str] = Field(None, description="DSS connection name")
    catalog: Optional[str] = Field(None, description="Database catalog name")
    schema: Optional[str] = Field(None, description="Database schema name")
    table: Optional[str] = Field(None, description="Table name")

    attributes: list[Attribute] = Field(default_factory=list, description="Available attributes")
    metrics: list[Metric] = Field(default_factory=list, description="Available metrics")
    filters: list[SemanticModelFilter] = Field(default_factory=list, description="Predefined filters")

    primaryKey: Key = Field(default_factory=Key)
    foreignKeys: list[ForeignKey] = Field(default_factory=list)


class Relationship(BaseModel):
    """DSS custom relationship definition."""

    name: Optional[str] = Field(None, description="Relationship name")
    description: Optional[str] = Field(None, description="Human-readable description")
    firstEntity: str = Field(..., description="First entity name")
    secondEntity: str = Field(..., description="Second entity name")
    pseudoSQLExpression: Optional[str] = Field(
        None, description="Pseudo-SQL join hint (e.g., 'orders.customer_id = customers.id')"
    )


class SemanticModelSpec(BaseModel):
    """
    The complete specification of a semantic model.

    This is the main structure stored in version spec_json fields.
    """

    entities: list[EntityClass] = Field(default_factory=list, description="Entity classes")
    relationships: list[Relationship] = Field(default_factory=list, description="Relationships between entities")
    instructions: Optional[str] = Field(
        None, description="Model-level LLM instructions from sqlGenerationConfig.instructions"
    )


# =============================================================================
# DSS-aligned Models
# =============================================================================


class DSSNiceVersionTag(BaseModel):
    """DSS NiceVersionTag structure used in version metadata."""

    on: Optional[str] = None
    by: Optional[str] = None
    version: Optional[int] = None


class DSSIndexingSettings(BaseModel):
    """Indexing settings for semantic model versions (DSS schema)."""

    maxDistinctValuesPerAttribute: int = 1000
    maxScannedRowsForSQLDatasets: int = -1
    maxScannedRowsForNonSQLDatasets: int = 1000000


class SemanticModelVersionBase(BaseModel):
    """Shared fields for semantic model versions."""

    id: str
    description: Optional[str] = None
    created: Optional[DSSNiceVersionTag] = None
    entities: list[dict[str, Any]] = Field(default_factory=list)
    relationships: list[dict[str, Any]] = Field(default_factory=list)
    goldenQueries: list[dict[str, Any]] = Field(default_factory=list)
    glossaryTerms: list[dict[str, Any]] = Field(default_factory=list)
    glossaryBindings: list[dict[str, Any]] = Field(default_factory=list)
    # vocabularyTermIds is stored in sqlGenerationConfig for DSS persistence
    indexingSettings: Optional[dict[str, Any] | DSSIndexingSettings] = None
    privateEditorData: dict[str, Any] = Field(default_factory=dict)
    sqlGenerationConfig: dict[str, Any] = Field(default_factory=dict)


class DSSSemanticModelVersion(SemanticModelVersionBase):
    """DSS-aligned semantic model version (flat, camelCase)."""

    indexingSettings: Optional[DSSIndexingSettings] = None


class SemanticModelBase(BaseModel):
    """Shared fields for semantic model metadata."""

    id: str
    projectKey: str
    name: str
    activeVersionId: Optional[str] = None
    versionTag: Optional[dict[str, Any]] = None
    creationTag: Optional[dict[str, Any]] = None
    tags: list[str] = Field(default_factory=list)
    customFields: dict[str, Any] = Field(default_factory=dict)
    checklists: dict[str, Any] = Field(default_factory=dict)
    description: Optional[str] = None
    shortDesc: Optional[str] = None


class DSSSemanticModel(SemanticModelBase):
    """DSS-aligned semantic model definition (camelCase)."""

    versions: list[DSSSemanticModelVersion] = Field(default_factory=list)


# =============================================================================
# API Request/Response Models
# =============================================================================


class SemanticModelCreate(BaseModel):
    """Request body for creating a new semantic model."""

    name: str = Field(..., min_length=1, description="Model name")
    description: Optional[str] = Field(None, description="Model description")


class SemanticModelUpdate(BaseModel):
    """Request body for updating a semantic model."""

    name: Optional[str] = Field(None, min_length=1, description="Model name")
    description: Optional[str] = Field(None, description="Model description")
    shortDesc: Optional[str] = Field(None, description="Short description")


class VersionCreate(BaseModel):
    """Request body for creating a new version."""

    fromVersion: Optional[str] = Field(
        None,
        description="Copy entities, relationships, glossary, vocabulary, and golden queries from this version",
    )
    version: Optional[dict[str, Any]] = Field(None, description="Version payload")


class GoldenQueryInput(BaseModel):
    """Golden query data (DSS schema)."""

    name: str = Field(..., min_length=1, description="Query name")
    description: Optional[str] = Field(None, description="Query description")
    question: str = Field(..., min_length=1, description="Natural language question")
    generatedSql: str = Field(..., description="Generated SQL query")
    created: Optional["DSSNiceVersionTag"] = Field(None, description="Creation tag")


class VersionUpdate(BaseModel):
    """Request body for updating a version."""

    version: dict[str, Any] = Field(..., description="Version patch payload")


class SetActiveVersion(BaseModel):
    """Request body for setting the active version."""

    versionId: str = Field(..., description="Version ID to set as active")


class PlaygroundFullRequest(BaseModel):
    """Request body for generate + execute in one call."""

    semantic_model_id: str = Field(..., description="Semantic model ID")
    version_id: Optional[str] = Field(None, description="Version ID (defaults to active)")
    llm_id: Optional[str] = Field(None, description="LLM to use")
    embedding_llm_id: Optional[str] = Field(None, description="Embedding LLM for semantic matching")
    max_rows_per_query: Optional[int] = Field(None, description="Override max rows per query")
    agent_recursion_limit: Optional[int] = Field(
        None, description="Override agent recursion limit"
    )
    sources_records_limit: Optional[int] = Field(
        None, description="Override max rows in sources"
    )
    artifacts_records_limit: Optional[int] = Field(
        None, description="Override max rows in artifacts"
    )
    question: str = Field(..., min_length=1, description="Natural language question")


# =============================================================================
# Response Models
# =============================================================================


class GoldenQueryResponse(BaseModel):
    """Response model for a golden query (DSS schema)."""

    name: str
    description: Optional[str]
    question: str
    generatedSql: str
    created: Optional["DSSNiceVersionTag"] = None


class VersionResponse(SemanticModelVersionBase):
    """Response model for a semantic model version."""

    indexingSettings: Optional[dict[str, Any]] = None


class SemanticModelResponse(SemanticModelBase):
    """Response model for a semantic model."""

    versions: list[VersionResponse] = Field(default_factory=list)


class SemanticModelListItem(SemanticModelBase):
    """Response model for semantic model list item."""


class TermCorrectionResponse(BaseModel):
    """A term correction applied during resolution."""

    original: str = Field(..., description="Original term from user query")
    resolved: str = Field(..., description="Resolved value from index")
    entity_id: str = Field(..., description="Entity the value belongs to")
    attribute: str = Field(..., description="Attribute name")
    score: float = Field(..., description="Match confidence (0-100)")
    match_type: str = Field(default="fuzzy", description="How the match was found: 'exact', 'fuzzy', or 'semantic'")
    source: Literal["index", "glossary"] = Field(
        default="index",
        description="Where the correction came from: 'index' for value resolution, 'glossary' for term mapping",
    )
    glossary_term_id: Optional[str] = Field(None, description="ID of the glossary term if source='glossary'")
    llm_inferred: bool = Field(
        False, description="True if LLM translated the user's term to a different candidate value before matching"
    )
    # Glossary binding info (for source='glossary')
    target_type: Optional[Literal["entity", "attribute", "metric", "filter"]] = Field(
        None, description="For glossary terms: binding target type"
    )
    target_name: Optional[str] = Field(None, description="For glossary terms: the name of the bound element")


class SourceItem(BaseModel):
    """A source item in the execution trace."""

    type: str = Field(..., description="Type: INFO, GENERATED_SQL_QUERY, RECORDS, ERROR")
    textSnippet: Optional[str] = Field(None, description="Text content for INFO/ERROR types")
    performedQuery: Optional[str] = Field(None, description="SQL query for GENERATED_SQL_QUERY type")
    records: Optional[dict] = Field(None, description="Records data for RECORDS type")


class Sources(BaseModel):
    """Execution trace sources (for agent tool compatibility)."""

    toolCallDescription: str = Field(..., description="Description of the tool call")
    items: list[SourceItem] = Field(default_factory=list, description="Source items")


class ArtifactPart(BaseModel):
    """A part of an artifact."""

    type: str = Field(..., description="Part type (e.g., RECORDS)")
    records: Optional[dict] = Field(None, description="Records data with columns and data")


class Artifact(BaseModel):
    """An artifact containing structured data (for agent tool compatibility)."""

    name: str = Field(..., description="Artifact name")
    description: str = Field(..., description="Artifact description")
    parts: list[ArtifactPart] = Field(default_factory=list, description="Artifact parts")


class ValueCorrection(BaseModel):
    """Value correction from resolve_values tool."""

    original: str = Field(..., description="Original term from user's question")
    resolved: str = Field(..., description="Resolved value found in data")
    entity_id: str = Field(..., description="Entity where value was found")
    attribute: str = Field(..., description="Attribute/column where value was found")
    score: float = Field(..., description="Match confidence score (0-100)")
    match_type: str = Field(..., description="Match type: exact, fuzzy, or semantic")


class GlossaryTermUsed(BaseModel):
    """Glossary term used during query execution (with resolved binding info)."""

    term: str = Field(..., description="The glossary term")
    description: Optional[str] = Field(None, description="Term description")
    synonyms: Optional[list[str]] = Field(None, description="Term synonyms")
    entity_id: Optional[str] = Field(None, description="Bound entity ID")
    entity_name: Optional[str] = Field(None, description="Bound entity name (resolved)")
    target_type: Optional[str] = Field(None, description="Binding target type: entity, attribute, metric, or filter")
    target_id: Optional[str] = Field(None, description="Bound element ID")
    target_name: Optional[str] = Field(None, description="Bound element name (resolved)")


class PlaygroundFullResponse(BaseModel):
    """Response for generate + execute."""

    success: bool = Field(False, description="Whether SQL execution succeeded")
    error: Optional[str] = Field(None, description="Error message if success=false")

    # Agent's natural language response
    response: Optional[str] = Field(None, description="Agent's natural language response summarizing the results")

    # Value corrections and glossary terms used
    corrections: list[ValueCorrection] = Field(
        default_factory=list,
        description="Value corrections from resolve_values tool",
    )
    glossary_terms_used: list[GlossaryTermUsed] = Field(
        default_factory=list,
        description="Glossary terms used during query execution",
    )

    sources: Optional[Sources] = Field(None, description="Execution trace sources for agent tool compatibility")
    artifacts: list[Artifact] = Field(
        default_factory=list,
        description="Artifacts containing structured data (RECORDS) for export",
    )
    trace: Optional[dict[str, Any]] = Field(None, description="Full execution trace tree")

    # Timing
    duration_ms: Optional[float] = Field(None, description="Total execution time in milliseconds")

    # Legacy field (kept for golden query saving)
    sql: Optional[str] = Field(None, description="Last successful SQL (for golden query saving)")


# =============================================================================
# Auto-Generation Models
# =============================================================================


class EntityGroupProposal(BaseModel):
    """Proposed entity grouping from LLM."""

    entity_id: str = Field(..., description="Unique identifier for the entity")
    entity_name: str = Field(..., description="Human-readable entity name")
    description: Optional[str] = Field(None, description="Entity description")
    columns: list[str] = Field(default_factory=list, description="Columns in this entity")


class CoverageReport(BaseModel):
    """Column coverage statistics."""

    total_columns: int = Field(..., description="Total columns in dataset")
    used_columns: int = Field(..., description="Columns assigned to entities")
    unassigned_columns: list[str] = Field(default_factory=list, description="Columns not assigned to any entity")
    shared_columns: list[str] = Field(default_factory=list, description="Columns assigned to multiple entities")
    shared_details: dict[str, list[str]] = Field(default_factory=dict, description="Map of shared column -> entity IDs")


class AutogenModelResponse(BaseModel):
    """Response from /model/generate endpoint."""

    status: Literal["completed", "failed"] = Field(..., description="Generation status")
    model: Optional[SemanticModelSpec] = Field(None, description="Generated semantic model (null on early failure)")
    coverage: Optional[CoverageReport] = Field(None, description="Coverage statistics (null on early failure)")
    warnings: list[str] = Field(default_factory=list, description="Validation warnings (if any)")
    error: Optional[str] = Field(None, description="Error message (only if status='failed')")


# =============================================================================
# Glossary Models
# =============================================================================

class BindingTargetType(str, Enum):
    """Type of element a glossary term can bind to (DSS)."""

    ENTITY_CLASS = "ENTITY_CLASS"
    ATTRIBUTE = "ATTRIBUTE"
    METRIC = "METRIC"
    FILTER = "FILTER"


class GlossaryBindingCreate(BaseModel):
    """Request body for creating a glossary binding (DSS schema)."""

    termId: str = Field(..., description="ID of the glossary term to bind")
    targetEntityClass: str = Field(..., description="Target entity name")
    targetType: BindingTargetType = Field(..., description="Type of element being bound")
    targetName: Optional[str] = Field(None, description="Target element name (None if targetType=ENTITY_CLASS)")


class GlossaryBindingResponse(BaseModel):
    """Response model for a glossary binding (DSS schema)."""

    termId: str
    targetEntityClass: str
    targetType: BindingTargetType
    targetName: Optional[str]
    created: Optional[DSSNiceVersionTag] = None


class GlossaryTermCreate(BaseModel):
    """Request body for creating a glossary term (DSS schema)."""

    id: Optional[str] = Field(None, description="Term ID (optional; generated if absent)")
    term: str = Field(..., min_length=1, description="The term/phrase")
    description: Optional[str] = Field(None, description="Explanation of the term")
    source: Optional[GlossaryTermSource] = Field(None, description="How the term was created")
    synonyms: Optional[list[str]] = Field(None, description="Alternative terms")
    userModified: Optional[bool] = Field(None, description="Whether the term was edited by a user")
    privateEditorData: Optional[dict[str, Any]] = Field(None, description="Private editor data")


class GlossaryTermUpdate(BaseModel):
    """Request body for updating a glossary term (DSS schema)."""

    term: Optional[str] = Field(None, min_length=1, description="The term/phrase")
    description: Optional[str] = Field(None, description="Explanation of the term")
    source: Optional[GlossaryTermSource] = Field(None, description="How the term was created")
    synonyms: Optional[list[str]] = Field(None, description="Alternative terms")
    userModified: Optional[bool] = Field(None, description="Whether the term was edited by a user")
    privateEditorData: Optional[dict[str, Any]] = Field(None, description="Private editor data")


class GlossaryTermResponse(BaseModel):
    """Response model for a glossary term (DSS schema)."""

    id: str
    term: str
    description: Optional[str]
    source: GlossaryTermSource
    userModified: bool = Field(False, description="Whether the term has been manually edited")
    synonyms: list[str] = Field(default_factory=list)
    created: Optional[DSSNiceVersionTag] = None
    modified: Optional[DSSNiceVersionTag] = None
    privateEditorData: dict[str, Any] = Field(default_factory=dict)


class GlossaryExtractionRequest(BaseModel):
    """Request body for extracting glossary terms from a document draft."""

    llm_id: Optional[str] = Field(None, description="LLM to use for extraction")
    max_chunks: Optional[int] = Field(None, description="Maximum number of chunks to process")
    chunk_size: Optional[int] = Field(None, description="Target size in characters per chunk")
    chunk_overlap: Optional[int] = Field(None, description="Overlap size in characters between chunks")
    spec: Optional[dict[str, Any]] = Field(None, description="Semantic model spec for extraction context")
    terms: Optional[list[dict[str, Any]]] = Field(None, description="Existing draft terms for dedup/conflict checks")


class ExtractedTermData(BaseModel):
    """Data for an extracted term (used in conflict resolution)."""

    term: str
    description: Optional[str] = None
    synonyms: Optional[list[str]] = None
    privateEditorData: Optional[dict[str, Any]] = None


class GlossaryConflict(BaseModel):
    """A conflict between an existing user-edited term and an extracted term."""

    existing_term: GlossaryTermResponse
    extracted_term: ExtractedTermData


class GlossaryExtractionResponse(BaseModel):
    """Response for glossary extraction (draft, not persisted)."""

    documentName: str
    sectionsProcessed: int = Field(0, description="Number of document sections processed")
    termsCreated: list[GlossaryTermResponse] = Field(default_factory=list, description="New terms created")
    termsUpdated: list[GlossaryTermResponse] = Field(
        default_factory=list, description="Existing non-edited terms silently updated"
    )
    conflicts: list[GlossaryConflict] = Field(
        default_factory=list, description="Conflicts with user-edited terms requiring resolution"
    )


# Rebuild models to resolve forward references
VersionResponse.model_rebuild()
