import os
import re
from contextlib import contextmanager
from typing import Any, Literal

import numpy as np
import openai
import ruptures as rpt
import tiktoken
from jinja2 import Template
from packaging import version
from pydantic import BaseModel

from common.utils import get_path, format_string
from common.config import common_settings


def parse_prompts(input_data: str, delim: str = "user") -> tuple[str, str | None]:
    # Check if the input is a file path
    if os.path.exists(input_data):
        with open(input_data) as file:
            markdown_text = file.read()
    else:
        markdown_text = input_data

    # Use a regular expression to split on ':::user' only when it is on a line by itself
    parts = re.split(rf"\n\s*:::{delim}\s*\n", markdown_text, maxsplit=1)

    # Strip leading and trailing whitespace from each part
    system_prompt = parts[0].strip()
    user_prompt = parts[1].strip() if len(parts) > 1 else None

    return system_prompt, user_prompt


OPENAI_MODELS = {
    "gpt-3.5-turbo-0125",
    "gpt-4",
    "gpt-4-turbo-2024-04-09",
    "gpt-4o-2024-05-13",
    "gpt-4-0125-preview",
    "gpt-4-1106-preview",
    "gpt-4-vision-preview",
}

TOGETHER_COMPUTE_MODELS = {
    "mistralai/Mistral-7B-Instruct-v0.2",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
}

CEREBRAS_MODELS = {
    "qwen/qwen3-235b-a22b-2507",
    "cerebras/qwen/qwen3-235b-a22b-2507",
    "cerebras/zai-glm-4.7"
}

ANTHROPIC_MODELS = [
    "claude-3-5-sonnet-20240620",
    "claude-3-5-sonnet-20241022",
    "claude-3-5-haiku-20241022",
    "claude-sonnet-4-20250514",
    "claude-opus-4-20250514",
    "claude-opus-4-1-20250805",
    "claude-sonnet-4-5-20250929",
    "claude-haiku-4-5-20251001",
]

DEFAULT_FLAGSHIP_MODEL = common_settings.DEFAULT_FLAGSHIP_MODEL
DEFAULT_VISION_MODEL = common_settings.DEFAULT_FLAGSHIP_VISION_MODEL  # Deprecated
DEFAULT_FLAGSHIP_VISION_MODEL = common_settings.DEFAULT_FLAGSHIP_VISION_MODEL

DEFAULT_FAST_LONG_CONTEXT_MODEL = common_settings.DEFAULT_FAST_LONG_CONTEXT_MODEL
DEFAULT_FAST_MODEL = common_settings.DEFAULT_FAST_LONG_CONTEXT_MODEL
DEFAULT_FLAGSHIP_LONG_CONTEXT_MODEL = common_settings.DEFAULT_FLAGSHIP_MODEL

DEFAULT_SMALL_EMBEDDING_MODEL = common_settings.DEFAULT_SMALL_EMBEDDING_MODEL
DEFAULT_LARGE_EMBEDDING_MODEL = common_settings.DEFAULT_LARGE_EMBEDDING_MODEL

DEFAULT_TOKENIZER_MODEL = "gpt-4"

TOKENIZER_MODEL_MAPS = {
    "mistralai/Mistral-7B-Instruct-v0.2": "gpt-4",
    "mistralai/Mixtral-8x7B-Instruct-v0.1": "gpt-4",
    "gpt-4-0125-preview": "gpt-4",
    "gpt-4-1106-preview": "gpt-4",
    "gpt-4-turbo-2024-04-09": "gpt-4",
    "gpt-4o-2024-05-13": "gpt-4",
    "gpt-4o-mini-2024-07-18": "gpt-4",
    "claude-3-5-sonnet-20240620": "gpt-4",
    "claude-sonnet-4-20250514": "gpt-4",
    "claude-opus-4-20250514": "gpt-4",
    "claude-opus-4-1-20250805": "gpt-4",
    "claude-sonnet-4-5-20250929": "gpt-4",
    "claude-haiku-4-5-20251001": "gpt-4",
    "gemini/gemini-2.5-pro": "gpt-4",
    "gemini/gemini-3-pro-preview": "gpt-4",
    "vertex_ai/gemini-3-pro-preview": "gpt-4",
    "qwen/qwen3-235b-a22b-2507": "gpt-4",
    "cerebras/qwen/qwen3-235b-a22b-2507": "gpt-4",
}

CONTEXT_LENGTHS = {
    "gpt-3.5-turbo-0125": 16384,
    "gpt-4": 8192,
    "gpt-4-1106-preview": 32768,
    "gpt-4-0125-preview": 32768,
    "gpt-4-turbo-2024-04-09": 32768,
    "gpt-4o-2024-05-13": 65536,
    "gpt-4o-mini-2024-07-18": 65536,
    "gpt-4-vision-preview": 32768,
    "mistralai/Mistral-7B-Instruct-v0.2": 32768,
    "mistralai/Mixtral-8x7B-Instruct-v0.1": 32768,
    "claude-3-5-sonnet-20240620": 65536,
    "claude-sonnet-4-20250514": 65536,
    "claude-opus-4-20250514": 65536,
    "claude-opus-4-1-20250805": 65536,
    "gemini/gemini-2.5-pro": 131072,
    "gemini/gemini-3-pro-preview": 131072,
    "vertex_ai/gemini-3-pro-preview": 131072,
    "claude-sonnet-4-5-20250929": 131072,
    "claude-haiku-4-5-20251001": 131072,
    "qwen/qwen3-235b-a22b-2507": 262144,
    "cerebras/qwen/qwen3-235b-a22b-2507": 262144,
}

DEFAULT_CONTEXT_LENGTH_FALLBACKS = {
    "gpt-4": "gpt-4-turbo-2024-04-09",
}

MODEL_TYPES = {
    "claude-3-5-sonnet-20240620": "instruct",
    "claude-sonnet-4-20250514": "hybrid",
    "claude-opus-4-20250514": "hybrid",
    "claude-opus-4-1-20250805": "hybrid",
    "claude-sonnet-4-5-20250929": "hybrid",
    "claude-haiku-4-5-20251001": "hybrid",
    "gemini/gemini-2.5-pro": "reasoning",
    "gemini/gemini-3-pro-preview": "reasoning",
    "vertex_ai/gemini-3-pro-preview": "reasoning",
    "qwen/qwen3-235b-a22b-2507": "reasoning",
    "cerebras/qwen/qwen3-235b-a22b-2507": "reasoning",
}


def encode(text: str, model_name: str | None = None) -> list[int]:
    model_name = TOKENIZER_MODEL_MAPS.get(model_name, DEFAULT_TOKENIZER_MODEL)
    encoding = tiktoken.encoding_for_model(model_name)
    return encoding.encode(text, allowed_special=encoding.special_tokens_set)


def decode(tokens: list[int], model_name: str | None = None) -> str:
    model_name = TOKENIZER_MODEL_MAPS.get(model_name, DEFAULT_TOKENIZER_MODEL)
    return tiktoken.encoding_for_model(model_name).decode(tokens)


def count_tokens_in_str(text: str, model_name: str | None = None) -> int:
    assert isinstance(text, str)
    return len(encode(text, model_name=model_name))


def num_tokens_from_functions(
    functions: list[dict[str, Any]], model_name: str = DEFAULT_TOKENIZER_MODEL
):
    """Return the number of tokens used by a list of functions."""
    try:
        encoding = tiktoken.encoding_for_model(model_name)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")

    num_tokens = 0
    for function in functions:
        function_tokens = len(encoding.encode(function["name"]))
        function_tokens += len(encoding.encode(function["description"]))

        if "parameters" in function:
            parameters = function["parameters"]
            if "properties" in parameters:
                for propertiesKey in parameters["properties"]:
                    function_tokens += len(encoding.encode(propertiesKey))
                    v = parameters["properties"][propertiesKey]
                    for field in v:
                        if field == "type":
                            function_tokens += 2
                            function_tokens += len(encoding.encode(v["type"]))
                        elif field == "description":
                            function_tokens += 2
                            function_tokens += len(encoding.encode(v["description"]))
                        elif field == "enum":
                            function_tokens -= 3
                            for o in v["enum"]:
                                function_tokens += 3
                                function_tokens += len(encoding.encode(o))
                function_tokens += 11

        num_tokens += function_tokens

    num_tokens += 12
    return num_tokens


def clip_text(
    text: str,
    num_tokens: int,
    model_name: str = DEFAULT_TOKENIZER_MODEL,
    add_truncation_marker: bool = True,
) -> str:
    if text is None:
        return None
    tokens = encode(text, model_name)
    if len(tokens) <= num_tokens:
        return text
    if add_truncation_marker:
        truncation_marker = "[...]"
        num_tokens -= count_tokens_in_str(truncation_marker, model_name)
        return decode(tokens[:num_tokens], model_name) + " " + truncation_marker
    else:
        return decode(tokens[:num_tokens], model_name)


@contextmanager
def using_openai_credentials(api_key: str | None = None, endpoint: str | None = None):
    old_api_key = openai.api_key
    old_endpoint = openai.api_base
    if api_key is not None:
        openai.api_key = api_key
    if endpoint is not None:
        openai.api_base = endpoint
    yield
    openai.api_key = old_api_key
    openai.api_base = old_endpoint


def get_endpoint_and_key_for_model(
    model_name: str,
) -> tuple[str | None, str | None]:
    if model_name in OPENAI_MODELS:
        return None, None
    if model_name in TOGETHER_COMPUTE_MODELS:
        from common.config import settings

        return (
            "https://api.together.xyz",
            settings.TOGETHER_COMPUTE_API_KEY,
        )
    if model_name in CEREBRAS_MODELS:
        from common.config import settings

        return (
            "https://api.cerebras.ai/v1",
            getattr(settings, "CEREBRAS_API_KEY", None),
        )
    raise ValueError(f"Model {model_name} not found")


def to_normed_array(x: list[float] | list[list[float]]) -> np.ndarray:
    x = np.array(x)
    if x.ndim == 2:
        return x / np.linalg.norm(x, axis=1, keepdims=True)
    else:
        return x / np.linalg.norm(x)


def similarity_matrix(
    x: list[float] | list[list[float]], y: list[float] | list[list[float]]
) -> list[float] | list[list[float]]:
    x = to_normed_array(x)
    y = to_normed_array(y)
    sim_mat = x @ y.T
    return sim_mat.tolist()


def normalize_scores(scores: np.ndarray) -> np.ndarray:
    min_score = scores.min()
    score_range = scores.max() - min_score
    if score_range == 0:
        return np.ones_like(scores)
    else:
        return (scores - min_score) / score_range


def find_outlier_threshold(scores: np.ndarray) -> float:
    # Find the histogram of scores
    hist, bin_edges = np.histogram(scores.flatten(), bins=200)
    # Define the algo
    algo = rpt.Pelt(model="l2", min_size=1).fit(hist)
    # Find the change points
    result = algo.predict(pen=10 * hist.mean())
    # Find the last change point
    last_change_point = result[-2]
    # Find the threshold
    threshold = bin_edges[last_change_point]
    return threshold


class OutOfTokenCapError(Exception):
    pass


def find_prompt_path(
    name: str,
    prompt_dir: str,
    version_string: str | None = None,
    extensions: list[str] | None = None
) -> str:
    # Default to .md and .md.j2 if no extensions specified
    if extensions is None:
        extensions = [".md", ".md.j2"]

    # First case: name is a path - check for all provided extensions
    for ext in extensions:
        candidate_path = os.path.join(prompt_dir, f"{name}{ext}")
        if os.path.exists(candidate_path):
            return candidate_path

    # Second case: name is a directory
    candidate_path = os.path.join(prompt_dir, name)
    if version_string is None and os.path.isdir(candidate_path):
        # List all files in the directory that match the provided extensions
        files = [
            f
            for f in os.listdir(candidate_path)
            if any(f.endswith(ext) for ext in extensions)
        ]

        # Sort files based on version numbers
        def extract_version(filename):
            # Remove any of the provided extensions from the filename
            base = filename
            for ext in sorted(extensions, key=len, reverse=True):  # Longest first to handle nested extensions
                if base.endswith(ext):
                    base = base[:-len(ext)]
                    break
            # Extract version after 'v'
            return version.parse(base.split("v")[-1])

        sorted_files = sorted(files, key=extract_version)

        # Return the path of the most recent version
        if sorted_files:
            return os.path.join(candidate_path, sorted_files[-1])
        else:
            raise ValueError(f"No prompt files found in {candidate_path}")

    # Third case: name is a directory with a version preference
    elif version_string is not None:
        assert os.path.isdir(candidate_path)
        # Check for files with the specified version and any of the provided extensions
        for ext in extensions:
            prompt_path = os.path.join(candidate_path, f"v{version_string}{ext}")
            if os.path.exists(prompt_path):
                return prompt_path

        # If no file found with any extension, raise an error
        raise ValueError(f"Version {version_string} not found in {candidate_path}")

    else:
        raise ValueError(
            f"Prompt {name} (version {version_string}) not found in {prompt_dir}"
        )


def resolve_prompt_version_string(
    name: str,
    prompt_dir: str | None,
    version_string: str | None = None,
    extensions: list[str] | None = None,
) -> str | None:
    """Resolve a canonical prompt version string for tracing/telemetry.

    Priority:
    1) Explicit version_string if provided
    2) Infer from resolved prompt filename (e.g., v0.0.1.md -> 0.0.1)
    """
    if version_string is not None:
        return version_string
    if prompt_dir is None:
        return None

    try:
        prompt_path = find_prompt_path(
            name=name,
            prompt_dir=prompt_dir,
            version_string=version_string,
            extensions=extensions,
        )
    except Exception:
        return None

    filename = os.path.basename(prompt_path)
    # Handle nested extensions like .md.j2/.yml.j2, then remove the base extension.
    if filename.endswith(".j2"):
        filename = filename[: -len(".j2")]
    stem, _ = os.path.splitext(filename)
    if not stem:
        return None
    if stem.startswith("v") and len(stem) > 1:
        return stem[1:]
    return stem


class Prompt(BaseModel):
    prompt_type: Literal["string", "jinja"] = "string"
    system_prompt: str
    user_prompt: str | None = None

    def format(self, **kwargs) -> tuple[str, str | None]:
        if self.prompt_type == "string":
            system_prompt = format_string(self.system_prompt, **kwargs)
            user_prompt = (
                format_string(self.user_prompt, **kwargs) if self.user_prompt else None
            )
            return system_prompt, user_prompt
        elif self.prompt_type == "jinja":
            # Render Jinja2 templates
            system_template = Template(self.system_prompt)
            system_prompt = system_template.render(**kwargs)

            user_prompt = None
            if self.user_prompt:
                user_template = Template(self.user_prompt)
                user_prompt = user_template.render(**kwargs)

            return system_prompt, user_prompt
        else:
            raise ValueError(f"Unknown prompt type: {self.prompt_type}")

    def format_system_prompt(self, **kwargs) -> str:
        system_prompt, _ = self.format(**kwargs)
        return system_prompt

    def __str__(self) -> str:
        # This will make it play well with format_string function in common.utils
        return self.system_prompt

    @classmethod
    def as_prompt(cls, obj: Any) -> "Prompt":
        if isinstance(obj, cls):
            return obj
        elif isinstance(obj, str):
            return cls(prompt_type="string", system_prompt=obj)
        elif isinstance(obj, tuple) and len(obj) == 2 and isinstance(obj[0], str):
            return cls(prompt_type="string", system_prompt=obj[0], user_prompt=obj[1])
        elif isinstance(obj, dict) and "system_prompt" in obj:
            return cls(**obj)
        else:
            raise ValueError(f"Cannot convert {obj} to Prompt")

    @classmethod
    def load(
        cls, name: str, *, version_string: str | None = None, prompt_dir: str
    ) -> "Prompt":
        return load_prompt(
            name=name,
            version_string=version_string,
            prompt_dir=prompt_dir,
            as_pydantic_model=True,
        )


def load_prompt(
    name: str,
    version_string: str | None = None,
    system_only: bool = False,
    system_prompt_format_kwargs: dict[str, Any] | None = None,
    user_prompt_format_kwargs: dict[str, Any] | None = None,
    as_pydantic_model: bool | None = None,
    prompt_dir: str = "backend/app/core/intel/prompts",
) -> str | tuple[str, str | None] | Prompt:
    # Load the prompts
    # check if it's a relative path
    if not os.path.isabs(prompt_dir):
        prompt_dir = get_path(prompt_dir)
    prompt_path = find_prompt_path(name, prompt_dir, version_string)

    # Check if it's a Jinja template
    is_jinja_template = prompt_path.endswith(".j2")

    system_prompt, user_prompt = parse_prompts(prompt_path)

    # If it's a Jinja template, we need to return a Prompt model with prompt_type="jinja"
    if is_jinja_template:
        assert as_pydantic_model in (
            True,
            None,
        ), "Jinja templates must be returned as pydantic models"
        # For Jinja templates, we don't apply format_string here
        # The templates will be rendered later when used
        return Prompt(
            prompt_type="jinja",
            system_prompt=system_prompt,
            user_prompt=user_prompt,
        )

    # For non-Jinja templates, apply formatting as before
    if system_prompt_format_kwargs is not None:
        system_prompt = format_string(system_prompt, **system_prompt_format_kwargs)
    if user_prompt is not None and user_prompt_format_kwargs is not None:
        user_prompt = format_string(user_prompt, **user_prompt_format_kwargs)

    # Return based on flags
    if system_only:
        return system_prompt

    # If as_pydantic_model is explicitly set or None for Jinja (which was handled above)
    if as_pydantic_model:
        return Prompt(
            prompt_type="string",
            system_prompt=system_prompt,
            user_prompt=user_prompt,
        )

    # Default behavior for non-Jinja: return tuple
    return system_prompt, user_prompt
