import pytest
from typing import Optional, Annotated, List
from pydantic import BaseModel
from pydantic_resolve.utils.resolver_configurator import config_global_resolver
from pydantic_resolve import Entity, Relationship, MultipleRelationship, Link, LoadBy, DefineSubset, ErDiagram, Resolver
from aiodataloader import DataLoader



class Biz(BaseModel):
    id: int
    name: str
    user_id: int

class Foo(BaseModel):
    id: int
    name: str
    biz_id: int

class Bar(BaseModel):
    id: int
    name: str
    biz_id: int

class User(BaseModel):
    id: int
    name: str

class UserLoader(DataLoader):
    async def batch_load_fn(self, keys):
        users = [
            dict(id=1, name="a"),
            dict(id=2, name="b"),
            dict(id=3, name="c"),
        ]
        user_map = { u['id']: u for u in users}
        return [user_map.get(k, None) for k in keys]

class BarLoader(DataLoader):
    async def batch_load_fn(self, keys):
        bars = [
            dict(id=1, name="bar1", biz_id=1),
            dict(id=2, name="bar2", biz_id=1),
            dict(id=3, name="bar3", biz_id=2),
        ]
        bar_map = {}
        for b in bars:
            bar_map.setdefault(b['biz_id'], []).append(b)
        return [bar_map.get(k, []) for k in keys]

class SpecialBarLoader(DataLoader):
    async def batch_load_fn(self, keys):
        bars = [
            dict(id=1, name="special-bar1", biz_id=1),
            dict(id=2, name="special-bar2", biz_id=1),
            dict(id=3, name="special-bar3", biz_id=2),
        ]
        bar_map = {}
        for b in bars:
            bar_map.setdefault(b['biz_id'], []).append(b)
        return [bar_map.get(k, []) for k in keys]

class FooLoader(DataLoader):
    async def batch_load_fn(self, keys):
        foos = [
            dict(id=1, name="foo1", biz_id=1),
            dict(id=2, name="foo2", biz_id=1),
            dict(id=3, name="foo3", biz_id=2),
        ]
        foo_map = {}
        for f in foos:
            foo_map.setdefault(f['biz_id'], []).append(f)
        return [foo_map.get(k, []) for k in keys]

diagram = ErDiagram(
    configs=[
        Entity(kls=Biz, relationships=[
            Relationship(field='user_id', target_kls=User, loader=UserLoader),
            Relationship(field='id', target_kls=List[Foo], loader=FooLoader),
            MultipleRelationship(
                field='id',
                target_kls=list[Bar],
                links=[
                    Link(biz='a', loader=BarLoader),
                    Link(biz='special', loader=SpecialBarLoader),
                ]
            )
        ])
    ]
)


@pytest.fixture(autouse=True)
def setup_global_resolver():
    """Set up global resolver for this test module."""
    config_global_resolver(er_diagram=diagram)
    yield

class BizCase1(DefineSubset):
    __pydantic_resolve_subset__ = (Biz, ('id', 'name', 'user_id'))

    user: Annotated[Optional[User], LoadBy('user_id')] = None
    foos: Annotated[List[Foo], LoadBy('id')] = []
    bars: Annotated[List[Bar], LoadBy('id', biz='a')] = []
    special_bars: Annotated[list[Bar], LoadBy('id', biz='special')] = []
    

@pytest.mark.asyncio
async def test_resolver_factory_with_er_configs_inherit():
    d = [BizCase1(id=1, name="qq", user_id=1), BizCase1(id=2, name="ww", user_id=2)]
    d = await Resolver().resolve(d)

    assert d[0].user.name == "a"
    assert d[0].bars == [Bar(id=1, name="bar1", biz_id=1), Bar(id=2, name="bar2", biz_id=1)]
    assert d[0].special_bars == [Bar(id=1, name="special-bar1", biz_id=1), Bar(id=2, name="special-bar2", biz_id=1)]

    assert d[1].user.name == "b"
    assert d[1].foos == [Foo(id=3, name="foo3", biz_id=2)]