import json
import logging
import threading
import time
from dataclasses import dataclass
from typing import Dict, List, Optional

from .apis import MaximAPI
from .cache import MaximCache, MaximInMemoryCache
from .filter_objects import (
    IncomingQuery,
    QueryObject,
    find_all_matches,
    find_best_match,
    parse_incoming_query,
)
from .logger import Logger, LoggerConfig
from .models import (
    Folder,
    FolderEncoder,
    Prompt,
    PromptChain,
    PromptChainVersion,
    PromptChainVersionsAndRules,
    PromptVersion,
    QueryRule,
    RuleGroupType,
    RuleType,
    VersionAndRulesWithPromptChainIdEncoder,
    VersionAndRulesWithPromptId,
    VersionAndRulesWithPromptIdEncoder,
    VersionsAndRules,
)
from .test_runs import TestRunBuilder


@dataclass
class Config:
    """
    A class representing the configuration for the Maxim SDK.

    Attributes:
        api_key (str): The API key for the Maxim instance.
        base_url (Optional[str], optional): The base URL for the Maxim instance. Defaults to "https://app.getmaxim.ai".
        cache (Optional[MaximCache], optional): The cache to use for the Maxim instance. Defaults to None.
        debug (Optional[bool], optional): Whether to enable debug logging. Defaults to False.
        raise_exceptions (Optional[bool], optional): Whether to raise exceptions during logging operations. Defaults to False.
        prompt_management (Optional[bool], optional): Whether to enable prompt management. Defaults to False.
    """

    api_key: str
    base_url: Optional[str] = None
    cache: Optional[MaximCache] = None
    debug: Optional[bool] = False
    raise_exceptions: Optional[bool] = False
    prompt_management: Optional[bool] = False


EntityType = {"PROMPT": "PROMPT", "FOLDER": "FOLDER", "PROMPT_CHAIN": "PROMPT_CHAIN"}


class Maxim:
    """
    A class representing the Maxim SDK.

    This class provides methods for interacting with the Maxim API.
    """

    def __init__(self, config: Config):
        """
        Initializes a new instance of the Maxim class.

        Args:
            config (Config): The configuration for the Maxim instance.
        """
        self.base_url = config.base_url or "https://app.getmaxim.ai"
        self.api_key = config.api_key
        self.is_running = True
        self.raise_exceptions = config.raise_exceptions or False
        self.maxim_api = MaximAPI(self.base_url, self.api_key)
        self.__is_debug = config.debug or False
        self.__loggers: Dict[str, Logger] = {}
        self.prompt_management = config.prompt_management
        if self.prompt_management:
            self.__cache = config.cache or MaximInMemoryCache()
            self.__sync_thread = threading.Thread(target=self.__sync_timer)
            self.__sync_thread.daemon = True
            self.__sync_thread.start()

    def enable_prompt_management(self) -> "Maxim":
        self.prompt_management = True
        if not self.__sync_thread.is_alive():
            self.__sync_thread = threading.Thread(target=self.__sync_timer)
            self.__sync_thread.daemon = True
            self.__sync_thread.start()
        return self

    def enable_exceptions(self, val: bool) -> "Maxim":
        self.raise_exceptions = val
        return self

    def __sync_timer(self):
        if not self.prompt_management:
            return
        while self.is_running:
            self.__sync_entities()
            time.sleep(60)

    def __sync_entities(self):
        logging.debug(" Syncing prompts and folders")
        if not self.prompt_management:
            return
        if not self.is_running:
            return
        self.__sync_prompts()
        self.__sync_prompt_chains()
        self.__syncFolders()
        logging.debug("Syncing completed")

    def __sync_prompts(self):
        logging.debug("Syncing prompts")
        try:
            prompts = self.maxim_api.get_prompts()
            logging.debug(f"Found {len(prompts)} prompts")
            for prompt in prompts:
                try:
                    self.__cache.set(
                        self.__get_cache_key(EntityType["PROMPT"], prompt.promptId),
                        json.dumps(prompt, cls=VersionAndRulesWithPromptIdEncoder),
                    )
                except Exception as err:
                    logging.error(err)
                    if self.raise_exceptions:
                        raise err
        except Exception as err:
            logging.error(f"Error while syncing prompts: {err}")
            if self.raise_exceptions:
                raise err

    def __sync_prompt_chains(self):
        logging.debug("Syncing prompt Chains")
        try:
            prompt_chains = self.maxim_api.get_prompt_chains()
            logging.debug(f"Found {len(prompt_chains)} prompt chains")
            for prompt_chain in prompt_chains:
                try:
                    self.__cache.set(
                        self.__get_cache_key(
                            EntityType["PROMPT_CHAIN"], prompt_chain.promptChainId
                        ),
                        json.dumps(
                            prompt_chain, cls=VersionAndRulesWithPromptChainIdEncoder
                        ),
                    )
                except Exception as err:
                    logging.error(err)
                    if self.raise_exceptions:
                        raise err
        except Exception as err:
            logging.error(f"Error while syncing prompts: {err}")
            if self.raise_exceptions:
                raise err

    def __syncFolders(self):
        logging.debug("Syncing folders")
        try:
            folders = self.maxim_api.get_folders()
            logging.debug(f"Found {len(folders)} folders")
            for folder in folders:
                try:
                    self.__cache.set(
                        self.__get_cache_key(EntityType["FOLDER"], folder.id),
                        json.dumps(folder, cls=FolderEncoder),
                    )
                except Exception as err:
                    logging.error(err)
                    if self.raise_exceptions:
                        raise err
        except Exception as err:
            logging.error(f"Error while syncing folders: {err}")
            if self.raise_exceptions:
                raise err

    def __get_cache_key(self, entity: str, id: str) -> str:
        if entity == "PROMPT":
            return f"prompt:{id}"
        else:
            return f"folder:{id}"

    def __get_prompt_from_cache(self, key: str) -> Optional[VersionsAndRules]:
        data = self.__cache.get(key)
        if not data:
            return None
        json_data = json.loads(data)
        return VersionAndRulesWithPromptId.from_dict(json_data)

    def __get_all_prompts_from_cache(self) -> Optional[List[VersionsAndRules]]:
        keys = self.__cache.get_all_keys()
        if not keys:
            return None
        data = [self.__cache.get(key) for key in keys if key.startswith("prompt:")]
        prompt_list = []
        for d in data:
            if d is not None:
                json_data = json.loads(d)
                if "promptId" in json_data:
                    json_data.pop("promptId")
                prompt_list.append(VersionsAndRules.from_dict(json_data))
        return prompt_list

    def __get_prompt_chain_from_cache(
        self, key: str
    ) -> Optional[PromptChainVersionsAndRules]:
        data = self.__cache.get(key)
        if not data:
            return None
        parsed_data = json.loads(data)
        return PromptChainVersionsAndRules.from_dict(parsed_data)

    def __get_all_prompt_chains_from_cache(
        self,
    ) -> Optional[List[PromptChainVersionsAndRules]]:
        keys = self.__cache.get_all_keys()
        if not keys:
            return None
        data = [
            self.__cache.get(key) for key in keys if key.startswith("prompt_chain:")
        ]
        prompt_chain_list = []
        for d in data:
            if d is not None:
                json_data = json.loads(d)
                if "promptChainId" in json_data:
                    json_data.pop("promptChainId")
                prompt_chain_list.append(
                    PromptChainVersionsAndRules.from_dict(json_data)
                )
        return prompt_chain_list

    def __get_folder_from_cache(self, key: str) -> Optional[Folder]:
        data = self.__cache.get(key)
        if not data:
            return None
        json_data = json.loads(data)
        return Folder(**json_data)

    def __get_all_folders_from_cache(self) -> Optional[List[Folder]]:
        keys = self.__cache.get_all_keys()
        if not keys:
            return None
        data = [self.__cache.get(key) for key in keys if key.startswith("folder:")]
        return [Folder(**json.loads(d)) for d in data if d is not None]

    def __format_prompt(self, prompt: PromptVersion) -> Prompt:
        json_data = {
            "promptId": prompt.promptId,
            "version": prompt.version,
            "versionId": prompt.id,
            "tags": prompt.config.tags if prompt.config else [],
            "model": prompt.config.model if prompt.config else None,
            "messages": prompt.config.messages if prompt.config else None,
            "modelParameters": prompt.config.modelParameters if prompt.config else None,
        }
        return Prompt(**json_data)

    def __format_prompt_chain(
        self, prompt_chain_version: PromptChainVersion
    ) -> PromptChain:
        json_data = {
            "promptChainId": prompt_chain_version.promptChainId,
            "versionId": prompt_chain_version.id,
            "version": prompt_chain_version.version,
            "nodes": (
                [
                    n
                    for n in prompt_chain_version.config.nodes
                    if hasattr(n, "prompt") or (isinstance(n, dict) and "prompt" in n)
                ]
                if prompt_chain_version.config
                else []
            ),
        }
        return PromptChain(**json_data)

    def __get_prompt_version_for_rule(
        self,
        prompt_version_and_rules: VersionsAndRules,
        rule: Optional[QueryRule] = None,
    ) -> Optional[Prompt]:
        if rule:
            incoming_query = IncomingQuery(
                query=rule.query, operator=rule.operator, exactMatch=rule.exactMatch
            )
            objects = []
            for version_id, version_rules in prompt_version_and_rules.rules.items():
                for version_rule in version_rules:
                    if not version_rule.rules.query:
                        continue
                    if rule.scopes:
                        for key in rule.scopes.keys():
                            if key != "folder":
                                raise ValueError("Invalid scope added")
                    version = next(
                        (
                            v
                            for v in prompt_version_and_rules.versions
                            if v.id == version_id
                        ),
                        None,
                    )
                    if not version:
                        continue
                    query = version_rule.rules.query
                    if version.config and version.config.tags:
                        parsed_incoming_query = parse_incoming_query(
                            incoming_query.query
                        )
                        tags = version.config.tags
                        for key, value in tags.items():
                            if value == None:
                                continue
                            if parsed_incoming_query is None:
                                continue
                            incomingQueryFields = [
                                q.field for q in parsed_incoming_query
                            ]
                            if key in incomingQueryFields:
                                query.rules.append(
                                    RuleType(field=key, operator="=", value=value)
                                )
                    objects.append(QueryObject(query=query, id=version_id))

            deployed_version_object = find_best_match(objects, incoming_query)
            if deployed_version_object:
                deployed_version = next(
                    (
                        v
                        for v in prompt_version_and_rules.versions
                        if v.id == deployed_version_object.id
                    ),
                    None,
                )
                if deployed_version:
                    return self.__format_prompt(deployed_version)

        else:
            if prompt_version_and_rules.rules:
                for version_id, version_rules in prompt_version_and_rules.rules.items():
                    logging.debug(f"type: {type(version_rules[0].rules)}")
                    does_query_exist = any(
                        ruleElm.rules.query is not None for ruleElm in version_rules
                    )
                    if does_query_exist:
                        deployed_version = next(
                            (
                                v
                                for v in prompt_version_and_rules.versions
                                if v.id == version_id
                            ),
                            None,
                        )
                        if deployed_version:
                            return self.__format_prompt(deployed_version)
            else:
                return self.__format_prompt(prompt_version_and_rules.versions[0])

        if prompt_version_and_rules.fallbackVersion:
            logging.info(
                f"************* FALLBACK TRIGGERED : VERSION : {prompt_version_and_rules.fallbackVersion} *************"
            )
            return self.__format_prompt(prompt_version_and_rules.fallbackVersion)
        return None

    def __get_prompt_chain_version_for_rule(
        self,
        prompt_chain_version_and_rules: PromptChainVersionsAndRules,
        rule: Optional[QueryRule] = None,
    ) -> Optional[PromptChain]:
        if rule:
            incoming_query = IncomingQuery(
                query=rule.query, operator=rule.operator, exactMatch=rule.exactMatch
            )
            objects = []
            for (
                version_id,
                version_rules,
            ) in prompt_chain_version_and_rules.rules.items():
                for version_rule in version_rules:
                    if not version_rule.rules.query:
                        continue
                    if rule.scopes:
                        for key in rule.scopes.keys():
                            if key != "folder":
                                raise ValueError("Invalid scope added")
                    version = next(
                        (
                            v
                            for v in prompt_chain_version_and_rules.versions
                            if v.id == version_id
                        ),
                        None,
                    )
                    if not version:
                        continue
                    query = RuleGroupType(**version_rule.rules.query.__dict__)
                    objects.append(QueryObject(query=query, id=version_id))

            deployed_version_object = find_best_match(objects, incoming_query)
            if deployed_version_object:
                deployed_version = next(
                    (
                        v
                        for v in prompt_chain_version_and_rules.versions
                        if v.id == deployed_version_object.id
                    ),
                    None,
                )
                if deployed_version:
                    return self.__format_prompt_chain(deployed_version)

        else:
            if prompt_chain_version_and_rules.rules:
                for (
                    version_id,
                    version_rules,
                ) in prompt_chain_version_and_rules.rules.items():
                    logging.debug(f"type: {type(version_rules[0].rules)}")
                    does_query_exist = any(
                        ruleElm.rules.query != None for ruleElm in version_rules
                    )
                    if does_query_exist:
                        deployed_version = next(
                            (
                                v
                                for v in prompt_chain_version_and_rules.versions
                                if v.id == version_id
                            ),
                            None,
                        )
                        if deployed_version:
                            return self.__format_prompt_chain(deployed_version)
            else:
                return self.__format_prompt_chain(
                    prompt_chain_version_and_rules.versions[0]
                )

        if prompt_chain_version_and_rules.fallbackVersion:
            logging.info(
                f"************* FALLBACK TRIGGERED : VERSION : {prompt_chain_version_and_rules.fallbackVersion} *************"
            )
            return self.__format_prompt_chain(
                prompt_chain_version_and_rules.fallbackVersion
            )
        return None

    def __get_folders_for_rule(
        self, folders: List[Folder], rule: QueryRule
    ) -> List[Folder]:
        incoming_query = IncomingQuery(
            query=rule.query, operator=rule.operator, exactMatch=rule.exactMatch
        )
        objects = []
        for folder in folders:
            query = RuleGroupType(rules=[], combinator="AND")
            if not folder.tags:
                continue
            parsed_incoming_query = parse_incoming_query(incoming_query.query)
            tags = folder.tags
            for key, value in tags.items():
                if key in [q.field for q in parsed_incoming_query]:
                    # if not isinstance(value,None):
                    query.rules.append(RuleType(field=key, operator="=", value=value))
            if not query.rules:
                continue
            objects.append(QueryObject(query=query, id=folder.id))

        folder_objects = find_all_matches(objects, incoming_query)
        ids = [fo.id for fo in folder_objects]
        return [f for f in folders if f.id in ids]

    def get_prompt(self, id: str, rule: QueryRule) -> Optional[Prompt]:
        """
        Retrieves a prompt based on the provided id and rule.

        Args:
            id (str): The id of the prompt.
            rule (QueryRule): The rule to match the prompt against.

        Returns:
            Optional[Prompt]: The prompt object if found, otherwise None.
        """
        if self.prompt_management is False:
            raise Exception(
                "prompt_management is disabled. You can enable it by initializing Maxim with Config(...prompt_management=True)."
            )
        logging.debug(f"getPrompt: id: {id} for rule: {rule}")
        key = self.__get_cache_key("PROMPT", id)
        version_and_rules_with_prompt_id = self.__get_prompt_from_cache(key)
        if version_and_rules_with_prompt_id is None:
            version_and_rules_with_prompt_id = self.maxim_api.get_prompt(id)
            if len(version_and_rules_with_prompt_id.versions) == 0:
                return None
            self.__cache.set(
                id,
                json.dumps(
                    version_and_rules_with_prompt_id,
                    cls=VersionAndRulesWithPromptIdEncoder,
                ),
            )
        if not version_and_rules_with_prompt_id:
            return None
        prompt = self.__get_prompt_version_for_rule(
            version_and_rules_with_prompt_id, rule
        )
        if not prompt:
            return None
        return prompt

    def get_prompts(self, rule: QueryRule) -> List[Prompt]:
        """
        Retrieves all prompts that match the given rule.

        Args:
            rule (QueryRule): The rule to match the prompts against.

        Returns:
            List[Prompt]: A list of prompts that match the given rule.
        """
        if self.prompt_management is False:
            raise Exception(
                "prompt_management is disabled. Please enable it in config."
            )
        version_and_rules = self.__get_all_prompts_from_cache()
        prompts = []
        if version_and_rules is None or len(version_and_rules) == 0:
            self.__sync_entities()
            version_and_rules = self.__get_all_prompts_from_cache()
        if not version_and_rules:
            return []
        for v in version_and_rules:
            if rule.scopes:
                if "folder" not in rule.scopes.keys():
                    return []
                else:
                    if rule.scopes["folder"] != v.folderId:
                        continue
            prompt = self.__get_prompt_version_for_rule(v, rule)
            if prompt is not None:
                prompts.append(prompt)
        if len(prompts) == 0:
            return []
        return prompts

    def get_prompt_chain(self, id: str, rule: QueryRule) -> Optional[PromptChain]:
        """
        Retrieves a prompt chain based on the provided id and rule.

        Args:
            id (str): The id of the prompt chain.
            rule (QueryRule): The rule to match the prompt chain against.

        Returns:
            Optional[PromptChain]: The prompt chain object if found, otherwise None.
        """
        if self.prompt_management is False:
            raise Exception(
                "prompt_management is disabled. Please enable it in config by setting Config(...,prompt_management=True)."
            )
        key = self.__get_cache_key("PROMPT_CHAIN", id)
        version_and_rules = self.__get_prompt_chain_from_cache(key)
        if version_and_rules is None:
            version_and_rules = self.maxim_api.getPromptChain(id)
            if len(version_and_rules.versions) == 0:
                return None
            self.__cache.set(
                id,
                json.dumps(
                    version_and_rules, cls=VersionAndRulesWithPromptChainIdEncoder
                ),
            )
        if not version_and_rules:
            return None
        prompt_chains = self.__get_prompt_chain_version_for_rule(
            version_and_rules, rule
        )
        if not prompt_chains:
            return None
        return prompt_chains

    def get_folder_by_id(self, id: str) -> Optional[Folder]:
        """
        Retrieves a folder based on the provided id.

        Args:
            id (str): The id of the folder.

        Returns:
            Optional[Folder]: The folder object if found, otherwise None.
        """
        if self.prompt_management is False:
            raise Exception(
                "prompt_management is disabled. Please enable it in config."
            )
        key = self.__get_cache_key("FOLDER", id)
        folder = self.__get_folder_from_cache(key)
        if folder is None:
            try:
                folder = self.maxim_api.get_folder(id)
                if not folder:
                    return None
                self.__cache.set(key, json.dumps(folder, cls=FolderEncoder))
            except Exception:
                return None
        return folder

    def get_folders(self, rule: QueryRule) -> List[Folder]:
        """
        Retrieves all folders that match the given rule.

        Args:
            rule (QueryRule): The rule to match the folders against.

        Returns:
            List[Folder]: A list of folders that match the given rule.
        """
        if self.prompt_management is False:
            raise Exception(
                "prompt_management is disabled. Please enable it in config."
            )
        folders = self.__get_all_folders_from_cache()
        if folders is None or len(folders) == 0:
            self.__sync_entities()
            folders = self.__get_all_folders_from_cache()
        if not folders:
            return []
        return self.__get_folders_for_rule(folders, rule)

    def logger(self, config: LoggerConfig) -> Logger:
        """
        Creates a logger based on the provided configuration.

        Args:
            config (LoggerConfig): The configuration for the logger.

        Returns:
            Logger: The logger object.
        """
        # Checking if this log repository exist on server
        exists = self.maxim_api.does_log_repository_exist(config.id)
        if not exists:
            if config.id:
                logging.error("Log repository not found")
                if self.raise_exceptions:
                    raise Exception("Log repository not found")
        if config.id in self.__loggers:
            return self.__loggers[config.id]
        logger = Logger(
            config=config,
            api_key=self.api_key,
            base_url=self.base_url,
            is_debug=self.__is_debug,
            raise_exceptions=self.raise_exceptions,
        )
        self.__loggers[config.id] = logger
        return logger

    def create_test_run(self, name: str, in_workspace_id: str) -> TestRunBuilder:
        """
        Creates a test run builder based on the provided name and workspace id.

        Args:
            name (str): The name of the test run.
            in_workspace_id (str): The workspace id to create the test run in.

        Returns:
            TestRunBuilder: The test run builder object.
        """
        return TestRunBuilder(
            name=name,
            workspace_id=in_workspace_id,
            api_key=self.api_key,
            base_url=self.base_url,
            evaluators=[],
        )

    def cleanup(self):
        """
        Cleans up the Maxim sync thread.
        """
        # We sleep to allow all pending writes in the queue to come in and then we can flush
        time.sleep(2)
        logging.debug("Cleaning up Maxim sync thread")
        self.is_running = False
        for logger in self.__loggers.values():
            logger.cleanup()
        logging.debug("Cleanup done")
