"""
BitBucket prompt manager that integrates with LiteLLM's prompt management system.
Fetches .prompt files from BitBucket repositories and provides team-based access control.
"""

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment

from litellm.integrations.custom_prompt_management import CustomPromptManagement

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
else:
    LiteLLMLoggingObj = Any
from litellm.integrations.prompt_management_base import (
    PromptManagementBase,
    PromptManagementClient,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.utils import StandardCallbackDynamicParams

from .bitbucket_client import BitBucketClient


class BitBucketPromptTemplate:
    """
    Represents a prompt template loaded from BitBucket.
    """

    def __init__(
        self,
        template_id: str,
        content: str,
        metadata: Dict[str, Any],
        model: Optional[str] = None,
    ):
        self.template_id = template_id
        self.content = content
        self.metadata = metadata
        self.model = model or metadata.get("model")
        self.temperature = metadata.get("temperature")
        self.max_tokens = metadata.get("max_tokens")
        self.input_schema = metadata.get("input", {}).get("schema", {})
        self.optional_params = {
            k: v for k, v in metadata.items() if k not in ["model", "input", "content"]
        }

    def __repr__(self):
        return f"BitBucketPromptTemplate(id='{self.template_id}', model='{self.model}')"


class BitBucketTemplateManager:
    """
    Manager for loading and rendering .prompt files from BitBucket repositories.

    Supports:
    - Fetching .prompt files from BitBucket repositories
    - Team-based access control through BitBucket permissions
    - YAML frontmatter for metadata
    - Handlebars-style templating (using Jinja2)
    - Input/output schema validation
    - Model configuration
    """

    def __init__(
        self,
        bitbucket_config: Dict[str, Any],
        prompt_id: Optional[str] = None,
    ):
        self.bitbucket_config = bitbucket_config
        self.prompt_id = prompt_id
        self.prompts: Dict[str, BitBucketPromptTemplate] = {}
        self.bitbucket_client = BitBucketClient(bitbucket_config)

        # Templates fetched from a BitBucket repo are not trustworthy:
        # anyone with repo write access can ship Jinja syntax that, in a
        # plain `Environment()`, would reach `__class__.__init__.__globals__`
        # and pivot into RCE on the proxy host. The sandbox blocks that
        # attribute traversal while leaving normal `{{ var }}` substitution
        # intact. Matches the dotprompt manager's hardening.
        self.jinja_env = ImmutableSandboxedEnvironment(
            loader=DictLoader({}),
            autoescape=select_autoescape(["html", "xml"]),
            # Use Handlebars-style delimiters to match Dotprompt spec
            variable_start_string="{{",
            variable_end_string="}}",
            block_start_string="{%",
            block_end_string="%}",
            comment_start_string="{#",
            comment_end_string="#}",
        )

        # Load prompts from BitBucket if prompt_id is provided
        if self.prompt_id:
            self._load_prompt_from_bitbucket(self.prompt_id)

    def _load_prompt_from_bitbucket(self, prompt_id: str) -> None:
        """Load a specific .prompt file from BitBucket."""
        try:
            # Fetch the .prompt file from BitBucket
            prompt_content = self.bitbucket_client.get_file_content(
                f"{prompt_id}.prompt"
            )

            if prompt_content:
                template = self._parse_prompt_file(prompt_content, prompt_id)
                self.prompts[prompt_id] = template
        except Exception as e:
            raise Exception(f"Failed to load prompt '{prompt_id}' from BitBucket: {e}")

    def _parse_prompt_file(
        self, content: str, prompt_id: str
    ) -> BitBucketPromptTemplate:
        """Parse a .prompt file content and extract metadata and template."""
        # Split frontmatter and content
        if content.startswith("---"):
            parts = content.split("---", 2)
            if len(parts) >= 3:
                frontmatter_str = parts[1].strip()
                template_content = parts[2].strip()
            else:
                frontmatter_str = ""
                template_content = content
        else:
            frontmatter_str = ""
            template_content = content

        # Parse YAML frontmatter
        metadata: Dict[str, Any] = {}
        if frontmatter_str:
            try:
                import yaml

                metadata = yaml.safe_load(frontmatter_str) or {}
            except ImportError:
                # Fallback to basic parsing if PyYAML is not available
                metadata = self._parse_yaml_basic(frontmatter_str)
            except Exception:
                metadata = {}

        return BitBucketPromptTemplate(
            template_id=prompt_id,
            content=template_content,
            metadata=metadata,
        )

    def _parse_yaml_basic(self, yaml_str: str) -> Dict[str, Any]:
        """Basic YAML parser for simple cases when PyYAML is not available."""
        result: Dict[str, Any] = {}
        for line in yaml_str.split("\n"):
            line = line.strip()
            if ":" in line and not line.startswith("#"):
                key, value = line.split(":", 1)
                key = key.strip()
                value = value.strip()

                # Try to parse value as appropriate type
                if value.lower() in ["true", "false"]:
                    result[key] = value.lower() == "true"
                elif value.isdigit():
                    result[key] = int(value)
                elif value.replace(".", "").isdigit():
                    result[key] = float(value)
                else:
                    result[key] = value.strip("\"'")
        return result

    def render_template(
        self, template_id: str, variables: Optional[Dict[str, Any]] = None
    ) -> str:
        """Render a template with the given variables."""
        if template_id not in self.prompts:
            raise ValueError(f"Template '{template_id}' not found")

        template = self.prompts[template_id]
        jinja_template = self.jinja_env.from_string(template.content)

        return jinja_template.render(**(variables or {}))

    def get_template(self, template_id: str) -> Optional[BitBucketPromptTemplate]:
        """Get a template by ID."""
        return self.prompts.get(template_id)

    def list_templates(self) -> List[str]:
        """List all available template IDs."""
        return list(self.prompts.keys())


class BitBucketPromptManager(CustomPromptManagement):
    """
    BitBucket prompt manager that integrates with LiteLLM's prompt management system.

    This class enables using .prompt files from BitBucket repositories with the
    litellm completion() function by implementing the PromptManagementBase interface.

    Usage:
        # Configure BitBucket access
        bitbucket_config = {
            "workspace": "your-workspace",
            "repository": "your-repo",
            "access_token": "your-token",
            "branch": "main"  # optional, defaults to main
        }

        # Use with completion
        response = litellm.completion(
            model="bitbucket/gpt-4",
            prompt_id="my_prompt",
            prompt_variables={"variable": "value"},
            bitbucket_config=bitbucket_config,
            messages=[{"role": "user", "content": "This will be combined with the prompt"}]
        )
    """

    def __init__(
        self,
        bitbucket_config: Dict[str, Any],
        prompt_id: Optional[str] = None,
    ):
        self.bitbucket_config = bitbucket_config
        self.prompt_id = prompt_id
        self._prompt_manager: Optional[BitBucketTemplateManager] = None

    @property
    def integration_name(self) -> str:
        """Integration name used in model names like 'bitbucket/gpt-4'."""
        return "bitbucket"

    @property
    def prompt_manager(self) -> BitBucketTemplateManager:
        """Get or create the prompt manager instance."""
        if self._prompt_manager is None:
            self._prompt_manager = BitBucketTemplateManager(
                bitbucket_config=self.bitbucket_config,
                prompt_id=self.prompt_id,
            )
        return self._prompt_manager

    def get_prompt_template(
        self,
        prompt_id: str,
        prompt_variables: Optional[Dict[str, Any]] = None,
    ) -> Tuple[str, Dict[str, Any]]:
        """
        Get a prompt template and render it with variables.

        Args:
            prompt_id: The ID of the prompt template
            prompt_variables: Variables to substitute in the template

        Returns:
            Tuple of (rendered_prompt, metadata)
        """
        template = self.prompt_manager.get_template(prompt_id)
        if not template:
            raise ValueError(f"Prompt template '{prompt_id}' not found")

        # Render the template
        rendered_prompt = self.prompt_manager.render_template(
            prompt_id, prompt_variables or {}
        )

        # Extract metadata
        metadata = {
            "model": template.model,
            "temperature": template.temperature,
            "max_tokens": template.max_tokens,
            **template.optional_params,
        }

        return rendered_prompt, metadata

    def pre_call_hook(
        self,
        user_id: Optional[str],
        messages: List[AllMessageValues],
        function_call: Optional[Union[Dict[str, Any], str]] = None,
        litellm_params: Optional[Dict[str, Any]] = None,
        prompt_id: Optional[str] = None,
        prompt_variables: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]:
        """
        Pre-call hook that processes the prompt template before making the LLM call.
        """
        if not prompt_id:
            return messages, litellm_params

        try:
            # Get the rendered prompt and metadata
            rendered_prompt, prompt_metadata = self.get_prompt_template(
                prompt_id, prompt_variables
            )

            # Parse the rendered prompt into messages
            parsed_messages = self._parse_prompt_to_messages(rendered_prompt)

            # Merge with existing messages
            if parsed_messages:
                # If we have parsed messages, use them instead of the original messages
                final_messages: List[AllMessageValues] = parsed_messages
            else:
                # If no messages were parsed, prepend the prompt to existing messages
                final_messages = [
                    {"role": "user", "content": rendered_prompt}  # type: ignore
                ] + messages

            # Update litellm_params with prompt metadata
            if litellm_params is None:
                litellm_params = {}

            # Apply model and parameters from prompt metadata
            if prompt_metadata.get("model"):
                litellm_params["model"] = prompt_metadata["model"]

            for param in [
                "temperature",
                "max_tokens",
                "top_p",
                "frequency_penalty",
                "presence_penalty",
            ]:
                if param in prompt_metadata:
                    litellm_params[param] = prompt_metadata[param]

            return final_messages, litellm_params

        except Exception as e:
            # Log error but don't fail the call
            import litellm

            litellm._logging.verbose_proxy_logger.error(
                f"Error in BitBucket prompt pre_call_hook: {e}"
            )
            return messages, litellm_params

    def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]:
        """
        Parse prompt content into a list of messages.
        Handles both simple prompts and multi-role conversations.
        """
        messages = []
        lines = prompt_content.strip().split("\n")
        current_role = None
        current_content = []

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # Check for role indicators
            if line.lower().startswith("system:"):
                if current_role and current_content:
                    messages.append(
                        {
                            "role": current_role,
                            "content": "\n".join(current_content).strip(),
                        }  # type: ignore
                    )
                current_role = "system"
                current_content = [line[7:].strip()]  # Remove "System:" prefix
            elif line.lower().startswith("user:"):
                if current_role and current_content:
                    messages.append(
                        {
                            "role": current_role,
                            "content": "\n".join(current_content).strip(),
                        }  # type: ignore
                    )
                current_role = "user"
                current_content = [line[5:].strip()]  # Remove "User:" prefix
            elif line.lower().startswith("assistant:"):
                if current_role and current_content:
                    messages.append(
                        {
                            "role": current_role,
                            "content": "\n".join(current_content).strip(),
                        }  # type: ignore
                    )
                current_role = "assistant"
                current_content = [line[10:].strip()]  # Remove "Assistant:" prefix
            else:
                # Continue building current message
                current_content.append(line)

        # Add the last message
        if current_role and current_content:
            messages.append(
                {"role": current_role, "content": "\n".join(current_content).strip()}
            )

        # If no role indicators found, treat as a single user message
        if not messages and prompt_content.strip():
            messages = [{"role": "user", "content": prompt_content.strip()}]  # type: ignore

        return messages  # type: ignore

    def post_call_hook(
        self,
        user_id: Optional[str],
        response: Any,
        input_messages: List[AllMessageValues],
        function_call: Optional[Union[Dict[str, Any], str]] = None,
        litellm_params: Optional[Dict[str, Any]] = None,
        prompt_id: Optional[str] = None,
        prompt_variables: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Any:
        """
        Post-call hook for any post-processing after the LLM call.
        """
        return response

    def get_available_prompts(self) -> List[str]:
        """Get list of available prompt IDs."""
        return self.prompt_manager.list_templates()

    def reload_prompts(self) -> None:
        """Reload prompts from BitBucket."""
        if self.prompt_id:
            self._prompt_manager = None  # Reset to force reload
            self.prompt_manager  # This will trigger reload

    def should_run_prompt_management(
        self,
        prompt_id: Optional[str],
        prompt_spec: Optional[PromptSpec],
        dynamic_callback_params: StandardCallbackDynamicParams,
    ) -> bool:
        """
        Determine if prompt management should run based on the prompt_id.

        For BitBucket, we always return True and handle the prompt loading
        in the _compile_prompt_helper method.
        """
        return prompt_id is not None

    def _compile_prompt_helper(
        self,
        prompt_id: Optional[str],
        prompt_spec: Optional[PromptSpec],
        prompt_variables: Optional[dict],
        dynamic_callback_params: StandardCallbackDynamicParams,
        prompt_label: Optional[str] = None,
        prompt_version: Optional[int] = None,
    ) -> PromptManagementClient:
        """
        Compile a BitBucket prompt template into a PromptManagementClient structure.

        This method:
        1. Loads the prompt template from BitBucket
        2. Renders it with the provided variables
        3. Converts the rendered text into chat messages
        4. Extracts model and optional parameters from metadata
        """
        if prompt_id is None:
            raise ValueError("prompt_id is required for BitBucket prompt manager")

        try:
            # Load the prompt from BitBucket if not already loaded
            if prompt_id not in self.prompt_manager.prompts:
                self.prompt_manager._load_prompt_from_bitbucket(prompt_id)

            # Get the rendered prompt and metadata
            rendered_prompt, prompt_metadata = self.get_prompt_template(
                prompt_id, prompt_variables
            )

            # Convert rendered content to chat messages
            messages = self._parse_prompt_to_messages(rendered_prompt)

            # Extract model from metadata (if specified)
            template_model = prompt_metadata.get("model")

            # Extract optional parameters from metadata
            optional_params = {}
            for param in [
                "temperature",
                "max_tokens",
                "top_p",
                "frequency_penalty",
                "presence_penalty",
            ]:
                if param in prompt_metadata:
                    optional_params[param] = prompt_metadata[param]

            return PromptManagementClient(
                prompt_id=prompt_id,
                prompt_template=messages,
                prompt_template_model=template_model,
                prompt_template_optional_params=optional_params,
                completed_messages=None,
            )

        except Exception as e:
            raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")

    async def async_compile_prompt_helper(
        self,
        prompt_id: Optional[str],
        prompt_variables: Optional[dict],
        dynamic_callback_params: StandardCallbackDynamicParams,
        prompt_spec: Optional[PromptSpec] = None,
        prompt_label: Optional[str] = None,
        prompt_version: Optional[int] = None,
    ) -> PromptManagementClient:
        """
        Async version of compile prompt helper. Since BitBucket operations use sync client,
        this simply delegates to the sync version.
        """
        if prompt_id is None:
            raise ValueError("prompt_id is required for BitBucket prompt manager")

        return self._compile_prompt_helper(
            prompt_id=prompt_id,
            prompt_spec=prompt_spec,
            prompt_variables=prompt_variables,
            dynamic_callback_params=dynamic_callback_params,
            prompt_label=prompt_label,
            prompt_version=prompt_version,
        )

    def get_chat_completion_prompt(
        self,
        model: str,
        messages: List[AllMessageValues],
        non_default_params: dict,
        prompt_id: Optional[str],
        prompt_variables: Optional[dict],
        dynamic_callback_params: StandardCallbackDynamicParams,
        prompt_spec: Optional[PromptSpec] = None,
        prompt_label: Optional[str] = None,
        prompt_version: Optional[int] = None,
        ignore_prompt_manager_model: Optional[bool] = False,
        ignore_prompt_manager_optional_params: Optional[bool] = False,
    ) -> Tuple[str, List[AllMessageValues], dict]:
        """
        Get chat completion prompt from BitBucket and return processed model, messages, and parameters.
        """
        return PromptManagementBase.get_chat_completion_prompt(
            self,
            model,
            messages,
            non_default_params,
            prompt_id,
            prompt_variables,
            dynamic_callback_params,
            prompt_spec=prompt_spec,
            prompt_label=prompt_label,
            prompt_version=prompt_version,
        )

    async def async_get_chat_completion_prompt(
        self,
        model: str,
        messages: List[AllMessageValues],
        non_default_params: dict,
        prompt_id: Optional[str],
        prompt_variables: Optional[dict],
        dynamic_callback_params: StandardCallbackDynamicParams,
        litellm_logging_obj: LiteLLMLoggingObj,
        prompt_spec: Optional[PromptSpec] = None,
        tools: Optional[List[Dict]] = None,
        prompt_label: Optional[str] = None,
        prompt_version: Optional[int] = None,
        ignore_prompt_manager_model: Optional[bool] = False,
        ignore_prompt_manager_optional_params: Optional[bool] = False,
    ) -> Tuple[str, List[AllMessageValues], dict]:
        """
        Async version - delegates to PromptManagementBase async implementation.
        """
        return await PromptManagementBase.async_get_chat_completion_prompt(
            self,
            model,
            messages,
            non_default_params,
            prompt_id=prompt_id,
            prompt_variables=prompt_variables,
            litellm_logging_obj=litellm_logging_obj,
            dynamic_callback_params=dynamic_callback_params,
            prompt_spec=prompt_spec,
            tools=tools,
            prompt_label=prompt_label,
            prompt_version=prompt_version,
            ignore_prompt_manager_model=ignore_prompt_manager_model,
            ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
        )
