from __future__ import print_function

import six

from unittest import TestCase
from disco.types.base import Model, Field, enum, snowflake, ConversionError


class _A(Model):
    a = Field(int)
    b = Field(float)
    c = Field(str)


class _B(Model):
    a = Field(int)
    b = Field(float)
    c = Field(str)


class _C(Model):
    a = Field(_A)
    b = Field(_B)


class TestModel(TestCase):
    def test_model_simple_loading(self):
        inst = _A(dict(a=1, b=1.1, c='test'))
        self.assertEqual(inst.a, 1)
        self.assertEqual(inst.b, 1.1)
        self.assertEqual(inst.c, 'test')

    def test_model_load_into(self):
        inst = _A()
        _A.load_into(inst, dict(a=1, b=1.1, c='test'))
        self.assertEqual(inst.a, 1)
        self.assertEqual(inst.b, 1.1)
        self.assertEqual(inst.c, 'test')

    def test_model_loading_consume(self):
        obj = {
            'a': {
                'a': 1,
                'b': 2.2,
                'c': '3',
                'd': 'wow',
            },
            'b': {
                'a': 3,
                'b': 2.2,
                'c': '1',
                'z': 'wtf'
            },
            'g': 'lmao'
        }

        inst = _C()
        inst.load(obj, consume=True)

        self.assertEqual(inst.a.a, 1)
        self.assertEqual(inst.b.c, '1')

        self.assertEqual(obj, {'a': {'d': 'wow'}, 'b': {'z': 'wtf'}, 'g': 'lmao'})

    def test_model_field_enum(self):
        class en(object):
            A = 1
            B = 2
            C = 3

        class _M(Model):
            field = Field(enum(en))

        self.assertEqual(_M(field=en.A).field, en.A)
        self.assertEqual(_M(field=2).field, en.B)
        self.assertEqual(_M(field='3').field, None)
        self.assertEqual(_M(field='a').field, en.A)
        self.assertEqual(_M(field='A').field, en.A)

    def test_model_field_snowflake(self):
        class _M(Model):
            field = Field(snowflake)

        self.assertEqual(_M(field=327936274851954688).field, 327936274851954688)
        self.assertEqual(_M(field='327936274851954688').field, 327936274851954688)

        with self.assertRaises(ConversionError):
            _M(field='asdf')

    def test_model_field_cast(self):
        class Object(object):
            def __init__(self, v):
                self.v = v

            def __str__(self):
                return self.v

        class _M(Model):
            field = Field(Object, cast=six.text_type)

        inst = _M(field=u'wowza')
        self.assertEqual(inst.to_dict(), {'field': u'wowza'})