from __future__ import unicode_literals

from django.contrib.contenttypes.models import ContentType
from django.core.exceptions import ObjectDoesNotExist
from django.db import connection
from django.db.models import Prefetch
from django.db.models.query import get_prefetcher
from django.test import TestCase, override_settings
from django.utils import six
from django.utils.encoding import force_text

from .models import (
    Author, Author2, AuthorAddress, AuthorWithAge, Bio, Book, Bookmark,
    BookReview, BookWithYear, Comment, Department, Employee, FavoriteAuthors,
    House, LessonEntry, Person, Qualification, Reader, Room, TaggedItem,
    Teacher, WordEntry,
)


class PrefetchRelatedTests(TestCase):
    def setUp(self):
        self.book1 = Book.objects.create(title="Poems")
        self.book2 = Book.objects.create(title="Jane Eyre")
        self.book3 = Book.objects.create(title="Wuthering Heights")
        self.book4 = Book.objects.create(title="Sense and Sensibility")

        self.author1 = Author.objects.create(name="Charlotte",
                                             first_book=self.book1)
        self.author2 = Author.objects.create(name="Anne",
                                             first_book=self.book1)
        self.author3 = Author.objects.create(name="Emily",
                                             first_book=self.book1)
        self.author4 = Author.objects.create(name="Jane",
                                             first_book=self.book4)

        self.book1.authors.add(self.author1, self.author2, self.author3)
        self.book2.authors.add(self.author1)
        self.book3.authors.add(self.author3)
        self.book4.authors.add(self.author4)

        self.reader1 = Reader.objects.create(name="Amy")
        self.reader2 = Reader.objects.create(name="Belinda")

        self.reader1.books_read.add(self.book1, self.book4)
        self.reader2.books_read.add(self.book2, self.book4)

    def test_m2m_forward(self):
        with self.assertNumQueries(2):
            lists = [list(b.authors.all()) for b in Book.objects.prefetch_related('authors')]

        normal_lists = [list(b.authors.all()) for b in Book.objects.all()]
        self.assertEqual(lists, normal_lists)

    def test_m2m_reverse(self):
        with self.assertNumQueries(2):
            lists = [list(a.books.all()) for a in Author.objects.prefetch_related('books')]

        normal_lists = [list(a.books.all()) for a in Author.objects.all()]
        self.assertEqual(lists, normal_lists)

    def test_foreignkey_forward(self):
        with self.assertNumQueries(2):
            books = [a.first_book for a in Author.objects.prefetch_related('first_book')]

        normal_books = [a.first_book for a in Author.objects.all()]
        self.assertEqual(books, normal_books)

    def test_foreignkey_reverse(self):
        with self.assertNumQueries(2):
            [list(b.first_time_authors.all())
             for b in Book.objects.prefetch_related('first_time_authors')]

        self.assertQuerysetEqual(self.book2.authors.all(), ["<Author: Charlotte>"])

    def test_onetoone_reverse_no_match(self):
        # Regression for #17439
        with self.assertNumQueries(2):
            book = Book.objects.prefetch_related('bookwithyear').all()[0]
        with self.assertNumQueries(0):
            with self.assertRaises(BookWithYear.DoesNotExist):
                book.bookwithyear

    def test_survives_clone(self):
        with self.assertNumQueries(2):
            [list(b.first_time_authors.all())
             for b in Book.objects.prefetch_related('first_time_authors').exclude(id=1000)]

    def test_len(self):
        with self.assertNumQueries(2):
            qs = Book.objects.prefetch_related('first_time_authors')
            len(qs)
            [list(b.first_time_authors.all()) for b in qs]

    def test_bool(self):
        with self.assertNumQueries(2):
            qs = Book.objects.prefetch_related('first_time_authors')
            bool(qs)
            [list(b.first_time_authors.all()) for b in qs]

    def test_count(self):
        with self.assertNumQueries(2):
            qs = Book.objects.prefetch_related('first_time_authors')
            [b.first_time_authors.count() for b in qs]

    def test_exists(self):
        with self.assertNumQueries(2):
            qs = Book.objects.prefetch_related('first_time_authors')
            [b.first_time_authors.exists() for b in qs]

    def test_in_and_prefetch_related(self):
        """
        Regression test for #20242 - QuerySet "in" didn't work the first time
        when using prefetch_related. This was fixed by the removal of chunked
        reads from QuerySet iteration in
        70679243d1786e03557c28929f9762a119e3ac14.
        """
        qs = Book.objects.prefetch_related('first_time_authors')
        self.assertIn(qs[0], qs)

    def test_clear(self):
        """
        Test that we can clear the behavior by calling prefetch_related()
        """
        with self.assertNumQueries(5):
            with_prefetch = Author.objects.prefetch_related('books')
            without_prefetch = with_prefetch.prefetch_related(None)
            [list(a.books.all()) for a in without_prefetch]

    def test_m2m_then_m2m(self):
        """
        Test we can follow a m2m and another m2m
        """
        with self.assertNumQueries(3):
            qs = Author.objects.prefetch_related('books__read_by')
            lists = [[[six.text_type(r) for r in b.read_by.all()]
                      for b in a.books.all()]
                     for a in qs]
            self.assertEqual(lists,
            [
                [["Amy"], ["Belinda"]],  # Charlotte - Poems, Jane Eyre
                [["Amy"]],                # Anne - Poems
                [["Amy"], []],            # Emily - Poems, Wuthering Heights
                [["Amy", "Belinda"]],    # Jane - Sense and Sense
            ])

    def test_overriding_prefetch(self):
        with self.assertNumQueries(3):
            qs = Author.objects.prefetch_related('books', 'books__read_by')
            lists = [[[six.text_type(r) for r in b.read_by.all()]
                      for b in a.books.all()]
                     for a in qs]
            self.assertEqual(lists,
            [
                [["Amy"], ["Belinda"]],  # Charlotte - Poems, Jane Eyre
                [["Amy"]],                # Anne - Poems
                [["Amy"], []],            # Emily - Poems, Wuthering Heights
                [["Amy", "Belinda"]],    # Jane - Sense and Sense
            ])
        with self.assertNumQueries(3):
            qs = Author.objects.prefetch_related('books__read_by', 'books')
            lists = [[[six.text_type(r) for r in b.read_by.all()]
                      for b in a.books.all()]
    
                 for a 
in qs]
            self.assertEqual(lists,
            [
                [["Amy"], ["Belinda"]],  # Charlotte - Poems, Jane Eyre
                [["Amy"]],                # Anne - Poems
                [["Amy"], []],            # Emily - Poems, Wuthering Heights
                [["Amy", "Belinda"]],    # Jane - Sense and Sense
            ])

    def test_get(self):
        """
        Test that objects retrieved with .get() get the prefetch behavior.
        """
        # Need a double
        with self.assertNumQueries(3):
            author = Author.objects.prefetch_related('books__read_by').get(name="Charlotte")
            lists = [[six.text_type(r) for r in b.read_by.all()]
                     for b in author.books.all()]
            self.assertEqual(lists, [["Amy"], ["Belinda"]])  # Poems, Jane Eyre

    def test_foreign_key_then_m2m(self):
        """
        Test we can follow an m2m relation after a relation like ForeignKey
        that doesn't have many objects
        """
        with self.assertNumQueries(2):
            qs = Author.objects.select_related('first_book').prefetch_related('first_book__read_by')
            lists = [[six.text_type(r) for r in a.first_book.read_by.all()]
                     for a in qs]
            self.assertEqual(lists, [["Amy"],
                                     ["Amy"],
                                