# -*- coding: utf-8 -*-
from decimal import Decimal

from mongoengine import *

from tests.utils import MongoDBTestCase


class TestCachedReferenceField(MongoDBTestCase):
    def test_get_and_save(self):
        """
        Tests #1047: CachedReferenceField creates DBRefs on to_python,
        but can't save them on to_mongo.
        """

        class Animal(Document):
            name = StringField()
            tag = StringField()

        class Ocorrence(Document):
            person = StringField()
            animal = CachedReferenceField(Animal)

        Animal.drop_collection()
        Ocorrence.drop_collection()

        Ocorrence(
            person="testte", animal=Animal(name="Leopard", tag="heavy").save()
        ).save()
        p = Ocorrence.objects.get()
        p.person = "new_testte"
        p.save()

    def test_general_things(self):
        class Animal(Document):
            name = StringField()
            tag = StringField()

        class Ocorrence(Document):
            person = StringField()
            animal = CachedReferenceField(Animal, fields=["tag"])

        Animal.drop_collection()
        Ocorrence.drop_collection()

        a = Animal(name="Leopard", tag="heavy")
        a.save()

        self.assertEqual(Animal._cached_reference_fields, [Ocorrence.animal])
        o = Ocorrence(person="teste", animal=a)
        o.save()

        p = Ocorrence(person="Wilson")
        p.save()

        self.assertEqual(Ocorrence.objects(animal=None).count(), 1)

        self.assertEqual(a.to_mongo(fields=["tag"]), {"tag": "heavy", "_id": a.pk})

        self.assertEqual(o.to_mongo()["animal"]["tag"], "heavy")

        # counts
        Ocorrence(person="teste 2").save()
        Ocorrence(person="teste 3").save()

        count = Ocorrence.objects(animal__tag="heavy").count()
        self.assertEqual(count, 1)

        ocorrence = Ocorrence.objects(animal__tag="heavy").first()
        self.assertEqual(ocorrence.person, "teste")
        self.assertIsInstance(ocorrence.animal, Animal)

    def test_with_decimal(self):
        class PersonAuto(Document):
            name = StringField()
            salary = DecimalField()

        class SocialTest(Document):
            group = StringField()
            person = CachedReferenceField(PersonAuto, fields=("salary",))

        PersonAuto.drop_collection()
        SocialTest.drop_collection()

        p = PersonAuto(name="Alberto", salary=Decimal("7000.00"))
        p.save()

        s = SocialTest(group="dev", person=p)
        s.save()

        self.assertEqual(
            SocialTest.objects._collection.find_one({"person.salary": 7000.00}),
            {"_id": s.pk, "group": s.group, "person": {"_id": p.pk, "salary": 7000.00}},
        )

    def test_cached_reference_field_reference(self):
        class Group(Document):
            name = StringField()

        class Person(Document):
            name = StringField()
            group = ReferenceField(Group)

        class SocialData(Document):
            obs = StringField()
            tags = ListField(StringField())
            person = CachedReferenceField(Person, fields=("group",))

        Group.drop_collection()
        Person.drop_collection()
        SocialData.drop_collection()

        g1 = Group(name="dev")
        g1.save()

        g2 = Group(name="designers")
        g2.save()

        p1 = Person(name="Alberto", group=g1)
        p1.save()

        p2 = Person(name="Andre", group=g1)
        p2.save()

        p3 = Person(name="Afro design", group=g2)
        p3.save()

        s1 = SocialData(obs="testing 123", person=p1, tags=["tag1", "tag2"])
        s1.save()

        s2 = SocialData(obs="testing 321", person=p3, tags=["tag3", "tag4"])
        s2.save()

        self.assertEqual(
            SocialData.objects._collection.find_one({"tags": "tag2"}),
            {
                "_id": s1.pk,
                "obs": "testing 123",
                "tags": ["tag1", "tag2"],
                "person": {"_id": p1.pk, "group": g1.pk},
            },
        )

        self.assertEqual(SocialData.objects(person__group=g2).count(), 1)
        self.assertEqual(SocialData.objects(person__group=g2).first(), s2)

    def test_cached_reference_field_push_with_fields(self):
        class Product(Document):
            name = StringField()

        Product.drop_collection()

        class Basket(Document):
            products = ListField(CachedReferenceField(Product, fields=["name"]))

        Basket.drop_collection()
        product1 = Product(name="abc").save()
        product2 = Product(name="def").save()
        basket = Basket(products=[product1]).save()
        self.assertEqual(
            Basket.objects._collection.find_one(),
            {
                "_id": basket.pk,
                "products": [{"_id": product1.pk, "name": product1.name}],
            },
        )
        # push to list
        basket.update(push__products=product2)
        basket.reload()
        self.assertEqual(
            Basket.objects._collection.find_one(),
            {
                "_id": basket.pk,
                "products": [
                    {"_id": product1.pk, "name": product1.name},
                    {"_id": product2.pk, "name": product2.name},
                ],
            },
        )

    def test_cached_reference_field_update_all(self):
        class Person(Document):
            TYPES = (("pf", "PF"), ("pj", "PJ"))
            name = StringField()
            tp = StringField(choices=TYPES)
            father = CachedReferenceField("self", fields=("tp",))

        Person.drop_collection()

        a1 = Person(name="Wilson Father", tp="pj")
        a1.save()

        a2 = Person(name="Wilson Junior", tp="pf", father=a1)
        a2.save()

        a2 = Person.objects.with_id(a2.id)
        self.assertEqual(a2.father.tp, a1.tp)

        self.assertEqual(
            dict(a2.to_mongo()),
            {
                "_id": a2.pk,
                "name": u"Wilson Junior",
                "tp": u"pf",
                "father": {"_id": a1.pk, "tp": u"pj"},
            },
        )

        self.assertEqual(Person.objects(father=a1)._query, {"father._id": a1.pk})
        self.assertEqual(Person.objects(father=a1).count(), 1)

        Person.objects.update(set__tp="pf")
        Person.father.sync_all()

        a2.reload()
        self.assertEqual(
            dict(a2.to_mongo()),
            {
                "_id": a2.pk,
                "name": u"Wilson Junior",
                "tp": u"pf",
                "father": {"_id": a1.pk, "tp": u"pf"},
            },
        )

    def test_cached_reference_fields_on_embedded_documents(self):
        with self.assertRaises(InvalidDocumentError):

            class Test(Document):
                name = StringField()

            type(
                "WrongEmbeddedDocument",
                (EmbeddedDocument,),
                {"test": CachedReferenceField(Test)},
            )

    def test_cached_reference_auto_sync(self):
        class Person(Document):
            TYPES = (("pf", "PF"), ("pj", "PJ"))
            name = StringField()
            tp = StringField(choices=TYPES)

            father = CachedReferenceField("self", fields=("tp",))

        Person.drop_collection()

        a1 = Person(name="Wilson Father", tp="pj")
        a1.save()

        a2 = Person(name="Wilson Junior", tp="pf", father=a1)
        a2.save()

        a1.tp = "pf"
        a1.save()

        a2.reload()
        self.assertEqual(
            dict(a2.to_mongo()),
            {
                "_id": a2.pk,
                "name": "Wilson Junior",
                "tp": "pf",
                "father": {"_id": a1.pk, "tp": "pf"},
            },
        )

    def test_cached_reference_auto_sync_disabled(self):
        class Persone(Document):
            TYPES = (("pf", "PF"), ("pj", "PJ"))
            name = StringField()
            tp = StringField(choices=TYPES)

            father = CachedReferenceField("self", fields=("tp",), auto_sync=False)

        Persone.drop_collection()

        a1 = Persone(name="Wilson Father", tp="pj")
        a1.save()

        a2 = Persone(name="Wilson Junior", tp="pf", father=a1)
        a2.save()

        a1.tp = "pf"
        a1.save()

        self.assertEqual(
            Persone.objects._collection.find_one({"_id": a2.pk}),
            {
                "_id": a2.pk,
                "name": "Wilson Junior",
                "tp": "pf",
                "father": {"_id": a1.pk, "tp": "pj"},
            },
        )

    def test_cached_reference_embedded_fields(self):
        class Owner(EmbeddedDocument):
            TPS = (("n", "Normal"), ("u", "Urgent"))
            name = StringField()
            tp = StringField(verbose_name="Type", db_field="t", choices=TPS)

        class Animal(Document):
            name = StringField()
            tag = StringField()

            owner = EmbeddedDocumentField(Owner)

        class Ocorrence(Document):
            person = StringField()
            animal = CachedReferenceField(Animal, fields=["tag", "owner.tp"])

        Animal.drop_collection()
        Ocorrence.drop_collection()

        a = Animal(
            name="Leopard", tag="heavy", owner=Owner(tp="u", name="Wilson Júnior")
        )
        a.save()

        o = Ocorrence(person="teste", animal=a)
        o.save()
        self.assertEqual(
            dict(a.to_mongo(fields=["tag", "owner.tp"])),
            {"_id": a.pk, "tag": "heavy", "owner": {"t": "u"}},
        )
        self.assertEqual(o.to_mongo()["animal"]["tag"], "heavy")
        self.assertEqual(o.to_mongo()["animal"]["owner"]["t"], "u")

        # Check to_mongo with fields
        self.assertNotIn("animal", o.to_mongo(fields=["person"]))

        # counts
        Ocorrence(person="teste 2").save()
        Ocorrence(person="teste 3").save()

        count = Ocorrence.objects(animal__tag="heavy", animal__owner__tp="u").count()
        self.assertEqual(count, 1)

        ocorrence = Ocorrence.objects(
            animal__tag="heavy", animal__owner__tp="u"
        ).first()
        self.assertEqual(ocorrence.person, "teste")
        self.assertIsInstance(ocorrence.animal, Animal)

    def test_cached_reference_embedded_list_fields(self):
        class Owner(EmbeddedDocument):
            name = StringField()
            tags = ListField(StringField())

        class Animal(Document):
            name = StringField()
            tag = StringField()

            owner = EmbeddedDocumentField(Owner)

        class Ocorrence(Document):
            person = StringField()
            animal = CachedReferenceField(Animal, fields=["tag", "owner.tags"])

        Animal.drop_collection()
        Ocorrence.drop_collection()

        a = Animal(
            name="Leopard",
            tag="heavy",
            owner=Owner(tags=["cool", "funny"], name="Wilson Júnior"),
        )
        a.save()

        o = Ocorrence(person="teste 2", animal=a)
        o.save()
        self.assertEqual(
            dict(a.to_mongo(fields=["tag", "owner.tags"])),
            {"_id": a.pk, "tag": "heavy", "owner": {"tags": ["cool", "funny"]}},
        )

        self.assertEqual(o.to_mongo()["animal"]["tag"], "heavy")
        self.assertEqual(o.to_mongo()["animal"]["owner"]["tags"], ["cool", "funny"])

        # counts
        Ocorrence(person="teste 2").save()
        Ocorrence(person="teste 3").save()

        query = Ocorrence.objects(
            animal__tag="heavy", animal__owner__tags="cool"
        )._query
        self.assertEqual(query, {"animal.owner.tags": "cool", "animal.tag": "heavy"})

        ocorrence = Ocorrence.objects(
            animal__tag="heavy", animal__owner__tags="cool"
        ).first()
        self.assertEqual(ocorrence.person, "teste 2")
        self.assertIsInstance(ocorrence.animal, Animal)