from tinydb import TinyDB, Query, where
from tinydb.storages import MemoryStorage
import chromadb
from chromadb.api.types import EmbeddingFunction, Embeddings, Image, Images
from keras_facenet import FaceNet
from typing import Any
from datetime import datetime, timedelta

CHROMADB_LOC = "/home/user/data/chromadb"

class TinyDBHelper:
    def __init__(self):
        self.db = TinyDB(storage=MemoryStorage)
        self.tokens_table = self.db.table('tokens')

    def insert_token(self, user_id: str, token: str, expires_at: str):
        self.tokens_table.insert({'user_id': user_id, 'token': token, 'expires_at': expires_at})

    def query_token(self, user_id: str, token: str) -> bool:
        """Query to check if the token exists and is valid."""
        User = Query()
        # Assuming our tokens table contains 'user_id', 'token', and 'expires_at'
        result = self.tokens_table.search((User.user_id == user_id) & (User.token == token))
        # Optionally, check if the token is expired
        expires_at = datetime.fromisoformat(result[0]['expires_at'])
        if datetime.utcnow() > expires_at:
            return False
        
        return bool(result)

    def remove_token_by_value(self, token: str):
        """Remove a token based on its value."""
        self.tokens_table.remove((where('token') == token))

###### Class implementing Custom Embedding function for chroma db
#
class UserFaceEmbeddingFunction(EmbeddingFunction[Images]):
    def __init__(self):
        # Intitialize the FaceNet model
        self.facenet = FaceNet()

    def __call__(self, input: Images) -> Embeddings:
        # Since the input images are assumed to be `numpy.ndarray` objects already,
        # we can directly use them for embeddings extraction without additional processing.
        # Ensure the input images are pre-cropped face images ready for embedding extraction.

        # Extract embeddings using FaceNet for the pre-cropped face images.
        embeddings_array = self.facenet.embeddings(input)

        # Convert numpy array of embeddings to list of lists, as expected by Chroma.
        return embeddings_array.tolist()


# Usage example:
# user_face_embedding_function = UserFaceEmbeddingFunction()
# Assuming `images` is a list of `numpy.ndarray` objects where each represents a pre-cropped face image ready for embedding extraction.
# embeddings = user_face_embedding_function(images)


class ChromaDBFaceHelper:
    def __init__(self, db_path: str):
        self.client = chromadb.PersistentClient(db_path)
        self.user_faces_db = self.client.get_or_create_collection(name="user_faces_db", embedding_function=UserFaceEmbeddingFunction())

    def query_user_face(self, presented_face: Any, n_results: int = 1):
        return self.user_faces_db.query(query_images=[presented_face], n_results=n_results)
    
    def print_query_results(self, query_results: dict) -> None:
        for id, distance, metadata in zip(query_results["ids"][0], query_results['distances'][0], query_results['metadatas'][0]):
            print(f'id: {id}, distance: {distance}, metadata: {metadata}')


# Initialize these helpers globally if they are to be used across multiple modules
tinydb_helper = TinyDBHelper()
chromadb_face_helper = ChromaDBFaceHelper(CHROMADB_LOC)  # Initialization requires db_path