import os
import sys
from chorm import Table, Column, MaterializedView, MergeTree, create_engine, Session, select, MetaData
from chorm.sql.expression import func
from chorm.types import UInt64, String
from chorm.introspection import TableIntrospector, ModelGenerator

# 1. Setup Models
# Use a shared MetaData instance
metadata = MetaData()

class UserEvents(Table):
    """Source table (Table 1)."""
    __tablename__ = "user_events"
    __engine__ = MergeTree()
    __order_by__ = ("event_id",)
    metadata = metadata
    
    event_id = Column(UInt64())
    user_id = Column(UInt64())
    event_type = Column(String())

class UserStats(Table):
    """Target table (Table 2). Aggregated data."""
    __tablename__ = "user_stats"
    __engine__ = MergeTree() # Standard storage for the aggregation
    __order_by__ = ("user_id",)
    metadata = metadata
    
    user_id = Column(UInt64())
    count = Column(UInt64())

class UserEventsMV(Table):
    """Materialized View moving data from Events to Stats."""
    __tablename__ = "user_events_mv"
    # ClickHouse MV: "TO user_stats" means it pushes results into user_stats table
    __engine__ = MaterializedView(to_table="user_stats")
    metadata = metadata
    
    # The transformation query.
    # Note: In ClickHouse MVs with TO table, the aggregation logic must generate columns matching the Target table.
    __select__ = select(
        UserEvents.user_id,
        func.count().label("count")
    ).select_from(UserEvents).group_by(UserEvents.user_id)


def test_mv_lifecycle():
    print("\n=== Starting MV Lifecycle Integration Test (with Metadata) ===")
    
    # Connect
    dsn = os.getenv("CLICKHOUSE_DSN", "clickhouse://default:@localhost:8123/default")
    engine = create_engine(dsn)
    session = Session(engine)
    
    # 2. Cleanup & Create using Metadata
    print("\n[Step 1] Creating Tables via MetaData...")
    # Clean first
    metadata.drop_all(engine)
    
    # Create all
    # Note: dictionary insertion order in 'metadata.tables' should preserve definition order.
    # UserEvents -> UserStats -> UserEventsMV
    # However, strictly speaking, UserEventsMV depends on UserStats (TO table) and UserEvents (SELECT FROM).
    # If order is wrong, it might fail. But in Python 3.7+ dicts preserve insertion order,
    # and we defined classes in order.
    metadata.create_all(engine)
        
    print("✓ Tables created via metadata.create_all().")

    # 3. Introspection (Requested check BEFORE data insertion)
    print("\n[Step 2] Verifying Introspection...")
    with engine.connect() as conn:
        introspector = TableIntrospector(conn.client)
        generator = ModelGenerator()
        
        print("  Introspecting 'user_events_mv'...")
        info = introspector.get_table_info("user_events_mv")
        code = generator.generate_model(info)
        
        print("  Generated Code snippet:")
        print("-" * 40)
        print(code)
        print("-" * 40)
        
        # Verify engine type and arguments
        assert "MaterializedView" in code
        assert 'to_table="default.user_stats"' in code or 'to_table="user_stats"' in code
        assert "__select__" in code
        assert "SELECT" in code
        # Check if column presence is introspected (MVs often don't expose columns in system.columns easily or match target)
        # But 'system.tables' is where we get the metadata.
        print("✓ Introspection looks correct.")

    # 4. Insert Data into Source
    print("\n[Step 3] Inserting Data into Source (user_events)...")
    events = [
        UserEvents(event_id=1, user_id=101, event_type="login"),
        UserEvents(event_id=2, user_id=101, event_type="click"),
        UserEvents(event_id=3, user_id=102, event_type="login"),
        UserEvents(event_id=4, user_id=101, event_type="logout"),
    ]
    for e in events:
        session.add(e)
    session.commit()
    print(f"✓ Inserted {len(events)} events.")

    # 5. Retrieve Data from Target
    print("\n[Step 4] Retrieving Data from Target (user_stats)...")
    # MV trigger should have populated user_stats
    # User 101: 3 events
    # User 102: 1 event
    
    stmt = select(UserStats.user_id, UserStats.count).order_by(UserStats.user_id)
    results = session.execute(stmt).all()
    
    print("  Results from user_stats:", results)
    
    # Verification
    # Note: Results format depends on result proxy, likely tuples or Rows
    user_counts = {row.user_id: row.count for row in results}
    
    assert user_counts.get(101) == 3, f"Expected user 101 to have 3 events, got {user_counts.get(101)}"
    assert user_counts.get(102) == 1, f"Expected user 102 to have 1 event, got {user_counts.get(102)}"
    
    print("✓ Data verification passed.")
    
    # Final Cleanup
    print("\n[Step 5] Cleaning up...")
    with engine.connect() as conn:
        conn.execute("DROP TABLE IF EXISTS user_events_mv")
        conn.execute("DROP TABLE IF EXISTS user_stats")
        conn.execute("DROP TABLE IF EXISTS user_events")
    print("✓ Done.")

if __name__ == "__main__":
    test_mv_lifecycle()
