import json
import os
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 (
    ChatCompletionMessage,
    Folder,
    FolderEncoder,
    Node,
    Prompt,
    PromptChain,
    PromptChainVersion,
    PromptChainVersionsAndRules,
    PromptNode,
    PromptResponse,
    PromptVersion,
    QueryRule,
    RuleGroupType,
    RuleType,
    Tool,
    VersionAndRulesWithPromptChainIdEncoder,
    VersionAndRulesWithPromptId,
    VersionAndRulesWithPromptIdEncoder,
    VersionsAndRules,
)
from .runnable import RunnablePrompt, RunnablePromptChain
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: Optional[str] = os.environ.get("MAXIM_API_KEY")
    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 = Config()):
        """
        Initializes a new instance of the Maxim class.

        Args:
            config (Config): The configuration for the Maxim instance.
        """
        if config.api_key is None:
            raise ValueError("API key is required. Either set it in the environment variable MAXIM_API_KEY or pass it as a parameter to the Config object.")

        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_version: PromptVersion) -> Prompt:
        return Prompt(prompt_id=prompt_version.promptId,version_id=prompt_version.id,version=prompt_version.version,messages=prompt_version.config.messages if prompt_version.config else [],tags=prompt_version.config.tags if prompt_version.config else {},model_parameters=prompt_version.config.modelParameters if prompt_version.config else {},model=prompt_version.config.model if prompt_version.config else None)

    def __format_prompt_chain(
        self, prompt_chain_version: PromptChainVersion
    ) -> PromptChain:
        prompt_nodes: List[Node] = []
        if prompt_chain_version.config is not None:
            for node in prompt_chain_version.config.nodes:
                if isinstance(node.content, PromptNode):
                    prompt_nodes.append(node)
        return PromptChain(
            prompt_chain_id=prompt_chain_version.promptChainId,
            version_id=prompt_chain_version.id,
            version=prompt_chain_version.version,
            nodes=prompt_nodes,
        )

    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.exact_match
            )
            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 is 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.exact_match
            )
            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: Optional[PromptChainVersion] = None
                    for v in prompt_chain_version_and_rules.versions:
                        if v.id != version_id:
                            continue
                        version = v
                        break
                    if not version:
                        continue
                    if version_rule.rules.query is not None:
                        query = version_rule.rules.query
                        if isinstance(query, RuleGroupType):
                            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 is not 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.exact_match
        )
        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[RunnablePrompt]:
        """
        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
        # wrapping this prompt with Runnable
        return RunnablePrompt(prompt,self.maxim_api)

    def get_prompts(self, rule: QueryRule) -> List[RunnablePrompt]:
        """
        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:list[RunnablePrompt] = []
        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(RunnablePrompt(prompt,self.maxim_api))
        if len(prompts) == 0:
            return []
        return prompts

    def get_prompt_chain(
        self, id: str, rule: QueryRule
    ) -> Optional[RunnablePromptChain]:
        """
        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_chain = self.__get_prompt_chain_version_for_rule(version_and_rules, rule)
        if not prompt_chain:
            return None
        # The reason we have created a separate class
        # called agents is for two reasons
        # 1. We are not able to import MaximAPI in models fir cicular deps
        # 2. We want to set the foundations of agents for future
        return RunnablePromptChain(prompt_chain, maxim_api=self.maxim_api)

    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: Optional[LoggerConfig] = None) -> Logger:
        """
        Creates a logger based on the provided configuration.

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

        Returns:
            Logger: The logger object.
        """
        if config is None:
            # Fetching log repo id from env
            repo_id = os.environ.get("MAXIM_LOG_REPO_ID")
            if repo_id is None:
                raise ValueError("Log repo id is required. Either set environment variable MAXIM_LOG_REPO_ID or pass it as a parameter in LoggerConfig.")
            config = LoggerConfig(id=repo_id)
        # 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 chat_completion(
        self,
        model: str,
        messages: List[ChatCompletionMessage],
        tools: Optional[List[Tool]] = None,
        **kwargs,
    ) -> Optional[PromptResponse]:
        """
        Performs a chat completion request using the specified model and messages.

        Args:
            model (str): The model name to use for completion. The expected format is "provider/model_name". Example "openai/gpt-3.5-turbo".
            messages (List[ChatCompletionMessage]): List of chat messages in the conversation
            tools (Optional[List[Tool]], optional): List of tools available to the model. Defaults to None.
            **kwargs: Additional model parameters to pass to the completion request

        Returns:
            Optional[PromptResponse]: The completion response if successful, None otherwise
        """
        return self.maxim_api.run_prompt(model,messages,tools,**kwargs)

    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")
