import pytest
from sqlmodel import Session, create_engine, SQLModel, select
from sqldim import DimensionModel, FactModel, Field, SCD2Mixin
from sqldim.loaders.dimensional import DimensionalLoader

class ProductDim(DimensionModel, SCD2Mixin, table=True):
    __natural_key__ = ["sku"]
    id: int = Field(primary_key=True, surrogate_key=True)
    sku: str
    price: float

class SaleFact(FactModel, table=True):
    id: int = Field(primary_key=True)
    product_id: int = Field(foreign_key="productdim.id", dimension=ProductDim)
    quantity: int

@pytest.fixture
def session():
    engine = create_engine("sqlite:///:memory:")
    SQLModel.metadata.create_all(engine)
    with Session(engine) as session:
        yield session

@pytest.mark.asyncio
async def test_dimensional_load_with_sk_resolution(session):
    loader = DimensionalLoader(session, models=[ProductDim, SaleFact])
    
    # Register Dimension data
    products = [{"sku": "P1", "price": 10.0}]
    loader.register(ProductDim, products)
    
    # Register Fact data using Natural Key "P1" for "product_id"
    sales = [{"product_id": "P1", "quantity": 5}]
    loader.register(
        SaleFact, 
        sales, 
        key_map={"product_id": (ProductDim, "sku")}
    )
    
    # Run load
    await loader.run()
    
    # Verify SK Resolution
    product = session.exec(select(ProductDim).where(ProductDim.sku == "P1")).one()
    sale = session.exec(select(SaleFact)).one()
    
    assert sale.product_id == product.id
    assert sale.quantity == 5

@pytest.mark.asyncio
async def test_sk_resolver_cache(session):
    from sqldim.loaders.dimensional import SKResolver
    
    # Setup dimension
    p = ProductDim(sku="P2", price=20.0, is_current=True)
    session.add(p)
    session.commit()
    
    resolver = SKResolver(session)
    sk1 = resolver.resolve(ProductDim, "sku", "P2")
    assert sk1 == p.id
    
    # Manually change DB without clearing cache to verify cache is hit
    # (In real use, cache is per-load-run)
    assert (ProductDim, "sku", "P2") in resolver._cache
    sk2 = resolver.resolve(ProductDim, "sku", "P2")
    assert sk2 == sk1

@pytest.mark.asyncio
async def test_sk_resolution_missing_member(session):
    loader = DimensionalLoader(session, models=[ProductDim, SaleFact])
    
    # Register Fact data using a Natural Key that doesn't exist in the Dimension
    sales = [{"product_id": "MISSING", "quantity": 1}]
    loader.register(
        SaleFact, 
        sales, 
        key_map={"product_id": (ProductDim, "sku")}
    )
    
    # Run load - should not crash during SK resolution, but may fail during DB commit 
    # if types mismatch. In SQLite, it might actually store the string "MISSING" 
    # unless strict mode is on.
    await loader.run()
    
    # Verify that it stayed as "MISSING" because resolution failed
    sale = session.exec(select(SaleFact)).one()
    assert sale.product_id == "MISSING"

@pytest.mark.asyncio
async def test_loader_strategies(session):
    # Triggers the remaining strategy branches in loaders/dimensional.py
    class StratFact(FactModel, table=True):
        __strategy__ = "bulk"
        id: int = Field(primary_key=True)
        val: int

    StratFact.__table__.create(session.bind)
    loader = DimensionalLoader(session, models=[StratFact])
    loader.register(StratFact, [{"id": 1, "val": 10}, {"id": 2, "val": 20}])
    await loader.run()
    assert len(session.exec(select(StratFact)).all()) == 2

    class UpsertFact(FactModel, table=True):
        __strategy__ = "upsert"
        __natural_key__ = ["id"]
        id: int = Field(primary_key=True)
        val: int = Field(default=0)
    
    UpsertFact.__table__.create(session.bind)
    loader2 = DimensionalLoader(session, models=[UpsertFact])
    loader2.register(UpsertFact, [{"id": 1, "val": 100}])
    await loader2.run()
    assert session.exec(select(UpsertFact)).one().val == 100

    class MergeFact(FactModel, table=True):
        __strategy__ = "merge"
        __natural_key__ = ["id"]
        id: int = Field(primary_key=True)
        val: int = Field(default=0)
    
    MergeFact.__table__.create(session.bind)
    loader3 = DimensionalLoader(session, models=[MergeFact])
    loader3.register(MergeFact, [{"id": 1, "val": 1000}])
    await loader3.run()
    assert session.exec(select(MergeFact)).one().val == 1000
