# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
# See the LICENSE file in the root of this repository for details.
# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial

import os
import re
from concurrent.futures import ThreadPoolExecutor
from typing import Union

from jinja2 import Template

from validmind.utils import md_to_html

from ..client_config import client_config
from ..logging import get_logger

__executor = ThreadPoolExecutor()
__prompt = None

logger = get_logger(__name__)


AI_REVISION_NAME = "Generated by ValidMind AI"
DEFAULT_REVISION_NAME = "Default Description"


def _load_prompt():
    global __prompt

    if not __prompt:
        folder_path = os.path.join(os.path.dirname(__file__), "test_result_description")
        with open(os.path.join(folder_path, "system.jinja"), "r") as f:
            system_prompt = f.read()
        with open(os.path.join(folder_path, "user.jinja"), "r") as f:
            user_prompt = f.read()

        __prompt = (Template(system_prompt), Template(user_prompt))

    return __prompt


def prompt_to_message(role, prompt):
    if "[[IMAGE:" not in prompt:
        return {"role": role, "content": prompt}

    content = []

    # Regex pattern to find [[IMAGE:<b64-data>]] markers
    pattern = re.compile(r"\[\[IMAGE:(.*?)\]\]", re.DOTALL)

    last_index = 0
    for match in pattern.finditer(prompt):
        # Text before the image marker
        start, end = match.span()
        if start > last_index:
            content.append({"type": "text", "text": prompt[last_index:start]})

        # Image
        content.append({"type": "image_url", "image_url": {"url": match.group(1)}})

        last_index = end

    # Text after the last image
    if last_index < len(prompt):
        content.append({"type": "text", "text": prompt[last_index:]})

    return {"role": role, "content": content}


class DescriptionFuture:
    """This will be immediately returned from generate_description so that
    the tests can continue to be run in parallel while the description is
    retrieved asynchronously.

    The value will be retrieved later and if its not ready yet, it should
    block until it is.
    """

    def __init__(self, future):
        self._future = future

    def get_description(self):
        if isinstance(self._future, str):
            description = self._future
        else:
            # This will block until the future is completed
            description = self._future.result()

        return md_to_html(description, mathml=True)


def generate_description(
    test_id: str,
    test_description: str,
    test_summary: str,
    metric: Union[float, int] = None,
    figures: list = None,
):
    """Generate the description for the test results"""
    if not test_summary and not figures and not metric:
        raise ValueError(
            "No summary, unit metric or figures provided - cannot generate description"
        )

    # TODO: fix circular import
    from validmind.ai.utils import get_client_and_model

    client, model = get_client_and_model()

    # get last part of test id
    test_name = test_id.split(".")[-1]
    # truncate the test description to save time
    test_description = (
        f"{test_description[:500]}..."
        if len(test_description) > 500
        else test_description
    )

    if metric:
        metric_summary = f"**Metric Value**: {metric}"
        if test_summary:
            test_summary = metric_summary + "\n" + test_summary
        else:
            test_summary = metric_summary

    figures = [] if test_summary else figures

    input_data = {
        "test_name": test_name,
        "test_description": test_description,
        "summary": test_summary,
        "figures": [figure._get_b64_url() for figure in figures],
    }
    system, user = _load_prompt()

    response = client.chat.completions.create(
        model=model,
        temperature=0.0,
        messages=[
            prompt_to_message("system", system.render(input_data)),
            prompt_to_message("user", user.render(input_data)),
        ],
    )

    return response.choices[0].message.content


def background_generate_description(
    test_id: str,
    test_description: str,
    test_summary: str,
    figures: list = None,
    metric: Union[int, float] = None,
):
    def wrapped():
        try:
            return generate_description(
                test_id=test_id,
                test_description=test_description,
                test_summary=test_summary,
                figures=figures,
                metric=metric,
            )
        except Exception as e:
            logger.error(f"Failed to generate description: {e}")

            return test_description

    return DescriptionFuture(__executor.submit(wrapped))


def get_description_metadata(
    test_id,
    default_description,
    summary=None,
    figures=None,
    metric=None,
    prefix="metric_description",
    should_generate=True,
):
    """Get Metadata Dictionary for a Test or Metric Result

    Generates an LLM interpretation of the test results or uses the default
    description and returns a metadata object that can be logged with the test results.

    By default, the description is generated by an LLM that will interpret the test
    results and provide a human-readable description. If the summary or figures are
    not provided, or the `VALIDMIND_LLM_DESCRIPTIONS_ENABLED` environment variable is
    set to `0` or `false` or no LLM has been configured, the default description will
    be used as the test result description.

    Note: Either the summary or figures must be provided to generate the description.

    Args:
        test_id (str): The test ID
        default_description (str): The default description for the test
        summary (Any): The test summary or results to interpret
        figures (List[Figure]): The figures to attach to the test suite result
        metric (Union[int, float]): Unit metrics attached to the test result
        prefix (str): The prefix to use for the content ID (Default: "metric_description")
        should_generate (bool): Whether to generate the description or not (Default: True)

    Returns:
        dict: The metadata object to be logged with the test results
    """
    # Check the feature flag first, then the environment variable
    llm_descriptions_enabled = (
        client_config.can_generate_llm_test_descriptions()
        and os.getenv("VALIDMIND_LLM_DESCRIPTIONS_ENABLED", "1") not in ["0", "false"]
    )

    # TODO: fix circular import
    from validmind.ai.utils import is_configured

    if (
        should_generate
        and (summary or figures)
        and llm_descriptions_enabled
        and is_configured()
    ):
        revision_name = AI_REVISION_NAME

        # get description future and set it as the description in the metadata
        # this will lazily retrieved so it can run in the background in parallel
        description = background_generate_description(
            test_id=test_id,
            test_description=default_description,
            test_summary=summary,
            figures=figures,
            metric=metric,
        )

    else:
        revision_name = DEFAULT_REVISION_NAME
        description = md_to_html(default_description, mathml=True)

    return {
        "content_id": f"{prefix}:{test_id}::{revision_name}",
        "text": description,
    }
