"""
Example: How Level 1 and Level 2 agents use the semantic model.

This demonstrates the workflow from indexing a source through creating
cross-source semantic entities.
"""

from datetime import datetime

from semantic_model import (
    # Core model
    SemanticModel,
    create_empty_model,
    # Level 1 components
    DataSource,
    SourceType,
    Table,
    TableType,
    SemanticRole,
    TemporalInfo,
    TemporalGrain,
    PrimaryKey,
    Column,
    SemanticType,
    SemanticCategory,
    AllowedValue,
    ColumnProfiling,
    InternalRelationship,
    RelationshipType,
    JoinHint,
    # Level 2 components
    SemanticEntity,
    EntityManifestation,
    ManifestationRole,
    KeyMapping,
    KeyMappingType,
    UnifiedAttribute,
    AttributeSource,
    IdentityResolution,
    IdentityResolutionStrategy,
    EntityRelationship,
    Cardinality,
    JoinPath,
    JoinStep,
    # Shared components
    ConfidenceScore,
    LowConfidenceItem,
    ConfidenceObjectType,
    GlossaryTerm,
    GlossaryScope,
    GlossaryCreator,
    # Serialization
    save_model,
)


def level_1_index_source() -> DataSource:
    """
    Simulate Level 1 agent indexing a data source.

    The Level 1 agent:
    1. Reads schema metadata from Trino
    2. Profiles tables and columns
    3. Infers semantic types
    4. Detects relationships
    5. Generates descriptions
    6. Computes confidence scores
    """

    # Create the data source
    source = DataSource(
        id="ecommerce-analytics",
        name="E-Commerce Analytics",
        trino_catalog="analytics",
        trino_schema="ecommerce",
        fully_qualified_prefix="analytics.ecommerce",
        source_type=SourceType.ANALYTICAL,
        description="Pre-aggregated e-commerce data for analytics and reporting",
        business_context="Primary source for revenue, order, and customer analytics",
        domain="E-Commerce",
        owner_team="Data Platform",
        owner_contact="data-platform@company.com",
    )

    # Index the orders table
    orders_table = Table(
        id="orders",
        name="orders",
        fully_qualified_name="analytics.ecommerce.orders",
        table_type=TableType.TABLE,
        semantic_role=SemanticRole.FACT,
        description="Fact table containing one row per order placed on the platform",
        business_purpose="Use for order-level analysis, revenue reporting, and customer behavior",
        usage_notes="For daily aggregates, prefer daily_order_summary for better performance",
        caveats=[
            "Refunded orders remain with status='refunded' - exclude for net revenue",
            "total_amount includes tax but excludes shipping",
        ],
        temporal_info=TemporalInfo(
            grain=TemporalGrain.EVENT,
            primary_timestamp_column="order_created_at",
            event_time_column="order_created_at",
            timezone="UTC",
            partition_column="order_created_at",
        ),
        primary_key=PrimaryKey(columns=["order_id"]),
        columns=[
            Column(
                id="order_id",
                name="order_id",
                ordinal_position=0,
                data_type="VARCHAR(36)",
                is_nullable=False,
                is_primary_key=True,
                semantic_type=SemanticType.identifier(subtype="uuid", confidence=0.95),
                description="Unique identifier for the order",
                business_meaning="System-generated UUID assigned when order is placed",
            ),
            Column(
                id="customer_id",
                name="customer_id",
                ordinal_position=1,
                data_type="VARCHAR(36)",
                is_nullable=True,
                is_foreign_key=True,
                semantic_type=SemanticType.identifier(subtype="uuid", confidence=0.9),
                description="ID of the customer who placed the order",
                business_meaning="Links to customer record. NULL for guest checkouts (~2%)",
                common_filters=["customer_id IS NOT NULL"],
            ),
            Column(
                id="order_status",
                name="order_status",
                ordinal_position=2,
                data_type="VARCHAR(20)",
                is_nullable=False,
                is_categorical=True,
                semantic_type=SemanticType(
                    category=SemanticCategory.CATEGORICAL,
                    confidence=0.95,
                ),
                description="Current status in the order fulfillment pipeline",
                allowed_values=[
                    AllowedValue(value="pending", description="Order placed, not yet processed", frequency=0.05),
                    AllowedValue(value="processing", description="Order being prepared", frequency=0.08),
                    AllowedValue(value="shipped", description="Order has left warehouse", frequency=0.12),
                    AllowedValue(value="delivered", description="Order confirmed delivered", frequency=0.65),
                    AllowedValue(value="refunded", description="Order was refunded", frequency=0.07),
                    AllowedValue(value="cancelled", description="Order cancelled before shipment", frequency=0.03),
                ],
                common_filters=[
                    "order_status NOT IN ('refunded', 'cancelled')",
                    "order_status = 'delivered'",
                ],
            ),
            Column(
                id="total_amount",
                name="total_amount",
                ordinal_position=3,
                data_type="DECIMAL(12,2)",
                is_nullable=False,
                semantic_type=SemanticType(
                    category=SemanticCategory.CURRENCY,
                    confidence=0.98,
                ),
                description="Total order value including tax, excluding shipping",
                unit="USD",
                profiling=ColumnProfiling(
                    min_value="0.01",
                    max_value="15000.00",
                    avg_value=127.50,
                    median_value=89.99,
                ),
            ),
            Column(
                id="order_created_at",
                name="order_created_at",
                ordinal_position=4,
                data_type="TIMESTAMP(6) WITH TIME ZONE",
                is_nullable=False,
                is_partition_key=True,
                semantic_type=SemanticType.timestamp(subtype="event_time", confidence=0.99),
                description="When the order was placed",
            ),
        ],
        confidence=ConfidenceScore(
            overall=0.92,
            threshold=0.8,
            schema_understanding=0.95,
            semantic_typing=0.90,
            description_quality=0.88,
        ),
    )

    # Index the customers table
    customers_table = Table(
        id="customers",
        name="customers",
        fully_qualified_name="analytics.ecommerce.customers",
        table_type=TableType.TABLE,
        semantic_role=SemanticRole.DIMENSION,
        description="Dimension table with customer master data and computed metrics",
        columns=[
            Column(
                id="customer_id",
                name="customer_id",
                ordinal_position=0,
                data_type="VARCHAR(36)",
                is_nullable=False,
                is_primary_key=True,
                semantic_type=SemanticType.identifier(subtype="uuid"),
                description="Unique customer identifier",
            ),
            Column(
                id="customer_name",
                name="customer_name",
                ordinal_position=1,
                data_type="VARCHAR(255)",
                is_nullable=False,
                semantic_type=SemanticType(category=SemanticCategory.TEXT, subtype="person_name"),
                description="Customer's full name",
            ),
            Column(
                id="email",
                name="email",
                ordinal_position=2,
                data_type="VARCHAR(255)",
                is_nullable=False,
                is_sensitive=True,
                semantic_type=SemanticType(category=SemanticCategory.TEXT, subtype="email"),
                description="Customer's email address",
            ),
            Column(
                id="lifetime_value",
                name="lifetime_value",
                ordinal_position=3,
                data_type="DECIMAL(12,2)",
                is_nullable=False,
                is_derived=True,
                semantic_type=SemanticType(category=SemanticCategory.CURRENCY),
                unit="USD",
                description="Total revenue from this customer (computed daily)",
                derivation_logic="SUM(orders.total_amount) WHERE order_status NOT IN ('refunded', 'cancelled')",
            ),
            Column(
                id="segment",
                name="segment",
                ordinal_position=4,
                data_type="VARCHAR(20)",
                is_nullable=False,
                is_categorical=True,
                semantic_type=SemanticType(category=SemanticCategory.CATEGORICAL),
                description="Customer segment based on behavior",
                allowed_values=[
                    AllowedValue(value="new", description="First purchase within 30 days"),
                    AllowedValue(value="active", description="Purchase within last 90 days"),
                    AllowedValue(value="at_risk", description="No purchase in 90-180 days"),
                    AllowedValue(value="churned", description="No purchase in 180+ days"),
                ],
            ),
        ],
        confidence=ConfidenceScore(overall=0.88, threshold=0.8),
    )

    # Add tables to source
    source = source.add_table(orders_table)
    source = source.add_table(customers_table)

    # Add detected relationship
    source.internal_relationships.append(
        InternalRelationship(
            id="orders_customers",
            name="Orders to Customers",
            description="Links each order to the customer who placed it",
            from_table_id="orders",
            from_column_id="customer_id",
            to_table_id="customers",
            to_column_id="customer_id",
            relationship_type=RelationshipType.ONE_TO_MANY,
            join_hint=JoinHint.LEFT,
            cardinality_notes="~98% of orders have a customer_id. Guest checkouts are NULL.",
            confidence=0.95,
        )
    )

    # Set overall source confidence
    source.confidence = ConfidenceScore(
        overall=0.90,
        threshold=0.8,
    )

    # Mark as confirmed (above threshold)
    source.status = source.status.confirm()

    return source


def level_2_create_entities(model: SemanticModel) -> SemanticModel:
    """
    Simulate Level 2 agent creating cross-source semantic entities.

    The Level 2 agent:
    1. Analyzes confirmed sources from Level 1
    2. Identifies canonical business entities
    3. Maps manifestations across sources
    4. Creates unified attributes
    5. Builds entity relationships with join paths
    6. Generates the global glossary
    """

    # Create the Customer entity
    customer_entity = SemanticEntity(
        id="customer",
        name="Customer",
        description="A customer is any individual or organization that has created an account or made a purchase on the platform",
        business_context="Central entity for all customer-related analytics, marketing, and support",
        domain="E-Commerce",
        canonical_id_name="customer_id",
        canonical_id_format="UUID",
        canonical_id_description="System-generated UUID from the e-commerce platform",
        manifestations=[
            EntityManifestation(
                source_id="ecommerce-analytics",
                table_id="customers",
                fully_qualified_name="analytics.ecommerce.customers",
                role=ManifestationRole.DERIVED,
                key_column_id="customer_id",
                key_mapping=KeyMapping(mapping_type=KeyMappingType.IDENTICAL),
                usage_guidance="PREFERRED for analytics. Pre-joined with computed metrics like LTV.",
                confidence=0.95,
            ),
        ],
        unified_attributes=[
            UnifiedAttribute(
                id="customer_name",
                name="Customer Name",
                description="Full name of the customer",
                semantic_type=SemanticType(category=SemanticCategory.TEXT, subtype="person_name"),
                sources=[
                    AttributeSource(
                        source_id="ecommerce-analytics",
                        table_id="customers",
                        column_id="customer_name",
                        fully_qualified_column="analytics.ecommerce.customers.customer_name",
                        priority=1,
                        freshness="daily",
                    ),
                ],
            ),
            UnifiedAttribute(
                id="customer_ltv",
                name="Customer Lifetime Value",
                description="Total revenue from this customer excluding refunds",
                semantic_type=SemanticType(category=SemanticCategory.CURRENCY),
                sources=[
                    AttributeSource(
                        source_id="ecommerce-analytics",
                        table_id="customers",
                        column_id="lifetime_value",
                        fully_qualified_column="analytics.ecommerce.customers.lifetime_value",
                        priority=1,
                        freshness="daily",
                    ),
                ],
            ),
        ],
        identity_resolution=IdentityResolution(
            strategy=IdentityResolutionStrategy.SINGLE_KEY,
            resolution_notes="All sources use the same customer_id UUID",
            confidence=0.95,
        ),
        common_queries=[
            "Get customer details with their LTV",
            "Find customers by segment",
            "Customer order history",
        ],
    )

    # Create the Order entity
    order_entity = SemanticEntity(
        id="order",
        name="Order",
        description="An order is a purchase transaction made by a customer",
        domain="E-Commerce",
        canonical_id_name="order_id",
        canonical_id_format="UUID",
        manifestations=[
            EntityManifestation(
                source_id="ecommerce-analytics",
                table_id="orders",
                fully_qualified_name="analytics.ecommerce.orders",
                role=ManifestationRole.PRIMARY,
                key_column_id="order_id",
                usage_guidance="Primary source for order-level analysis",
            ),
        ],
    )

    # Add entities to model
    model = model.add_entity(customer_entity)
    model = model.add_entity(order_entity)

    # Create relationship between Customer and Order
    customer_order_relationship = EntityRelationship(
        id="customer_orders",
        name="Customer Orders",
        description="Links customers to their orders",
        from_entity_id="customer",
        to_entity_id="order",
        cardinality=Cardinality.ONE_TO_MANY,
        join_paths=[
            JoinPath(
                id="direct_join",
                name="Direct Join",
                description="Direct join via customer_id",
                steps=[
                    JoinStep(
                        from_table="analytics.ecommerce.customers",
                        from_column="customer_id",
                        to_table="analytics.ecommerce.orders",
                        to_column="customer_id",
                    ),
                ],
                use_when="Standard customer-order analysis",
                freshness="Real-time",
            ),
        ],
        recommended_path_id="direct_join",
        confidence=0.95,
    )

    model.entity_relationships.append(customer_order_relationship)

    # Add glossary terms
    model = model.add_glossary_term(
        GlossaryTerm(
            id="ltv",
            name="LTV Definition",
            term="LTV",
            definition="Lifetime Value - Total revenue from a customer excluding refunds and cancellations",
            scope=GlossaryScope.GLOBAL,
            synonyms=["Lifetime Value", "Customer Lifetime Value", "CLV"],
            abbreviation="LTV",
            related_entities=["customer"],
            created_by=GlossaryCreator.LEVEL_2_AGENT,
            confidence=0.9,
        )
    )

    model = model.add_glossary_term(
        GlossaryTerm(
            id="churned",
            name="Churned Definition",
            term="Churned",
            definition="A customer who has not made a purchase in the last 180 days",
            scope=GlossaryScope.GLOBAL,
            related_entities=["customer"],
            created_by=GlossaryCreator.LEVEL_2_AGENT,
            confidence=0.85,
        )
    )

    return model


def main():
    """Run the complete example workflow."""

    # Create empty model
    model = create_empty_model(
        model_id="ecommerce-semantic-model",
        organization_id="acme-corp",
    )

    print("=" * 60)
    print("LEVEL 1: Indexing Data Source")
    print("=" * 60)

    # Level 1: Index source
    source = level_1_index_source()
    model = model.add_source(source)

    print(f"Indexed source: {source.name}")
    print(f"  Tables: {source.table_count}")
    print(f"  Confidence: {source.confidence.overall}")
    print(f"  Status: {source.status.state.value}")
    print()

    print("=" * 60)
    print("LEVEL 2: Creating Semantic Entities")
    print("=" * 60)

    # Level 2: Create entities
    model = level_2_create_entities(model)

    print(f"Created entities: {[e.name for e in model.entities]}")
    print(f"Created relationships: {len(model.entity_relationships)}")
    print(f"Glossary terms: {len(model.glossary)}")
    print()

    # Validate the model
    errors = model.validate_references()
    if errors:
        print("Validation errors:")
        for error in errors:
            print(f"  - {error}")
    else:
        print("Model validation: PASSED")
    print()

    # Generate prompt context
    print("=" * 60)
    print("LLM PROMPT CONTEXT")
    print("=" * 60)
    print(model.to_prompt_format())

    # Save the model
    save_model(model, "/tmp/semantic_model_example.json")
    print(f"\nModel saved to /tmp/semantic_model_example.json")


if __name__ == "__main__":
    main()
