from typing import Any
from uuid import uuid4

import numpy
import pytest
from py10x_kernel import BTraitableProcessor, XCache

from core_10x.code_samples.person import Person as BasePerson
from core_10x.package_refactoring import PackageRefactoring
from core_10x.rc import RC_TRUE
from core_10x.trait_definition import T
from core_10x.ts_store import TsDuplicateKeyError


class EnhancedPerson(BasePerson):
    numpy_weight_lbs: Any = T()

    def numpy_weight_lbs_get(self):
        return numpy.float64(self.weight_lbs)


test_classes = {
    cls_name: type(
        cls_name,
        (EnhancedPerson,),
        {
            '__module__': __name__,
        },
        custom_collection=custom_collection
    )
    for custom_collection in (None,True)
    for cls_name in (f'Person#{uuid4().hex}{custom_collection and "#Custom" or ""}' for _ in range(2))
}

globals().update(test_classes)


@pytest.fixture(params=[True, False], ids=['custom_collection', 'default_collection'])
def ts_setup(ts_instance,request):
    Person, Person1 = list(test_classes.values())[request.param*2:][:2]
    c, c1 = [uuid4().hex if request.param else PackageRefactoring.find_class_id(cls) for cls in (Person,Person1)]

    with ts_instance:
        p = Person(first_name='John', last_name='Doe', _collection_name=c if request.param else None)
        p.set_values(age=30, weight_lbs=100)
        assert p._rev == 0
        assert p.save() == RC_TRUE
        assert p._rev == 1
        assert p.save() == RC_TRUE
        assert p._rev == 1
        assert p.id().collection_name == (c if request.param else None)


        p1 = Person1(first_name='Joe', last_name='Doe', _collection_name=c1 if request.param else None)
        p1.set_values(age=32, weight_lbs=200)
        assert p1.save()
        assert p1.age == 32
        assert p1._rev == 1
        assert p1.id().collection_name == (c1 if request.param else None)

    yield ts_instance, p, p1, c, c1

    # Cleanup
    XCache.clear()
    BTraitableProcessor.current().end_using()
    for cn in [c, c1]:
        ts_instance.delete_collection(cn)
    assert not {c, c1}.intersection(ts_instance.collection_names('.*'))


class TestTSStore:
    """Test class for TS Store functionality."""

    def test_collection(self, ts_setup):
        ts_store, _p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)
        assert collection is not None

    def test_save(self, ts_setup):
        ts_store, p, _p1, c, _c1= ts_setup

        collection = ts_store.collection(c)
        serialized_entity = p.serialize_object()
        _rev = collection.save(serialized_entity.copy())
        assert p._rev == _rev

        serialized_entity |= {'attr': {'nested': 'value'}}
        _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 1 == _rev
        assert collection.load(p.id().value) == serialized_entity

        # test that nested dictionary replaces rather than updates
        serialized_entity |= {'attr': {'nested1': 'value1'}}
        _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 2 == _rev
        # assert collection.load(_id) == serialized_entity|{'attr': {'nested': 'value', 'nested1': 'value1'}} #incorrect behavior
        assert collection.load(p.id().value) == serialized_entity

        # test that dots are not interpreted as nested fields at the top level
        serialized_entity |= {'attr.nested2': 'value2'}
        _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 3 == _rev
        assert collection.load(p.id().value) == serialized_entity

        # check that we can have dictionary keys starting with $
        serialized_entity |= {'foo': {'$foo': 1}}
        # with pytest.raises(pyts_store.errors.WriteError, match="Unrecognized expression '\\$foo',"): #incorrect behavior
        _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 4 == _rev
        assert collection.load(p.id().value) == serialized_entity

        # check that we can *not* have top level keys starting with $ (current behavior, not ideal, but prob. ok)
        serialized_entity |= {'$foo': 1}
        with pytest.raises(Exception, match='Use of undefined variable: foo'):
            _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 4 == _rev
        del serialized_entity['$foo']
        assert collection.load(p.id().value) == serialized_entity

        # test that dots are not interpreted as nested fields at the nested level
        serialized_entity |= {'attr': {'nested.value': 1}}
        _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 5 == _rev
        assert collection.load(p.id().value) == serialized_entity

        # check that we can unset fields
        del serialized_entity['attr']
        _rev = collection.save(serialized_entity.copy())
        serialized_entity |= dict(_rev=_rev)
        assert p._rev + 6 == _rev
        assert collection.load(p.id().value) == serialized_entity

    def test_delete_restore(self, ts_setup):
        ts_store, p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)
        serialized_entity = p.serialize_object()
        id_value = serialized_entity['_id']
        assert collection.delete(id_value)
        assert collection.load(id_value) is None
        # TODO: restore is not implemented so use save_new meanwhile
        # assert collection.restore(id_value)
        collection.save_new(serialized_entity)
        assert collection.load(id_value) == serialized_entity

    def test_find(self, ts_setup):
        ts_store, p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)
        result = collection.find()
        assert next(iter(result)) == p.serialize_object()
        assert list(result) == []

    def test_load(self, ts_setup):
        ts_store, p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)
        id_value = p.id().value
        result = collection.load(id_value)
        assert result == p.serialize_object()

    def test_delete(self, ts_setup):
        ts_store, _p, p1, _c, c1= ts_setup
        collection = ts_store.collection(c1)
        id_value = p1.id().value
        result = collection.delete(id_value)
        assert result
        result = collection.load(id_value)
        assert result is None

    def test_save_new_with_overwrite(self, ts_setup):
        """Test save_new with overwrite=True flag."""
        ts_store, p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)
        serialized_entity = p.serialize_object()
        id_value = serialized_entity['_id']

        # Try to save_new with overwrite=True on existing document
        result = collection.save_new(serialized_entity.copy(), overwrite=True)
        assert result == 1
        assert collection.load(id_value) == serialized_entity

    def test_save_new_with_set_operation(self, ts_setup):
        """Test save_new with $set MongoDB-style operation."""
        ts_store, _p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)

        # Create a new document with $set
        doc_id = 'test_doc_123'
        serialized_entity = {'$set': {'_id': doc_id, 'name': 'Test Document', 'value': 42}}

        result = collection.save_new(serialized_entity)
        assert result == 1

        loaded = collection.load(doc_id)
        assert loaded == serialized_entity['$set'] | {'_rev': 1}

    def test_save_new_duplicate_key_error(self, ts_setup):
        """Test that save_new raises TsDuplicateKeyError when inserting duplicate without overwrite."""
        ts_store, _p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)

        # Test 1: Duplicate without $set
        doc_id = 'duplicate_test_123'
        serialized_entity = {'_id': doc_id, 'name': 'First Document'}

        result = collection.save_new(serialized_entity)
        assert result == 1

        # Try to insert the same document again without overwrite (no $set)
        with pytest.raises(TsDuplicateKeyError, match=f'Duplicate key error collection.*dup key.*{doc_id}'):
            collection.save_new(serialized_entity, overwrite=False)

        # Test 2: Duplicate with $set
        doc_id2 = 'duplicate_test_456'
        serialized_entity2 = {'$set': {'_id': doc_id2, 'name': 'First Document with $set'}}

        result = collection.save_new(serialized_entity2)
        assert result == 1

        # Try to insert the same document again without overwrite (with $set)
        with pytest.raises(TsDuplicateKeyError, match=f'Duplicate key error collection.*dup key.*{doc_id2}'):
            collection.save_new(serialized_entity2, overwrite=False)

    def test_save_new_with_set_and_overwrite(self, ts_setup):
        """Test save_new with $set and overwrite=True."""
        ts_store, _p, _p1, c, _c1= ts_setup
        collection = ts_store.collection(c)

        doc_id = 'set_overwrite_test_123'
        # First insert
        serialized_entity1 = {'_id': doc_id, 'name': 'Original'}
        result = collection.save_new(serialized_entity1)
        assert result == 1

        # Update with $set and overwrite=True
        serialized_entity2 = {'$set': {'_id': doc_id, 'name': 'Updated', 'new_field': 'new_value'}}
        result = collection.save_new(serialized_entity2, overwrite=True)
        assert result == 1

        loaded = collection.load(doc_id)
        assert loaded['name'] == 'Updated'
        assert loaded['new_field'] == 'new_value'

    def test_ts_class_association_ts_uri_resolution(self, ts_instance):
        """Test that TsClassAssociation.ts_uri correctly resolves store URIs for classes and their subclasses."""
        from core_10x.py_class import PyClass
        from core_10x.traitable import NamedTsStore, T, Traitable, TsClassAssociation

        class Dummy1(Traitable):
            text: str = T(T.ID)

        class Dummy2(Traitable):
            text: str = T(T.ID)

        class Dummy3(Dummy2): ...

        dummy_uri1 = 'mongodb://localhost/dummy1'
        dummy_uri2 = 'mongodb://localhost/dummy2'

        with ts_instance:
            # Create and save NamedTsStore objects
            ns1 = NamedTsStore(logical_name='dummy1', uri=dummy_uri1, _replace=True)
            ns1.save()
            ns2 = NamedTsStore(logical_name='dummy2', uri=dummy_uri2, _replace=True)
            ns2.save()

            # Create and save TsClassAssociation objects
            name1 = PyClass.name(Dummy1)
            name2 = PyClass.name(Dummy2)
            TsClassAssociation(py_canonical_name=name1, ts_logical_name='dummy1', _replace=True).save()
            TsClassAssociation(py_canonical_name=name2, ts_logical_name='dummy2', _replace=True).save()

            # Class-specific association
            assert TsClassAssociation.ts_uri(Dummy1) == dummy_uri1

            # Subclass should inherit parent's association if it has none of its own
            assert TsClassAssociation.ts_uri(Dummy3) == dummy_uri2
