import os

import gridfs
import numpy as np
import cv2

from pymongo import MongoClient
from dotenv import load_dotenv

from schemas.base_schemas import *
from schemas.api_schemas import *

class MongoWorker():
    def __init__(self):
        load_dotenv()
        self.client = MongoClient(host=os.getenv('MONGO_HOST'),
                                  port=int(os.getenv('MONGO_PORT')),
                                  username=os.getenv('MONGO_USER'),
                                  password=os.getenv('MONGO_PASS'))
        self.db = self.client["Face_recognition"]
        self.persons = self.db["persons"]
        self.counters = self.db["counters"]
        self.gfs = gridfs.GridFS(self.db)

    def count_persons(self):
        docs_count = self.persons.count_documents({})
        return docs_count

    def get_and_update_counter(self, counter_name: str) -> int:
        counter = self.counters.find_one_and_update(
            {"counter_name": counter_name},
            {"$inc": {"counter": 1}},
            upsert=True,
            return_document=True)
        return counter["counter"]

    def add_person(self, person_name:str, encodings:ndarray[Any, dtype], image_name:str, face_img:np.ndarray) -> BaseResponse:
        try:
            new_person_id = self.get_and_update_counter("persons")
            new_person = Person(
                person_id = new_person_id,
                person_name = person_name,
                encodings = encodings.tolist(),
                image_name = f"{new_person_id}_{image_name}")

            file_id = self.save_image_to_gridfs(face_img, image_name)
            mongo_result = self.persons.insert_one(new_person.model_dump())
            return BaseResponse(result={"file_id": file_id, "mongo_result": mongo_result})
        except Exception as exception:
            print(exception)
            return BaseResponse(result=exception, error=True)

    def save_image_to_gridfs(self, image_array, filename):
        success, encoded_image = cv2.imencode('.png', image_array)
        if not success:
            raise ValueError("Не удалось закодировать изображение")
        file_id = self.gfs.put(encoded_image.tobytes(), filename=filename)
        return file_id

    def get_all_encodings(self):
        result = {}
        projection = {"person_id": 1, "encodings": 1, "_id": 0}

        for document in self.persons.find({}, projection):
            entry = {document["person_id"]: document["encodings"]}
            result.update(entry)

        return result

    def retrieve_image_from_gridfs(self, file_id):
        try:
            gridfs_file = self.gfs.get(file_id)
            image_bytes = gridfs_file.read()
            image_array = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
            if image_array is None:
                raise ValueError("Не удалось декодировать изображение")
            return image_array
        except gridfs.errors.NoFile:
            raise FileNotFoundError(f"Файл с ID {file_id} не найден в GridFS")