# python
import os
import threading
from config.paths import get_database_dir
from tinydb import TinyDB
from tinydb.queries import QueryInstance

from db.mongo_client import get_mongo_client, get_mongo_db_name

try:
    from tinydb.table import Document
except ImportError:  # pragma: no cover - tinydb version compatibility
    from tinydb.database import Document


def _stable_clause_key(clause: dict) -> str:
    """Return a stable string key for sorting MongoDB filter clauses."""
    import json
    return json.dumps(clause, sort_keys=True, default=str)


try:
    from pymongo import ASCENDING
except ImportError:  # pragma: no cover - pymongo is a required dependency
    ASCENDING = 1


USE_MONGODB = os.environ.get('USE_MONGODB', 'true').lower() == 'true'
# Resolve the MongoDB database name once at module load so runtime changes to
# DB_ENV/DATA_ENV in tests do not switch databases mid-process.
_mongo_db_name = get_mongo_db_name() if USE_MONGODB else None
base_dir = get_database_dir()
os.makedirs(base_dir, exist_ok=True)


# ---------------------------------------------------------------------------
# TinyDB-backed table wrapper
# ---------------------------------------------------------------------------

class LockedTable:
    """
    Thread-safe wrapper around a TinyDB table. All callable attribute access
    is wrapped to acquire a reentrant lock while calling the underlying method.
    Non-callable attributes are returned directly.
    """
    def __init__(self, table):
        self._table = table
        self._lock = threading.RLock()

    def __getattr__(self, name):
        # avoid proxying internal attrs
        if name in ('_table', '_lock'):
            return super().__getattribute__(name)

        attr = getattr(self._table, name)

        if callable(attr):
            def locked_call(*args, **kwargs):
                with self._lock:
                    return attr(*args, **kwargs)
            return locked_call
        return attr

    # convenience explicit methods (ensure these are class methods, not top-level)
    def insert(self, *args, **kwargs):
        with self._lock:
            return self._table.insert(*args, **kwargs)

    def insert_multiple(self, *args, **kwargs):
        with self._lock:
            return self._table.insert_multiple(*args, **kwargs)

    def search(self, *args, **kwargs):
        with self._lock:
            return self._table.search(*args, **kwargs)

    def get(self, *args, **kwargs):
        with self._lock:
            return self._table.get(*args, **kwargs)

    def all(self, *args, **kwargs):
        with self._lock:
            return self._table.all(*args, **kwargs)

    def remove(self, *args, **kwargs):
        with self._lock:
            return self._table.remove(*args, **kwargs)

    def update(self, *args, **kwargs):
        with self._lock:
            return self._table.update(*args, **kwargs)

    def truncate(self):
        with self._lock:
            return self._table.truncate()

    def close(self):
        with self._lock:
            return self._table.close()


# ---------------------------------------------------------------------------
# TinyDB Query -> MongoDB filter translator
# ---------------------------------------------------------------------------

_MONGO_OP_MAP = {
    '==': '$eq',
    '!=': '$ne',
    '<': '$lt',
    '<=': '$lte',
    '>': '$gt',
    '>=': '$gte',
}

_NEGATED_OPS = {
    '$eq': '$ne',
    '$ne': '$eq',
    '$lt': '$gte',
    '$lte': '$gt',
    '$gt': '$lte',
    '$gte': '$lt',
}


def _field_to_mongo(field_path: tuple) -> str:
    """Map a TinyDB field path to a MongoDB field name.

    The model ``id`` field is stored as the MongoDB ``_id`` field, so queries
    on ``id`` are translated to queries on ``_id``.
    """
    if len(field_path) == 1:
        return '_id' if field_path[0] == 'id' else field_path[0]
    return '.'.join('_id' if p == 'id' else p for p in field_path)


def _negate_condition(cond: dict) -> dict | None:
    """Return a MongoDB condition that negates a single-field condition."""
    if len(cond) != 1:
        return None
    field, inner = next(iter(cond.items()))
    if not isinstance(inner, dict) or len(inner) != 1:
        return None
    op, value = next(iter(inner.items()))
    if op in _NEGATED_OPS:
        return {field: {_NEGATED_OPS[op]: value}}
    return None


def _hash_to_mongo_filter(query_hash) -> dict | None:
    """Translate a TinyDB query hash tuple to a MongoDB filter document.

    Returns ``None`` when the query construct cannot be expressed as a native
    MongoDB filter, signalling that the caller should fall back to in-memory
    TinyDB evaluation.
    """
    if not isinstance(query_hash, tuple) or len(query_hash) == 0:
        return None

    op = query_hash[0]

    if op in _MONGO_OP_MAP:
        field = _field_to_mongo(query_hash[1])
        value = query_hash[2]
        return {field: {_MONGO_OP_MAP[op]: value}}

    if op == 'exists':
        field = _field_to_mongo(query_hash[1])
        return {field: {'$exists': True}}

    if op == 'one_of':
        field = _field_to_mongo(query_hash[1])
        return {field: {'$in': list(query_hash[2])}}

    if op == 'any':
        field = _field_to_mongo(query_hash[1])
        return {field: {'$in': list(query_hash[2])}}

    if op == 'all':
        field = _field_to_mongo(query_hash[1])
        return {field: {'$all': list(query_hash[2])}}

    if op == 'matches':
        field = _field_to_mongo(query_hash[1])
        # TinyDB matches() anchors the regex at the start of the string.
        return {field: {'$regex': f'^{query_hash[2]}'}}

    if op == 'search':
        field = _field_to_mongo(query_hash[1])
        return {field: {'$regex': query_hash[2]}}

    if op == 'and':
        merged: dict = {}
        for sub_hash in query_hash[1]:
            sub = _hash_to_mongo_filter(sub_hash)
            if sub is None:
                return None
            for field, inner in sub.items():
                if field in merged:
                    if isinstance(merged[field], dict) and isinstance(inner, dict):
                        merged[field].update(inner)
                    elif isinstance(merged[field], list) and isinstance(inner, list):
                        merged[field].extend(inner)
                    else:
                        return None
                else:
                    merged[field] = (
                        dict(inner) if isinstance(inner, dict)
                        else list(inner) if isinstance(inner, list)
                        else inner
                    )
        return merged

    if op == 'or':
        clauses = [_hash_to_mongo_filter(sub_hash) for sub_hash in query_hash[1]]
        if any(c is None for c in clauses):
            return None
        return {'$or': sorted(clauses, key=_stable_clause_key)}

    if op == 'not':
        inner = _hash_to_mongo_filter(query_hash[1])
        if inner is None:
            return None
        negated = _negate_condition(inner)
        if negated is not None:
            return negated
        return None

    # Unsupported operation (test, fragment, noop, etc.) -> fall back.
    return None


def _query_to_mongo_filter(query) -> dict | None:
    """Translate a TinyDB QueryInstance to a MongoDB filter, if possible."""
    if isinstance(query, dict):
        return query
    if isinstance(query, QueryInstance):
        return _hash_to_mongo_filter(query._hash)
    return None


def _evaluate_in_memory(docs, query) -> list:
    """Evaluate a TinyDB query against an in-memory list of documents."""
    if callable(query):
        return [doc for doc in docs if query(doc)]
    return docs


# ---------------------------------------------------------------------------
# MongoDB-backed table wrapper
# ---------------------------------------------------------------------------

class MongoLockedTable:
    """Drop-in replacement for ``LockedTable`` that delegates to MongoDB.

    The adapter preserves the existing table API while mapping the model
    ``id`` field to MongoDB's ``_id`` field on reads and writes.
    """

    def __init__(self, collection_name: str):
        self.collection_name = collection_name

    def _collection(self):
        client = get_mongo_client()
        return client[_mongo_db_name][self.collection_name]

    @staticmethod
    def _doc_to_mongo(doc: dict) -> dict:
        """Store a copy of ``doc`` with ``id`` promoted to MongoDB ``_id``.

        This avoids storing both ``_id`` and ``id`` with identical values.
        The original ``id`` field is removed from the stored document.
        """
        if doc is None:
            return None
        d = dict(doc)
        if 'id' in d:
            d['_id'] = d.pop('id')
        return d

    @staticmethod
    def _doc_from_mongo(doc: dict):
        """Return a TinyDB-compatible Document with ``doc_id`` set to ``_id``.

        Restores the model ``id`` field from MongoDB's ``_id`` and exposes
        TinyDB's ``doc_id`` attribute so callers that rely on it continue to
        work.
        """
        if doc is None:
            return None
        d = dict(doc)
        doc_id = d.pop('_id', None)
        if doc_id is not None:
            d['id'] = doc_id
        return Document(d, doc_id=doc_id)

    def _mongo_filter(self, cond):
        """Translate a TinyDB query or dict to a MongoDB filter."""
        return _query_to_mongo_filter(cond)

    def all(self):
        return [self._doc_from_mongo(doc) for doc in self._collection().find({})]

    def search(self, cond):
        mongo_filter = self._mongo_filter(cond)
        if mongo_filter is not None:
            cursor = self._collection().find(mongo_filter)
            return [self._doc_from_mongo(doc) for doc in cursor]

        # Fallback: fetch all and evaluate the TinyDB query in Python.
        docs = list(self._collection().find({}))
        matched = _evaluate_in_memory(
            [self._doc_from_mongo(doc) for doc in docs], cond
        )
        return matched

    def get(self, cond):
        mongo_filter = self._mongo_filter(cond)
        if mongo_filter is not None:
            doc = self._collection().find_one(mongo_filter)
            return self._doc_from_mongo(doc)

        docs = list(self._collection().find({}))
        for doc in docs:
            d = self._doc_from_mongo(doc)
            if callable(cond) and cond(d):
                return d
        return None

    def insert(self, document: dict):
        doc = self._doc_to_mongo(document)
        result = self._collection().insert_one(doc)
        return str(result.inserted_id)

    def insert_multiple(self, documents: list):
        if not documents:
            return []
        docs = [self._doc_to_mongo(d) for d in documents]
        result = self._collection().insert_many(docs)
        return [str(iid) for iid in result.inserted_ids]

    def update(self, fields, cond=None, doc_ids=None):
        is_callable = callable(fields)

        if doc_ids is not None:
            mongo_filter = {'_id': {'$in': list(doc_ids)}}
            target_ids = [str(did) for did in doc_ids]
            if not target_ids:
                return []

            if is_callable:
                # Fetch, apply callable in-memory, and replace each document.
                updated_ids = []
                for doc in self._collection().find(mongo_filter):
                    d = self._doc_from_mongo(doc)
                    fields(d)
                    new_doc = self._doc_to_mongo(d)
                    new_doc.pop('_id', None)
                    self._collection().update_one(
                        {'_id': doc['_id']}, {'$set': new_doc}
                    )
                    updated_ids.append(str(doc['_id']))
                return updated_ids

            update_doc = self._doc_to_mongo(fields) or {}
            update_doc.pop('_id', None)
            update_doc.pop('id', None)
            if update_doc:
                self._collection().update_many(mongo_filter, {'$set': update_doc})
            return target_ids

        mongo_filter = self._mongo_filter(cond)
        if mongo_filter is not None and not is_callable:
            update_doc = self._doc_to_mongo(fields) or {}
            update_doc.pop('_id', None)
            update_doc.pop('id', None)
            target_ids = [
                str(doc['_id'])
                for doc in self._collection().find(mongo_filter, {'_id': 1})
            ]
            if target_ids and update_doc:
                self._collection().update_many(
                    mongo_filter, {'$set': update_doc}
                )
            return target_ids

        # Fallback: evaluate the query in-memory and update one at a time.
        docs = list(self._collection().find({}))
        updated_ids = []
        for doc in docs:
            d = self._doc_from_mongo(doc)
            match = cond(d) if callable(cond) else (mongo_filter is not None)
            if not match:
                continue
            if is_callable:
                fields(d)
                new_doc = self._doc_to_mongo(d)
                new_doc.pop('_id', None)
                self._collection().update_one(
                    {'_id': doc['_id']}, {'$set': new_doc}
                )
            else:
                update_doc = self._doc_to_mongo(fields) or {}
                update_doc.pop('_id', None)
                update_doc.pop('id', None)
                if update_doc:
                    self._collection().update_one(
                        {'_id': doc['_id']}, {'$set': update_doc}
                    )
            updated_ids.append(str(doc['_id']))
        return updated_ids

    def remove(self, cond):
        mongo_filter = self._mongo_filter(cond)
        if mongo_filter is not None:
            target_ids = [
                str(doc['_id'])
                for doc in self._collection().find(mongo_filter, {'_id': 1})
            ]
            if target_ids:
                self._collection().delete_many(mongo_filter)
            return target_ids

        # Fallback: evaluate the query in-memory and delete one at a time.
        docs = list(self._collection().find({}))
        removed_ids = []
        for doc in docs:
            d = self._doc_from_mongo(doc)
            if callable(cond) and cond(d):
                self._collection().delete_one({'_id': doc['_id']})
                removed_ids.append(str(doc['_id']))
        return removed_ids

    def truncate(self):
        self._collection().delete_many({})

    def close(self):
        # MongoDB clients are shared and long-lived; nothing to close here.
        pass


# ---------------------------------------------------------------------------
# Collection factory
# ---------------------------------------------------------------------------

def _make_table(json_path: str, collection_name: str):
    if USE_MONGODB:
        return MongoLockedTable(collection_name)
    db = TinyDB(json_path, indent=2)
    return LockedTable(db)


# Setup DB files next to this module

child_path = os.path.join(base_dir, 'children.json')
task_path = os.path.join(base_dir, 'tasks.json')
routine_path = os.path.join(base_dir, 'routines.json')
routine_items_path = os.path.join(base_dir, 'routine_items.json')
routine_schedules_path = os.path.join(base_dir, 'routine_schedules.json')
routine_extensions_path = os.path.join(base_dir, 'routine_extensions.json')
reward_path = os.path.join(base_dir, 'rewards.json')
image_path = os.path.join(base_dir, 'images.json')
pending_reward_path = os.path.join(base_dir, 'pending_rewards.json')
pending_confirmations_path = os.path.join(base_dir, 'pending_confirmations.json')
users_path = os.path.join(base_dir, 'users.json')
tracking_events_path = os.path.join(base_dir, 'tracking_events.json')
child_overrides_path = os.path.join(base_dir, 'child_overrides.json')
chore_schedules_path = os.path.join(base_dir, 'chore_schedules.json')
task_extensions_path = os.path.join(base_dir, 'task_extensions.json')
refresh_tokens_path = os.path.join(base_dir, 'refresh_tokens.json')
push_subscriptions_path = os.path.join(base_dir, 'push_subscriptions.json')
digest_action_tokens_path = os.path.join(base_dir, 'digest_action_tokens.json')

# Expose table objects backed by TinyDB or MongoDB based on USE_MONGODB
child_db = _make_table(child_path, 'children')
task_db = _make_table(task_path, 'tasks')
routine_db = _make_table(routine_path, 'routines')
routine_items_db = _make_table(routine_items_path, 'routine_items')
routine_schedules_db = _make_table(routine_schedules_path, 'routine_schedules')
routine_extensions_db = _make_table(routine_extensions_path, 'routine_extensions')
reward_db = _make_table(reward_path, 'rewards')
image_db = _make_table(image_path, 'images')
pending_reward_db = _make_table(pending_reward_path, 'pending_rewards')
pending_confirmations_db = _make_table(pending_confirmations_path, 'pending_confirmations')
users_db = _make_table(users_path, 'users')
tracking_events_db = _make_table(tracking_events_path, 'tracking_events')
child_overrides_db = _make_table(child_overrides_path, 'child_overrides')
chore_schedules_db = _make_table(chore_schedules_path, 'chore_schedules')
task_extensions_db = _make_table(task_extensions_path, 'task_extensions')
refresh_tokens_db = _make_table(refresh_tokens_path, 'refresh_tokens')
push_subscriptions_db = _make_table(push_subscriptions_path, 'push_subscriptions')
digest_action_tokens_db = _make_table(digest_action_tokens_path, 'digest_action_tokens')


# ---------------------------------------------------------------------------
# Index management
# ---------------------------------------------------------------------------

COLLECTION_INDEXES = {
    # NOTE: The model ``id`` field is stored as MongoDB's primary key ``_id``,
    # so no separate unique index on ``id`` is needed.  Only secondary indexes
    # for frequently queried fields are defined here.
    'children': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'tasks': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'routines': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'routine_items': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'routine_schedules': [],
    'routine_extensions': [],
    'rewards': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'images': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'pending_rewards': [
        {'keys': [('child_id', ASCENDING)]},
    ],
    'pending_confirmations': [
        {'keys': [('user_id', ASCENDING)]},
        {'keys': [('child_id', ASCENDING)]},
        {'keys': [('entity_id', ASCENDING), ('entity_type', ASCENDING)]},
    ],
    'users': [],
    'tracking_events': [
        {'keys': [('user_id', ASCENDING)]},
        {'keys': [('child_id', ASCENDING)]},
        {'keys': [('entity_id', ASCENDING), ('entity_type', ASCENDING)]},
    ],
    'child_overrides': [
        {'keys': [('child_id', ASCENDING)]},
        {'keys': [('entity_id', ASCENDING), ('entity_type', ASCENDING)]},
    ],
    'chore_schedules': [
        {'keys': [('user_id', ASCENDING)]},
        {'keys': [('child_id', ASCENDING)]},
    ],
    'task_extensions': [
        {'keys': [('user_id', ASCENDING)]},
        {'keys': [('child_id', ASCENDING)]},
    ],
    'refresh_tokens': [
        {'keys': [('user_id', ASCENDING)]},
        {'keys': [('token', ASCENDING)], 'unique': True, 'sparse': True},
    ],
    'push_subscriptions': [
        {'keys': [('user_id', ASCENDING)]},
    ],
    'digest_action_tokens': [
        {'keys': [('user_id', ASCENDING)]},
        {'keys': [('token', ASCENDING)], 'unique': True, 'sparse': True},
    ],
}


def ensure_mongodb_indexes(client=None, db_name=None):
    """Create required indexes on all MongoDB collections.

    Safe to call repeatedly: MongoDB treats index creation as idempotent.
    """
    if not USE_MONGODB:
        return

    client = client or get_mongo_client()
    db_name = db_name or _mongo_db_name
    db = client[db_name]

    for collection_name, indexes in COLLECTION_INDEXES.items():
        coll = db[collection_name]
        for spec in indexes:
            keys = spec['keys']
            kwargs = {k: v for k, v in spec.items() if k != 'keys'}
            coll.create_index(keys, **kwargs)


# Clear test collections at import time so tests start with a clean slate.
if os.environ.get('DB_ENV', 'prod') == 'test':
    child_db.truncate()
    task_db.truncate()
    routine_db.truncate()
    routine_items_db.truncate()
    routine_schedules_db.truncate()
    routine_extensions_db.truncate()
    reward_db.truncate()
    image_db.truncate()
    pending_reward_db.truncate()
    pending_confirmations_db.truncate()
    users_db.truncate()
    tracking_events_db.truncate()
    child_overrides_db.truncate()
    chore_schedules_db.truncate()
    task_extensions_db.truncate()
    refresh_tokens_db.truncate()
    push_subscriptions_db.truncate()
    digest_action_tokens_db.truncate()