from django.core.exceptions import ImproperlyConfigured
from django.db import models

from polymodels.managers import PolymorphicManager

from .base import TestCase
from .models import Animal, BigSnake, HugeSnake, Mammal, Monkey, Snake, Zoo


class PolymorphicQuerySetTest(TestCase):
    def test_select_subclasses(self):
        Animal.objects.create(name="animal")
        Mammal.objects.create(name="mammal")
        Monkey.objects.create(name="monkey")
        Snake.objects.create(name="snake", length=10)
        BigSnake.objects.create(name="big snake", length=101)
        HugeSnake.objects.create(name="huge snake", length=155)
        # Assert `select_subclasses` correctly calls `select_related` and `filter`.
        animals = Animal.objects.select_subclasses()
        animals_expected_query_select_related = {
            "mammal": {"monkey": {}},
            "snake": {},
        }
        self.assertEqual(
            animals.query.select_related, animals_expected_query_select_related
        )
        with self.assertNumQueries(1):
            self.assertQuerySetEqual(
                animals.all(),
                [
                    "<Animal: animal>",
                    "<Mammal: mammal>",
                    "<Monkey: monkey>",
                    "<Snake: snake>",
                    "<BigSnake: big snake>",
                    "<HugeSnake: huge snake>",
                ],
                transform=repr,
            )
        with self.assertNumQueries(1):
            self.assertQuerySetEqual(
                list(animals.iterator()),
                [
                    "<Animal: animal>",
                    "<Mammal: mammal>",
                    "<Monkey: monkey>",
                    "<Snake: snake>",
                    "<BigSnake: big snake>",
                    "<HugeSnake: huge snake>",
                ],
                transform=repr,
            )
        # Filter out non-mammal (direct subclass)
        animal_mammals = Animal.objects.select_subclasses(Mammal)
        animal_mammals_expected_query_select_related = {"mammal": {"monkey": {}}}
        self.assertEqual(
            animal_mammals.query.select_related,
            animal_mammals_expected_query_select_related,
        )
        with self.assertNumQueries(1):
            self.assertQuerySetEqual(
                animal_mammals.all(),
                ["<Mammal: mammal>", "<Monkey: monkey>"],
                transform=repr,
            )
        # Filter out non-snake (subclass through an abstract one)
        animal_snakes = Animal.objects.select_subclasses(Snake)
        self.assertEqual(animal_snakes.query.select_related, {"snake": {}})
        with self.assertNumQueries(1):
            self.assertQuerySetEqual(
                animal_snakes.all(),
                ["<Snake: snake>", "<BigSnake: big snake>", "<HugeSnake: huge snake>"],
                transform=repr,
            )
        # Subclass with only proxies
        snakes = Snake.objects.select_subclasses()
        self.assertFalse(snakes.query.select_related)
        with self.assertNumQueries(1):
            self.assertQuerySetEqual(
                snakes.all(),
                ["<Snake: snake>", "<BigSnake: big snake>", "<HugeSnake: huge snake>"],
                transform=repr,
            )
        # Subclass filter proxies
        snake_bigsnakes = Snake.objects.select_subclasses(BigSnake)
        self.assertFalse(snakes.query.select_related)
        with self.assertNumQueries(1):
            self.assertQuerySetEqual(
                snake_bigsnakes.all(),
                ["<BigSnake: big snake>", "<HugeSnake: huge snake>"],
                transform=repr,
            )

    def test_select_subclasses_get(self):
        snake = Snake.objects.create(name="snake", length=10)
        self.assertEqual(Animal.objects.select_subclasses().get(), snake)

    def test_select_subclasses_values(self):
        Animal.objects.create(name="animal")
        self.assertQuerySetEqual(
            Animal.objects.select_subclasses().values_list("name", flat=True),
            ["animal"],
            lambda x: x,
        )

    def test_exclude_subclasses(self):
        Animal.objects.create(name="animal")
        Mammal.objects.create(name="first mammal")
        Mammal.objects.create(name="second mammal")
        Monkey.objects.create(name="donkey kong")
        self.assertQuerySetEqual(
            Animal.objects.exclude_subclasses(),
            ["<Animal: animal>"],
            transform=repr,
        )
        self.assertQuerySetEqual(
            Mammal.objects.exclude_subclasses(),
            ["<Mammal: first mammal>", "<Mammal: second mammal>"],
            transform=repr,
        )
        self.assertQuerySetEqual(
            Monkey.objects.exclude_subclasses(),
            ["<Monkey: donkey kong>"],
            transform=repr,
        )

    def test_select_subclasses_prefetch_related(self):
        zoo = Zoo.objects.create()
        animal = Animal.objects.create(name="animal")
        mammal = Mammal.objects.create(name="mammal")
        monkey = Monkey.objects.create(name="monkey")
        zoo.animals.add(animal, mammal, monkey)
        other_monkey = Monkey.objects.create(name="monkey")
        monkey.friends.add(other_monkey)
        queryset = Animal.objects.select_subclasses().prefetch_related("zoos")
        with self.assertNumQueries(2):
            self.assertSequenceEqual(
                queryset,
                [
                    animal,
                    mammal,
                    monkey,
                    other_monkey,
                ],
            )
            self.assertSequenceEqual(queryset[0].zoos.all(), [zoo])
            self.assertSequenceEqual(queryset[1].zoos.all(), [zoo])
            self.assertSequenceEqual(queryset[2].zoos.all(), [zoo])
        # Test prefetch related combination.
        queryset = Animal.objects.select_subclasses().prefetch_related(
            "zoos",
            "mammal__monkey__friends",
        )
        with self.assertNumQueries(3):
            self.assertSequenceEqual(
                queryset,
                [
                    animal,
                    mammal,
                    monkey,
                    other_monkey,
                ],
            )
            self.assertSequenceEqual(queryset[0].zoos.all(), [zoo])
            self.assertSequenceEqual(queryset[1].zoos.all(), [zoo])
            self.assertSequenceEqual(queryset[2].zoos.all(), [zoo])
            self.assertSequenceEqual(queryset[2].friends.all(), [other_monkey])
            self.assertSequenceEqual(queryset[3].friends.all(), [monkey])


class PolymorphicManagerTest(TestCase):
    def test_improperly_configured(self):
        with self.assertRaisesMessage(
            ImproperlyConfigured,
            "`PolymorphicManager` can only be used on `BasePolymorphicModel` subclasses.",
        ):

            class NonPolymorphicModel(models.Model):
                objects = PolymorphicManager()

                class Meta:
                    app_label = "polymodels"

    def test_proxy_filtering(self):
        """
        Make sure managers attached to proxy models returns a queryset of
        proxies only.
        """
        Snake.objects.create(name="snake", length=1)
        BigSnake.objects.create(name="big snake", length=10)
        HugeSnake.objects.create(name="huge snake", length=100)
        self.assertQuerySetEqual(
            Snake.objects.all(),
            ["<Snake: snake>", "<Snake: big snake>", "<Snake: huge snake>"],
            transform=repr,
        )
        self.assertQuerySetEqual(
            BigSnake.objects.all(),
            ["<BigSnake: big snake>", "<BigSnake: huge snake>"],
            transform=repr,
        )
        self.assertQuerySetEqual(
            HugeSnake.objects.all(),
            ["<HugeSnake: huge snake>"],
            transform=repr,
        )
