"""curation class """

import os
import time
import asyncio
import sys
import traceback
from bson import ObjectId
import requests
from loguru import logger
from aimmocore import config as conf
from aimmocore.core import load_zoo_dataset, list_zoo_datasets
from aimmocore.core import services as acs
from aimmocore.core import images as aci
from aimmocore.core import utils
from aimmocore.core.database import MongoDB
from aimmocore.core.event import SingletonEventLoop as sel
from aimmocore.core.storages import StorageConfig, StorageType
from aimmocore.server.schemas.datasets import ProcessStatus


class Curation:
    """Curation class for handling dataset curation tasks."""

    def __init__(self, api_key: str):
        """initiailize curation

        Args:
            api_key (str): _description_
        """
        self.subscription_id = "NoSubcriptionId"  # Subscription ID is set after validating the API key
        self.api_key = api_key
        self.version = utils.get_version()
        self.auth_headers = {"Authorization": f"Bearer {self.api_key}", "version": self.version}
        self._valid_api_key()
        self.db = self._get_db_connection()
        self.event_loop = sel.get_instance().get_loop()

        if api_key:
            os.environ["CURATION_API_KEY"] = api_key

    def _get_db_connection(self):
        # have to refactoring..
        db = MongoDB()
        db.connect()
        return db.get_engine()

    def curate_dataset(self, name: str, storage_config: StorageConfig) -> str:
        """
        Synchronously curates a dataset by internally calling an asynchronous method using the event loop.

        Args:
            name (str): The name of the dataset to be curated.
            storage_config (StorageConfig): Configuration object containing settings and parameters
                                            necessary for storing and handling the dataset.
        """
        return self.event_loop.run_until_complete(self._curate_dataset_async(name, storage_config))

    async def _curate_dataset_async(self, name: str, storage_config: StorageConfig) -> str:
        """
        Asynchronously curates a dataset by generating image URLs from a storage configuration,
        and processing these URLs to manage dataset curation.

        Args:
            name (str): The name of the dataset to be curated.
            storage_config (StorageConfig): The configuration object that includes details like storage type
                                            and credentials, which are essential for generating image URLs.

        Returns:
            str: A unique identifier for the curated dataset, typically a dataset ID generated by the system.
        """
        if not name:
            raise ValueError("Dataset name is required")

        self._valid_api_key()
        try:
            image_props = self.get_image_properties(storage_config)
            image_datas = self._generate_image_datas(image_props)
            dataset_id = self._process_dataset(name, image_datas, storage_config)
            if await self._prepare_request_data(dataset_id, image_datas):
                self._request_curation(dataset_id, len(image_datas))
                logger.info(f"Created dataset_id: {dataset_id}")
                self._post_process_dataset(dataset_id)
            return dataset_id
        except Exception as e:  # pylint: disable=broad-except
            logger.error(str(e) + " It will return 'Wrong' as dataset_id")
            logger.error(traceback.format_exc())
            return "Wrong"

    def _generate_image_datas(self, image_props: list):
        """_summary_"""
        return [
            {
                # Hash the filename to generate a unique ID, including the file size to avoid collisions
                "id": str(image["size"]) + utils.hash_filename(image["url"].split("?")[0]),
                "image_url": image["url"],
                "file_size": image["size"],
            }
            for image in image_props
        ]

    @utils.validate_dataset_id
    def _request_curation(self, dataset_id: str, request_file_count: int):
        """_summary_"""
        data = {"dataset_id": dataset_id, "request_file_count": request_file_count}
        file_path = dataset_id + ".txt"
        # Open the file in binary mode
        with open(f"{conf.AIMMOCORE_WORKDIR}/{dataset_id}.txt", "rb") as file:
            files = {"file": (file_path, file)}
            response = requests.post(
                conf.CURATION_UPLOAD_ENDPOINT, headers=self.auth_headers, data=data, files=files, timeout=30
            )
        # Check the response status code
        if response.status_code == 200 and response.text:
            response_data = response.json()
            self.event_loop.run_until_complete(acs.update_dataset_curation_status(self.db, response_data))
        else:
            logger.error(f"Failed to request curation for dataset {dataset_id}: {response.text}")
            sys.exit(1)

    def get_image_properties(self, storage_config: StorageConfig):
        """
        Generates and returns a list of file properties based on the provided storage configuration.
        This function currently supports only Azure storage types.

        Args:
            storage_config (StorageConfig): A configuration object that includes the storage type
                                            and credentials necessary for accessing the storage service.

        Returns:
            list[str]: A list of file properties({url, size}) generated from the storage configuration.

        Raises:
            NotImplementedError: If the storage type specified in the storage configuration is not supported.
                                Currently, only 'AZURE' is implemented.
        """
        if storage_config.storage_type != StorageType.AZURE:
            raise NotImplementedError(f"Image URL generation not implemented for {storage_config.storage_type}")
        return storage_config.generate_image_properties()

    @utils.validate_dataset_id
    async def _prepare_request_data(
        self, dataset_id: str, image_datas: list, model_id: str = conf.DEFAULT_CURATION_MODEL_ID
    ):
        """
        Prepare request data from image properties.

        Args:
            dataset_id (str): The unique identifier for the dataset.
            image_datas (list): List of dictionaries, where each dictionary contains
                                properties of an image such as 'id', 'url', and 'size'.
            model_id (str): The model ID to be used in the request. Defaults to
                            conf.DEFAULT_CURATION_MODEL_ID.

        Returns:
            list: A list of dictionaries formatted for the curation model request.
        """
        model_types = ["meta", "emd"]
        request_datas = [
            {
                "id": image["id"],
                "image_url": image["image_url"],
                "file_size": image["file_size"],
                "model_id": model_id,
                "model_type": model_types,
            }
            for image in image_datas
        ]
        file_path = f"{conf.AIMMOCORE_WORKDIR}/{dataset_id}.txt"
        return await utils.write_to_file(file_path, request_datas)

    async def _process_dataset_async(self, name: str, image_datas: list, storage_config: StorageConfig):
        """Asynchronous method to process dataset by inserting and writing to files."""
        dataset_id = str(ObjectId())

        tasks = [
            acs.insert_raw_files(self.db, image_datas),
            acs.insert_raw_dataset(self.db, name, dataset_id, image_datas, storage_config),
        ]
        await asyncio.gather(*tasks)
        return dataset_id

    def _process_dataset(self, name, request_datas, storage_config):
        """Process dataset by inserting and writing to files."""
        dataset_id = self.event_loop.run_until_complete(
            self._process_dataset_async(name, request_datas, storage_config)
        )
        return dataset_id

    @utils.validate_dataset_id
    def _post_process_dataset(self, dataset_id):
        """Handle post-processing tasks such as thumbnail generation."""
        logger.debug(f"Post-processing dataset {dataset_id}")
        # self.event_loop.run_until_complete(aci.generate_thumbnail(self.auth_headers, self.db, dataset_id))

    @utils.validate_dataset_id
    def tracing_curation_task(self, dataset_id: str, tracing: bool = False) -> str:
        """
        Tracks and updates the curation task status for a given dataset ID.

        Args:
            dataset_id (str): The unique identifier for the dataset whose curation status is to be tracked.
            tracing (bool, optional): If True, continues to trace until completion or until retries are exhausted.
                                    Defaults to False.

        Raises:
            RuntimeError: If the tracing does not complete successfully after the allowed number of retries.

        Returns:
            str: The final status of the curation task.
        """

        logger.debug(f"Tracing curation, curation_id is {dataset_id}")
        url = f"{conf.CURATION_STATUS_ENDPOINT}?dataset_id={dataset_id}"
        while True:
            json_results = utils.make_get_request(url, self.auth_headers)
            status_str = json_results.get("status")

            try:
                status = ProcessStatus(status_str)
            except ValueError:
                logger.error(f"Invalid status '{status_str}' received for dataset {dataset_id}")
                raise

            logger.info(f"Tracing curation task, dataset_id is {dataset_id}, status is {status_str}")

            if status in [ProcessStatus.COMPLETED, ProcessStatus.FAILED, ProcessStatus.ERROR]:
                if status == ProcessStatus.COMPLETED:
                    self.event_loop.run_until_complete(aci.generate_thumbnail(self.auth_headers, self.db, dataset_id))
                    self._update_results(dataset_id, json_results)
                self._update_status(dataset_id, status.value)
                tracing = False

            if not tracing:
                return status.value
            time.sleep(5)  # Wait before retrying if tracing is enabled

    def _valid_api_key(self):
        """valid api key"""
        response = requests.get(conf.CURATION_AUTH_ENDPOINT, headers=self.auth_headers, timeout=30)
        if response.status_code != 200 and response.text:
            raise RuntimeError("Invalid API Key")
        logger.debug("API Key is authorized")
        self.subscription_id = response.json().get("subscription_id")

    @utils.validate_dataset_id
    async def _update_results_async(self, dataset_id: str, results: dict):
        """
        Asynchronously update the curation results in the database.

        Args:
            dataset_id (str): The unique identifier for the dataset.
            results (dict): The results dictionary containing the curation outcomes,
                            specifically the 'emd_results' and the 'status' of the curation.
        """
        processed_count = len(results["emd_results"])
        meta_list = await acs.update_curation_results(self.db, dataset_id, results)
        await acs.update_dataset_info(self.db, dataset_id, results["status"], meta_list, processed_count)

    @utils.validate_dataset_id
    def _update_results(self, dataset_id: str, results: dict):
        """
        Update the curation results in the database by calling the asynchronous update method.


        Args:
            dataset_id (str): The unique identifier for the dataset.
            results (dict): The results dictionary containing the outcomes of the dataset curation.

        """
        # Schedule and run the asynchronous update operations as a single coroutine
        task = self.event_loop.create_task(self._update_results_async(dataset_id, results))
        self.event_loop.run_until_complete(task)

    @utils.validate_dataset_id
    def _update_status(self, dataset_id: str, status: str):
        """update status

        Args:
            dataset_id (str): dataset id
            status (str): curation process status
        """
        self.event_loop.run_until_complete(acs.update_dataset_info_status(self.db, dataset_id, status))

    def load_sample_dataset(self, dataset_name: str = "aimmo-ad-dataset"):
        """load sample dataset

        Args:
            dataset_name (str): sample dataset name
        """
        self.event_loop.run_until_complete(load_zoo_dataset(dataset_name))

    def list_sample_datasets(self) -> list:
        """list sample dataset

        Args:
            dataset_name (str): sample dataset name
        """
        return list_zoo_datasets()
