from pydantic import BaseModel
from typing import List
from pydantic_resolve import Resolver, Loader
from aiodataloader import DataLoader
import pytest
from pydantic_resolve.exceptions import GlobalLoaderFieldOverlappedError

class LoaderA(DataLoader):
    power: int
    async def batch_load_fn(self, keys: List[int]):
        return [ k** self.power for k in keys ]

class LoaderB(DataLoader):
    power: int
    async def batch_load_fn(self, keys: List[int]):
        return [ k** self.power for k in keys ]

class LoaderC(DataLoader):
    power: int
    add: int 
    async def batch_load_fn(self, keys: List[int]):
        return [ k** self.power + self.add for k in keys ]


async def loader_fn_a(keys):
    return [ k**2 for k in keys ]

class A(BaseModel):
    val: int

    a: int = 0
    def resolve_a(self, loader=Loader(LoaderA)):
        return loader.load(self.val)

    b: int = 0
    def resolve_b(self, loader=Loader(LoaderB)):
        return loader.load(self.val)

    c: int = 0
    def resolve_c(self, loader=Loader(LoaderC)):
        return loader.load(self.val)


@pytest.mark.asyncio
async def test_case_0():
    data = [A(val=n) for n in range(3)]
    data = await Resolver(global_loader_param={'power': 2}, 
                          loader_params={LoaderC:{'add': 1}}).resolve(data)
    assert data[2].a == 4
    assert data[2].b == 4
    assert data[2].c == 5


@pytest.mark.asyncio
async def test_case_1():
    data = [A(val=n) for n in range(3)]
    data = await Resolver(loader_params={LoaderA:{'power': 2}, 
                                          LoaderB:{'power': 3},
                                          LoaderC:{'power': 3, 'add': 1}}).resolve(data)
    assert data[2].a == 4
    assert data[2].b == 8
    assert data[2].c == 9


@pytest.mark.asyncio
async def test_case_3():
    data = [A(val=n) for n in range(3)]
    with pytest.raises(GlobalLoaderFieldOverlappedError):
        data = await Resolver(global_loader_param={'power': 2}, 
                            loader_params={LoaderC:{'add': 1, 'power': 3}}).resolve(data)