# This file was auto-generated by the framework. You can add data models or change the existing ones

from moose_lib import Key, IngestPipeline, IngestPipelineConfig, StringToEnumMixin, clickhouse_default, OlapTable, OlapConfig, MergeTreeEngine, ReplacingMergeTreeEngine
from datetime import datetime
from typing import Optional, Annotated, Any
from pydantic import BaseModel, BeforeValidator
from enum import IntEnum, auto



class Baz(StringToEnumMixin, IntEnum):
    QUX = auto()
    QUUX = auto()


class Foo(BaseModel):
    primary_key: Key[str]
    timestamp: float
    baz: Baz
    optional_text: Optional[str] = None


class Bar(BaseModel):
    primary_key: Key[str]
    utc_timestamp: datetime
    baz: Baz
    has_text: bool
    text_length: int


fooModel = IngestPipeline[Foo]("Foo", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=False,
    dead_letter_queue=True
))

barModel = IngestPipeline[Bar]("Bar", IngestPipelineConfig(
    ingest_api=False,
    stream=True,
    table=True,
    dead_letter_queue=True
))

# =======Comprehensive Type Testing Data Models=========

from typing import List, Tuple
from pydantic import Field

# Test 1: Basic primitive types
class BasicTypes(BaseModel):
    id: Key[str]
    timestamp: datetime
    string_field: str
    number_field: float
    boolean_field: bool
    optional_string: Optional[str] = None
    nullable_number: Optional[float] = None

# Test 2: Simple arrays of primitives
class SimpleArrays(BaseModel):
    id: Key[str]
    timestamp: datetime
    string_array: List[str]
    number_array: List[float]
    boolean_array: List[bool]
    optional_string_array: Optional[List[str]] = None
    mixed_optional_array: Optional[List[str]] = None

# Test 3: Nested objects
class Coordinates(BaseModel):
    lat: float
    lng: float

class Address(BaseModel):
    street: str
    city: str
    coordinates: Coordinates

class Settings(BaseModel):
    theme: str
    notifications: bool

class Config(BaseModel):
    enabled: bool
    settings: Settings

class Metadata(BaseModel):
    tags: List[str]
    priority: int
    config: Config

class NestedObjects(BaseModel):
    id: Key[str]
    timestamp: datetime
    address: Address
    metadata: Metadata

# Test 4: Arrays of objects
class User(BaseModel):
    name: str
    age: int
    active: bool

class TransactionMetadata(BaseModel):
    category: str
    tags: List[str]

class Transaction(BaseModel):
    id: str
    amount: float
    currency: str
    metadata: TransactionMetadata

class ArraysOfObjects(BaseModel):
    id: Key[str]
    timestamp: datetime
    users: List[User]
    transactions: List[Transaction]

# Test 5: Deeply nested arrays (main focus for ENG-875) - ClickHouse compatible
class SimplifiedItem(BaseModel):
    name: str
    values: List[float]
    data: List[str]
    # Flattened structure instead of deeply nested objects
    metric_names: List[str]
    metric_values: List[float]

class ComplexNestedCategory(BaseModel):
    category: str
    items: List[SimplifiedItem]

class DeeplyNestedArrays(BaseModel):
    id: Key[str]
    timestamp: datetime
    # Level 1: Array of arrays (safe)
    matrix_2d: List[List[float]]
    # Level 2: Array of arrays of arrays (safe)
    matrix_3d: List[List[List[float]]]
    # Level 3: Array of arrays of arrays of arrays (pushing limits but should work)
    matrix_4d: List[List[List[List[float]]]]
    # Simplified nested: Reduce nesting depth to avoid ClickHouse issues
    complex_nested: List[ComplexNestedCategory]

# Test 6: Mixed complex types (ClickHouse-safe)
class ClickEvent(BaseModel):
    type: str = Field(default="click")
    target: str
    coordinate_x: float
    coordinate_y: float

class FlattenedData(BaseModel):
    required: str
    optional_data: List[str]  # Flattened, no nested optionals
    tags: List[str]
    values: List[float]

class Column(BaseModel):
    col: int
    values: List[str]

class ComplexMatrixRow(BaseModel):
    row: int
    columns: List[Column]

class MixedComplexTypes(BaseModel):
    id: Key[str]
    timestamp: datetime
    # Arrays with click events (safe)
    events: List[ClickEvent]
    # Flattened optional structures (avoid nested optionals)
    nested_data: List[FlattenedData]
    # Multi-dimensional with objects (safe)
    complex_matrix: List[ComplexMatrixRow]

# Test 7: Edge cases and boundary conditions (ClickHouse-compatible)
class EmptyObject(BaseModel):
    id: str

class Property(BaseModel):
    key: str
    value: str
    tags: List[str]

class Metric(BaseModel):
    name: str
    values: List[float]

class ModerateLevel2(BaseModel):
    data: List[str]
    values: List[float]

class ModerateLevel1(BaseModel):
    level2: List[ModerateLevel2]

class ModerateNesting(BaseModel):
    level1: List[ModerateLevel1]

class SimplifiedComplexItem(BaseModel):
    id: str
    properties: List[Property]
    metrics: List[Metric]

class EdgeCases(BaseModel):
    id: Key[str]
    timestamp: datetime
    # Empty arrays (safe)
    empty_string_array: List[str]
    empty_object_array: List[EmptyObject]
    # Simplified nullable structures (avoid nested nullables)
    nullable_string: Optional[str] = None
    nullable_number: Optional[float] = None
    # Moderate depth nesting (3 levels max for safety)
    moderate_nesting: ModerateNesting
    # Simplified complex arrays
    complex_array: List[SimplifiedComplexItem]

# =======Pipeline Configurations for Test Models=========

basic_types_model = IngestPipeline[BasicTypes]("BasicTypes", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

simple_arrays_model = IngestPipeline[SimpleArrays]("SimpleArrays", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

nested_objects_model = IngestPipeline[NestedObjects]("NestedObjects", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

arrays_of_objects_model = IngestPipeline[ArraysOfObjects]("ArraysOfObjects", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

deeply_nested_arrays_model = IngestPipeline[DeeplyNestedArrays]("DeeplyNestedArrays", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

mixed_complex_types_model = IngestPipeline[MixedComplexTypes]("MixedComplexTypes", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

edge_cases_model = IngestPipeline[EdgeCases]("EdgeCases", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

# =======Optional Nested Fields with ClickHouse Defaults Test=========

# Test class with optional nested fields and ClickHouse defaults
class TestNested(BaseModel):
    name: Optional[str] = None
    age: Optional[float] = None

class OptionalNestedTest(BaseModel):
    id: Key[str]
    timestamp: datetime
    nested: List[TestNested]
    other: Optional[Annotated[str, clickhouse_default("''")]] = None

optional_nested_test_model = IngestPipeline[OptionalNestedTest]("OptionalNestedTest", IngestPipelineConfig(
    ingest_api=True,
    stream=True,
    table=True,
    dead_letter_queue=True
))

# =======Versioned OlapTables Test=========
# Test versioned OlapTables - same name, different versions
# This demonstrates the OlapTable versioning functionality

class UserEventV1(BaseModel):
    """Version 1.0 of user events - basic structure."""
    user_id: Key[str]
    event_type: str
    timestamp: float
    metadata: Optional[str] = None

class UserEventV2(BaseModel):
    """Version 2.0 of user events - enhanced with session tracking."""
    user_id: Key[str]
    event_type: str
    timestamp: float
    metadata: Optional[str] = None
    session_id: str
    user_agent: Optional[str] = None

# Version 1.0 - MergeTree engine
user_events_v1 = OlapTable[UserEventV1](
    "UserEvents",
    OlapConfig(
        version="1.0",
        engine=MergeTreeEngine(),
        order_by_fields=["user_id", "timestamp"]
    )
)

# Version 2.0 - ReplacingMergeTree engine with enhanced schema
user_events_v2 = OlapTable[UserEventV2](
    "UserEvents",
    OlapConfig(
        version="2.0",
        engine=ReplacingMergeTreeEngine(),
        order_by_fields=["user_id", "session_id", "timestamp"]
    )
)
