# langxchange/prompt_helper.py

from typing import Any, Dict, List, Optional, Callable, Union
from enum import Enum
import json
from .tracing_helper import trace_span


class PromptMode(Enum):
    """Enumeration for different prompt building modes."""
    BASIC = "basic"
    AUGMENTED = "augmented"
    CONTEXTUAL = "contextual"
    SUMMARIZED = "summarized"


class ResponseValidator:
    """
    Validates LLM responses against user intent, system prompt, and conversation history.
    Used by the prompt_pass feature to ensure response quality.
    """

    VALIDATION_PROMPT_TEMPLATE = """You are a response quality validator. Analyze if the given response appropriately addresses the user's request.

## Original System Context:
{system_prompt}

## Conversation History:
{history}

## User's Original Query:
{user_query}

## Generated Response:
{response}

## Your Task:
1. Evaluate if the response accurately addresses the user's query
2. Check if the response aligns with the system prompt guidelines
3. Consider the conversation history for context continuity
4. If the response is inadequate, provide an improved version

Respond in the following JSON format:
{{
    "is_valid": true/false,
    "issues": ["list of issues if any"],
    "refined_response": "improved response if is_valid is false, otherwise null"
}}
"""

    @staticmethod
    def build_validation_prompt(
        system_prompt: str,
        user_query: str,
        response: str,
        history: Optional[List[Dict[str, str]]] = None
    ) -> List[Dict[str, str]]:
        """Build a validation prompt for checking response quality."""
        history_str = ""
        if history:
            history_str = "\n".join([
                f"{msg['role'].upper()}: {msg['content']}"
                for msg in history
            ])
        else:
            history_str = "(No prior history)"

        validation_content = ResponseValidator.VALIDATION_PROMPT_TEMPLATE.format(
            system_prompt=system_prompt,
            history=history_str,
            user_query=user_query,
            response=response
        )

        return [
            {"role": "system", "content": "You are a precise response quality validator. Always respond with valid JSON."},
            {"role": "user", "content": validation_content}
        ]

    @staticmethod
    def parse_validation_result(llm_output: str) -> Dict[str, Any]:
        """Parse the validation result from LLM output."""
        try:
            # Try to extract JSON from the response
            start_idx = llm_output.find('{')
            end_idx = llm_output.rfind('}') + 1
            if start_idx != -1 and end_idx > start_idx:
                json_str = llm_output[start_idx:end_idx]
                result = json.loads(json_str)
                return {
                    "is_valid": result.get("is_valid", True),
                    "issues": result.get("issues", []),
                    "refined_response": result.get("refined_response")
                }
        except json.JSONDecodeError:
            pass

        # Default: assume valid if parsing fails
        return {"is_valid": True, "issues": [], "refined_response": None}


class PromptBuilder:
    """
    Handles different strategies for building prompts from retrieval results.
    """

    @staticmethod
    def basic_prompt(
        system_prompt: str,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        history: Optional[List[Dict[str, str]]] = None
    ) -> List[Dict[str, str]]:
        """Build a basic prompt with history."""
        messages = [{"role": "system", "content": system_prompt}]
        if history:
            messages.extend(history)
        messages.append({"role": "user", "content": user_query})
        return messages

    @staticmethod
    def augmented_prompt(
        system_prompt: str,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        max_context_length: int = 2000,
        history: Optional[List[Dict[str, str]]] = None
    ) -> List[Dict[str, str]]:
        """Build an augmented prompt with retrieval results as context."""
        messages = [{"role": "system", "content": system_prompt}]

        if history:
            messages.extend(history)

        if retrieval_results:
            context_parts = []
            current_length = 0

            for hit in retrieval_results:
                # Support both object attributes and dictionary keys
                if isinstance(hit, dict):
                    doc = hit.get("document") or hit.get("content") or hit.get("text") or ""
                else:
                    doc = getattr(hit, "document", None) or getattr(hit, "content", None) or getattr(hit, "text", "")

                if isinstance(doc, list):
                    doc = "\n".join(doc)

                context_part = f"{doc}"
                
                # Check if adding this part exceeds max_context_length (simple char count as proxy)
                if current_length + len(context_part) > max_context_length:
                    # If it's the first doc and it's already too big, truncate it
                    if not context_parts:
                        doc_truncated = context_part[:max_context_length] + "... [truncated]"
                        context_parts.append(doc_truncated)
                    break

                context_parts.append(context_part)
                current_length += len(context_part)

            if context_parts:
                context_content = "\n\n---\n\n".join(context_parts)
                augmented_query = f"""Below is some relevant context information that may help you address the query. Please use it if applicable, while keeping the conversation flow natural.

Relevant Context:
{context_content}

---

User Query: {user_query}"""

                messages.append({"role": "user", "content": augmented_query})
            else:
                messages.append({"role": "user", "content": user_query})
        else:
            messages.append({"role": "user", "content": user_query})

        return messages

    @staticmethod
    def contextual_prompt(
        system_prompt: str,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        history: Optional[List[Dict[str, str]]] = None
    ) -> List[Dict[str, str]]:
        """Build a contextual prompt with retrieval results as assistant messages."""
        messages = [{"role": "system", "content": system_prompt}]

        if history:
            messages.extend(history)

        if retrieval_results:
            current_length = 0
            for hit in retrieval_results:
                doc = hit.get("text") or hit.get("document", "")
                if isinstance(doc, list):
                    doc = "\n".join(doc)
                
                if current_length + len(doc) > 4000: # Default safe limit for contextual messages
                    break

                meta = hit.get("metadata", {})
                tag = ", ".join(f"{k}={v}" for k, v in meta.items())

                messages.append({
                    "role": "assistant",
                    "content": f"[Reference] {tag}:\n{doc}"
                })
                current_length += len(doc)

        messages.append({"role": "user", "content": user_query})
        return messages

    @staticmethod
    def summarized_prompt(
        system_prompt: str,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        max_snippets: int = 3,
        history: Optional[List[Dict[str, str]]] = None
    ) -> List[Dict[str, str]]:
        """Build a prompt with summarized retrieval results."""
        messages = [{"role": "system", "content": system_prompt}]

        if history:
            messages.extend(history)

        if retrieval_results:
            # Take top snippets and create a concise summary
            top_results = retrieval_results[:max_snippets]
            summaries = []

            for i, hit in enumerate(top_results, 1):
                doc = hit.get("text") or hit.get("document", "")
                if isinstance(doc, list):
                    doc = "\n".join(doc)

                # Truncate long documents
                if len(doc) > 300:
                    doc = doc[:300] + "..."

                meta = hit.get("metadata", {})
                source = meta.get("source", f"Document {i}")

                summaries.append(f"{i}. {source}: {doc}")

            if summaries:
                summary_content = "\n".join(summaries)
                enhanced_query = f"""Relevant information:
{summary_content}

Question: {user_query}

Please answer based on the information provided above."""

                messages.append({"role": "user", "content": enhanced_query})
            else:
                messages.append({"role": "user", "content": user_query})
        else:
            messages.append({"role": "user", "content": user_query})

        return messages


class EnhancedPromptHelper:
    """
    Builds prompts for LLMs and provides an interface to call injected tools.
    Supports multiple prompt building strategies for efficient retrieval augmentation.
    Includes prompt_pass feature for response validation and refinement.
    """

    def __init__(
        self,
        llm: Any,
        system_prompt: str,
        tools: Optional[Dict[str, Any]] = None,
        default_mode: PromptMode = PromptMode.BASIC,
        max_context_length: int = 2000,
        max_snippets: int = 3
    ):
        """
        Initialize the EnhancedPromptHelper.

        :param llm: an object with .chat(messages, **kwargs) -> str
        :param system_prompt: the initial system-level instruction
        :param tools: a dict mapping tool names to tool instances
        :param default_mode: default prompt building mode
        :param max_context_length: maximum context length for augmented mode
        :param max_snippets: maximum number of snippets for summarized mode
        """
        self.llm = llm
        self.system_prompt = system_prompt
        self.tools = tools or {}
        self.default_mode = default_mode
        self.max_context_length = max_context_length
        self.max_snippets = max_snippets

        # Map modes to builder functions
        self._builders = {
            PromptMode.BASIC: PromptBuilder.basic_prompt,
            PromptMode.AUGMENTED: PromptBuilder.augmented_prompt,
            PromptMode.CONTEXTUAL: PromptBuilder.contextual_prompt,
            PromptMode.SUMMARIZED: PromptBuilder.summarized_prompt
        }

    def call_tool(self, tool_name: str, *args, **kwargs) -> Any:
        """
        Call a registered tool by name.

        :param tool_name: name of the tool to call
        :param args: positional arguments for the tool
        :param kwargs: keyword arguments for the tool
        :return: result from the tool
        :raises ValueError: if tool is not found
        """
        if tool_name not in self.tools:
            raise ValueError(f"Tool '{tool_name}' not found. Available tools: {list(self.tools.keys())}")

        tool = self.tools[tool_name]

        # Check if the tool has a 'run' method first (prioritize run method)
        if hasattr(tool, 'run') and callable(tool.run):
            return tool.run(*args, **kwargs)
        elif callable(tool):
            return tool(*args, **kwargs)
        else:
            raise ValueError(f"Tool '{tool_name}' is not callable and has no 'run' method")

    def register_tool(self, name: str, tool: Any) -> None:
        """Register a new tool."""
        self.tools[name] = tool

    def list_tools(self) -> List[str]:
        """List available tool names."""
        return list(self.tools.keys())

    def build_prompt(
        self,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        mode: Optional[PromptMode] = None,
        custom_system_prompt: Optional[str] = None,
        history: Optional[List[Dict[str, str]]] = None
    ) -> List[Dict[str, str]]:
        """
        Build a prompt using the specified mode.

        :param user_query: the user's query
        :param retrieval_results: optional retrieval results to include
        :param mode: prompt building mode (uses default if None)
        :param custom_system_prompt: override system prompt for this request
        :param history: optional conversation history to include
        :return: list of message dictionaries
        """
        mode = mode or self.default_mode
        system_prompt = custom_system_prompt or self.system_prompt

        builder = self._builders.get(mode)
        if not builder:
            raise ValueError(f"Unknown prompt mode: {mode}")

        # Pass additional parameters based on mode
        if mode == PromptMode.AUGMENTED:
            return builder(
                system_prompt,
                user_query,
                retrieval_results,
                max_context_length=self.max_context_length,
                history=history
            )
        elif mode == PromptMode.SUMMARIZED:
            return builder(
                system_prompt,
                user_query,
                retrieval_results,
                max_snippets=self.max_snippets,
                history=history
            )
        else:
            return builder(system_prompt, user_query, retrieval_results, history=history)

    def _validate_and_refine(
        self,
        user_query: str,
        response: str,
        history: Optional[List[Dict[str, str]]] = None,
        custom_system_prompt: Optional[str] = None,
        validation_criteria: Optional[Callable[[str], bool]] = None,
        temperature: float = 0.3,
        max_tokens: int = 1024
    ) -> Dict[str, Any]:
        """
        Validate and optionally refine an LLM response.

        :param user_query: the original user query
        :param response: the LLM response to validate
        :param history: conversation history for context
        :param custom_system_prompt: override system prompt
        :param validation_criteria: optional custom validation function
        :param temperature: temperature for validation LLM call
        :param max_tokens: max tokens for validation response
        :return: dict with is_valid, issues, and refined_response
        """
        system_prompt = custom_system_prompt or self.system_prompt

        # If custom validation criteria provided, use it first
        if validation_criteria is not None:
            try:
                is_valid = validation_criteria(response)
                if is_valid:
                    return {"is_valid": True, "issues": [], "refined_response": None}
            except Exception:
                pass  # Fall through to LLM-based validation

        # Build validation prompt
        validation_messages = ResponseValidator.build_validation_prompt(
            system_prompt=system_prompt,
            user_query=user_query,
            response=response,
            history=history
        )

        # Call LLM for validation
        validation_response = self.llm.chat(
            messages=validation_messages,
            temperature=temperature,
            max_tokens=max_tokens
        )

        # Extract content if needed
        if hasattr(validation_response, 'content') and not isinstance(validation_response, (str, dict)):
            validation_response = str(validation_response.content) or ""

        # Parse and return result
        return ResponseValidator.parse_validation_result(validation_response)

    def run(
        self,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        temperature: float = 0.7,
        max_tokens: int = 512,
        mode: Optional[PromptMode] = None,
        custom_system_prompt: Optional[str] = None,
        history: Optional[List[Dict[str, str]]] = None,
        prompt_pass: int = 1,
        validation_criteria: Optional[Callable[[str], bool]] = None,
        **llm_kwargs
    ) -> Union[str, Dict[str, Any]]:
        """
        Run the LLM with the built prompt, with optional multi-pass validation.

        :param user_query: the user's query
        :param retrieval_results: optional retrieval results to include
        :param temperature: sampling temperature for the LLM
        :param max_tokens: maximum tokens to generate
        :param mode: prompt building mode (uses default if None)
        :param custom_system_prompt: override system prompt for this request
        :param history: optional conversation history to include
        :param prompt_pass: number of validation passes (1 = no validation, >1 = validate and refine)
        :param validation_criteria: optional custom validation function (str) -> bool
        :param llm_kwargs: additional keyword arguments for the LLM
        :return: LLM response string, or dict with response and pass_info if prompt_pass > 1
        """
        if prompt_pass < 1:
            raise ValueError("prompt_pass must be at least 1")

        try:
            messages = self.build_prompt(
                user_query=user_query,
                retrieval_results=retrieval_results,
                mode=mode,
                custom_system_prompt=custom_system_prompt,
                history=history
            )

            # Merge default parameters with custom ones
            chat_params = {
                'messages': messages,
                'temperature': temperature,
                'max_tokens': max_tokens,
                **llm_kwargs
            }

            response = self.llm.chat(**chat_params)

            # Extract content if it's an object (like DeepSeek's APIResponse)
            if hasattr(response, 'content') and not isinstance(response, (str, dict)):
                response = str(response.content) or ""

            # If only 1 pass, return immediately (backward compatible)
            if prompt_pass == 1:
                return response or ""

            # Multi-pass validation and refinement
            pass_info = {
                "total_passes": 1,
                "validation_history": []
            }

            current_response = response

            for pass_num in range(2, prompt_pass + 1):
                validation_result = self._validate_and_refine(
                    user_query=user_query,
                    response=current_response,
                    history=history,
                    custom_system_prompt=custom_system_prompt,
                    validation_criteria=validation_criteria
                )

                pass_info["validation_history"].append({
                    "pass": pass_num,
                    "is_valid": validation_result["is_valid"],
                    "issues": validation_result["issues"]
                })
                pass_info["total_passes"] = pass_num

                # If valid, stop early
                if validation_result["is_valid"]:
                    break

                # If refined response provided, use it
                if validation_result["refined_response"]:
                    current_response = validation_result["refined_response"]
                else:
                    # Re-generate with slightly adjusted prompt
                    retry_messages = messages.copy()
                    issues_str = ", ".join(validation_result["issues"]) if validation_result["issues"] else "unclear"
                    retry_messages.append({
                        "role": "assistant",
                        "content": current_response
                    })
                    retry_messages.append({
                        "role": "user",
                        "content": f"The previous response had the following issues: {issues_str}. Please provide an improved response that better addresses my original query."
                    })

                    retry_response = self.llm.chat(
                        messages=retry_messages,
                        temperature=max(0.3, temperature - 0.1),  # Lower temp for refinement
                        max_tokens=max_tokens,
                        **llm_kwargs
                    )

                    if hasattr(retry_response, 'content') and not isinstance(retry_response, (str, dict)):
                        retry_response = str(retry_response.content) or ""

                    current_response = retry_response or current_response

            return {
                "response": current_response,
                "pass_info": pass_info
            }

        except Exception as e:
            raise RuntimeError(f"Error during LLM execution: {str(e)}") from e

    def run_with_tools(
        self,
        user_query: str,
        tool_calls: Optional[List[Dict[str, Any]]] = None,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        prompt_pass: int = 1,
        validation_criteria: Optional[Callable[[str], bool]] = None,
        **run_kwargs
    ) -> Dict[str, Any]:
        """
        Run the LLM with optional tool calls and multi-pass validation.

        :param user_query: the user's query
        :param tool_calls: list of tool calls to execute before LLM
        :param retrieval_results: optional retrieval results to include
        :param prompt_pass: number of validation passes
        :param validation_criteria: optional custom validation function
        :param run_kwargs: arguments for the run method
        :return: dictionary with LLM response, tool results, and pass_info
        """
        tool_results = {}

        # Execute tool calls if provided
        if tool_calls:
            for tool_call in tool_calls:
                tool_name = tool_call.get('name')
                tool_args = tool_call.get('args', [])
                tool_kwargs = tool_call.get('kwargs', {})

                if tool_name:
                    try:
                        result = self.call_tool(tool_name, *tool_args, **tool_kwargs)
                        tool_results[tool_name] = result
                    except Exception as e:
                        tool_results[tool_name] = f"Error: {str(e)}"

        # Run LLM with prompt_pass support
        llm_result = self.run(
            user_query=user_query,
            retrieval_results=retrieval_results,
            prompt_pass=prompt_pass,
            validation_criteria=validation_criteria,
            **run_kwargs
        )

        # Handle both single-pass (str) and multi-pass (dict) returns
        if isinstance(llm_result, dict):
            return {
                'llm_response': llm_result["response"],
                'tool_results': tool_results,
                'pass_info': llm_result.get("pass_info")
            }
        else:
            return {
                'llm_response': llm_result,
                'tool_results': tool_results,
                'pass_info': None
            }

    def update_config(
        self,
        default_mode: Optional[PromptMode] = None,
        max_context_length: Optional[int] = None,
        max_snippets: Optional[int] = None,
        system_prompt: Optional[str] = None
    ) -> None:
        """Update configuration parameters."""
        if default_mode is not None:
            self.default_mode = default_mode
        if max_context_length is not None:
            self.max_context_length = max_context_length
        if max_snippets is not None:
            self.max_snippets = max_snippets
        if system_prompt is not None:
            self.system_prompt = system_prompt

    def get_config(self) -> Dict[str, Any]:
        """Get current configuration."""
        return {
            'default_mode': self.default_mode.value,
            'max_context_length': self.max_context_length,
            'max_snippets': self.max_snippets,
            'available_tools': list(self.tools.keys())
        }


class PromptHelper:
    """
    Builds prompts for LLMs and provides an interface to call injected tools.
    """

    def __init__(
        self,
        llm: Any,
        system_prompt: str,
        tools: Optional[Dict[str, Any]] = None
    ):
        """
        :param llm: an object with .chat(messages, **kwargs) -> str
        :param system_prompt: the initial system-level instruction
        :param tools: a dict mapping tool names to tool instances
        """
        self.llm = llm
        self.system_prompt = system_prompt


    @trace_span("PromptHelper.run")
    def run(
        self,
        user_query: str,
        retrieval_results: Optional[List[Dict[str, Any]]] = None,
        temperature: float = 0.7,
        max_tokens: int = 512
    ) -> str:

        # otherwise build a normal chat
        messages: List[Dict[str, str]] = [
            {"role": "system", "content": self.system_prompt}
        ]

        if retrieval_results:
            # include retrieved snippets as assistant messages
            for hit in retrieval_results:
                doc = hit.get("text") or hit.get("document", "")
                if isinstance(doc, list):
                    doc = "\n".join(doc)
                meta = hit.get("metadata", {})
                tag = ", ".join(f"{k}={v}" for k, v in meta.items())
                messages.append({
                    "role": "assistant",
                    "content": f"[Snippet] {tag}:\n{doc}"
                })

        messages.append({"role": "user", "content": user_query})

        return self.llm.chat(
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens
        )
