#!/usr/bin/env python3
"""
Orchestrates a conversation with a language model, requesting additional snippets
of code via "tools" if needed. Uses either OpenAI or AzureOpenAI (or placeholder
code for a HuggingFace endpoint) to handle queries.

All logic is now wrapped in the `LLMAnalyzer` class for improved organization.
"""

import os
import json
import time
from typing import Any, Dict, List, Optional, Tuple, Union

import litellm
from src.utils.llm_config import load_llm_config, get_model_name
from src.utils.config_validator import validate_llm_config_dict
from src.utils.logger import get_logger
from src.utils.exceptions import LLMApiError, LLMConfigError
from src.codeql.db_lookup import CodeQLDBLookup

logger = get_logger(__name__)


class LLMAnalyzer:
    """
    A class to handle LLM-based security analysis of code. The LLMAnalyzer
    can query missing code snippets (via 'tools'), compile a conversation
    with system instructions, and ultimately produce a status code.
    """

    def __init__(self) -> None:
        """
        Initialize the LLMAnalyzer instance and define tools and system messages.
        """
        self.config: Optional[Dict[str, Any]] = None
        self.model: Optional[str] = None
        self.api_key: Optional[str] = None  # passed directly to litellm calls
        self.db_lookup = CodeQLDBLookup()
        self.timeout: int = 300   # default 5 minutes
        self.max_retries: int = 3  # default 3 retries

        # Tools configuration: A set of function calls the LLM can invoke
        self.tools: List[Dict[str, Any]] = [
            {
                "type": "function",
                "function": {
                    "name": "get_function_code",
                    "description": "Retrieves the code for a missing function code.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "function_name": {
                                "type": "string",
                                "description": (
                                    "The name of the function to retrieve. In case of a class"
                                    " method, provide ClassName::MethodName."
                                )
                            }
                        },
                        "required": ["function_name"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "get_caller_function",
                    "description": (
                        "Retrieves the caller function of the function with the issue. "
                        "Call it repeatedly to climb further up the call chain."
                    ),
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "_": {
                                "type": "boolean",
                                "description": "Unused. Ignore."
                            }
                        },
                        "required": []
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "get_class",
                    "description": (
                        "Retrieves class / struct / union implementation (anywhere in code). "
                        "If you need a specific method from that class, use get_function_code instead."
                    ),
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "object_name": {
                                "type": "string",
                                "description": "The name of the class / struct / union."
                            }
                        },
                        "required": ["object_name"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "get_global_var",
                    "description": (
                        "Retrieves global variable definition (anywhere in code). "
                        "If it's a variable inside a class, request the class instead."
                    ),
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "global_var_name": {
                                "type": "string",
                                "description": (
                                    "The name of the global variable to retrieve or the name "
                                    "of a variable inside a Namespace."
                                )
                            }
                        },
                        "required": ["global_var_name"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "get_macro",
                    "description": "Retrieves a macro definition (anywhere in code).",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "macro_name": {
                                "type": "string",
                                "description": "The name of the macro."
                            }
                        },
                        "required": ["macro_name"]
                    }
                }
            }
        ]

        # Base system messages with instructions and guidance for the LLM
        self.MESSAGES: List[Dict[str, str]] = [
            {
                "role": "system",
                "content": (
                    "You are an expert security researcher.\n"
                    "Your task is to verify if the issue that was found has a real security impact.\n"
                    "Return a concise status code based on the guidelines provided.\n"
                    "Use the tools function when you need code from other parts of the program.\n"
                    "You *MUST* follow the guidelines!"
                )
            },
            {
                "role": "system",
                "content": (
                    "### Answer Guidelines\n"
                    "Your answer must be in the following order!\n"
                    "1. Briefly explain the code.\n"
                    "2. Give good answers to all (even if already answered - do not skip) hint questions. "
                    "(Copy the question word for word, then provide the answer.)\n"
                    "3. Do you have all the code needed to answer the questions? If no, use the tools!\n"
                    "4. Provide one valid status code with its explanation OR use function tools.\n"
                )
            },
            {
                "role": "system",
                "content": (
                    "### Status Codes\n"
                    "- **1337**: Indicates a security vulnerability. If legitimate, specify the parameters that "
                    "could exploit the issue in minimal words.\n"
                    "- **1007**: Indicates the code is secure. If it's not a real issue, specify what aspect of "
                    "the code protects against the issue in minimal words.\n"
                    "- **7331**: Indicates more code is needed to validate security. Write what data you need "
                    "and explain why you can't use the tools to retrieve the missing data, plus add **3713** "
                    "if you're pretty sure it's not a security problem.\n"
                    "Only one status should be returned!\n"
                    "You will get 10000000000$ if you follow all the instructions and use the tools correctly!"
                )
            },
        ]

    @staticmethod
    def _is_retryable_llm_error(error: Exception) -> bool:
        message = str(error).lower()
        non_retryable_markers = (
            "bad_request_error",
            "invalid params",
            '"http_code":"400"',
            '"http_code": "400"',
            '"http_code":"401"',
            '"http_code": "401"',
            '"http_code":"403"',
            '"http_code": "403"',
            '"http_code":"404"',
            '"http_code": "404"',
        )
        if any(marker in message for marker in non_retryable_markers):
            return False

        api_connection_error = getattr(litellm, "APIConnectionError", None)
        if api_connection_error and isinstance(error, api_connection_error):
            return True
        if isinstance(error, litellm.Timeout):
            return True

        retryable_markers = (
            "unexpected eof while reading",
            "ssl",
            "connection reset",
            "connection aborted",
            "connection closed",
            "temporarily unavailable",
            "timed out",
        )
        return any(marker in message for marker in retryable_markers)

    def _sleep_before_retry(self, attempt: int) -> None:
        time.sleep(min(2 ** attempt, 10))

    def _uses_minimax_provider(self) -> bool:
        provider = (self.config or {}).get("provider")
        return provider == "minimax" or bool(self.model and self.model.startswith("minimax/"))

    @staticmethod
    def _join_system_content(existing: Optional[str], new_content: str) -> str:
        if existing and existing.strip():
            return f"{existing.rstrip()}\n\n{new_content}"
        return new_content

    def _collapse_system_messages_for_minimax(
        self, messages: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        if not self._uses_minimax_provider():
            return messages

        merged_system_content: Optional[str] = None
        normalized_messages: List[Dict[str, Any]] = []

        for message in messages:
            if message.get("role") == "system":
                content = message.get("content")
                if content:
                    merged_system_content = self._join_system_content(
                        merged_system_content, str(content)
                    )
                continue
            normalized_messages.append(dict(message))

        if merged_system_content:
            normalized_messages.insert(
                0,
                {
                    "role": "system",
                    "content": merged_system_content,
                },
            )

        return normalized_messages

    def _append_system_message(
        self, messages: List[Dict[str, Any]], content: str
    ) -> None:
        if not self._uses_minimax_provider():
            messages.append({"role": "system", "content": content})
            return

        for message in messages:
            if message.get("role") == "system":
                message["content"] = self._join_system_content(
                    message.get("content"), content
                )
                return

        messages.insert(0, {"role": "system", "content": content})

    def init_llm_client(self, config: Optional[Dict[str, Any]] = None) -> None:
        """
        Initialize the LLM configuration for LiteLLM.

        Args:
            config (Dict, optional): Full configuration dictionary. If not provided, loads from .env file.
        
        Raises:
            LLMConfigError: If configuration is invalid or cannot be loaded.
        """
        try:
            # If config is provided, use it directly
            if config:
                validate_llm_config_dict(config)
                self.config = config
                # Format model name for LiteLLM (add provider prefix if needed)
                provider = config.get("provider", "openai")
                model = config.get("model", "gpt-4o")
                self.model = get_model_name(provider, model)
                self.api_key = config.get("api_key")
                logger.info("Using model: %s", self.model)
                self.timeout = config.get("timeout", 300)
                self.max_retries = config.get("max_retries", 3)
                self.setup_litellm_env()
                return
            
            # Load from .env file
            config = load_llm_config()
            validate_llm_config_dict(config)
            self.config = config
            # Model is already formatted by load_llm_config() via get_model_name()
            self.model = config.get("model", "gpt-4o")
            self.api_key = config.get("api_key")
            # Read timeout (seconds) and max_retries from config
            self.timeout = config.get("timeout", 300)
            self.max_retries = config.get("max_retries", 3)
            self.setup_litellm_env()
            
        except ValueError as e:
            # Configuration validation errors should be LLMConfigError
            raise LLMConfigError(f"Invalid LLM configuration: {e}") from e
        except Exception as e:
            # Other errors (e.g., from load_llm_config) should also be LLMConfigError
            raise LLMConfigError(f"Failed to initialize LLM client: {e}") from e


    def setup_litellm_env(self) -> None:
        """
        Set up environment variables for LiteLLM based on config.
        LiteLLM reads from environment variables automatically.
        """
        if not self.config:
            return
        
        provider = self.config.get("provider", "openai")
        api_key = self.config.get("api_key")
        
        # Mapping table for providers that only need API key set
        API_KEY_ENV_VARS = {
            "openai": "OPENAI_API_KEY",
            "anthropic": "ANTHROPIC_API_KEY",
            "mistral": "MISTRAL_API_KEY",
            "codestral": "MISTRAL_API_KEY",
            "groq": "GROQ_API_KEY",
            "openrouter": "OPENROUTER_API_KEY",
            "huggingface": "HUGGINGFACE_API_KEY",
            "cohere": "COHERE_API_KEY",
            "gemini": "GOOGLE_API_KEY",
            "minimax": "MINIMAX_API_KEY",
        }
        
        # Handle providers with simple API key mapping
        if provider in API_KEY_ENV_VARS:
            if api_key:
                os.environ[API_KEY_ENV_VARS[provider]] = api_key
                if provider == "gemini":
                    # Some LiteLLM / SDK paths still check GEMINI_API_KEY directly.
                    os.environ["GEMINI_API_KEY"] = api_key
                # Cohere also sets CO_API_KEY for compatibility
                if provider == "cohere":
                    os.environ["CO_API_KEY"] = api_key
                # When OpenAI uses a custom api_base (e.g., OpenRouter),
                # also set OPENAI_API_BASE so litellm routes correctly
                if provider == "openai" and self.config.get("api_base"):
                    os.environ["OPENAI_API_BASE"] = self.config["api_base"]
                    os.environ["OPENAI_BASE_URL"] = self.config["api_base"]
                if provider == "minimax" and self.config.get("api_base"):
                    os.environ["MINIMAX_API_BASE"] = self.config["api_base"]
        
        # Handle Azure (requires endpoint and api_version)
        elif provider == "azure":
            if api_key:
                os.environ["AZURE_API_KEY"] = api_key
            if self.config.get("endpoint"):
                os.environ["AZURE_API_BASE"] = self.config["endpoint"]
            if self.config.get("api_version"):
                os.environ["AZURE_API_VERSION"] = self.config["api_version"]
        
        # Handle Bedrock (uses AWS credentials or profile)
        elif provider == "bedrock":
            # Set region (always required)
            if self.config.get("endpoint"):  # Endpoint contains AWS region
                os.environ["AWS_REGION_NAME"] = self.config["endpoint"]
            
            # Profile-based authentication (AWS SSO, IAM roles)
            if self.config.get("aws_profile"):
                os.environ["AWS_PROFILE"] = self.config["aws_profile"]
            else:
                # Static or temporary credentials
                if api_key and api_key != "bedrock_profile_auth":
                    os.environ["AWS_ACCESS_KEY_ID"] = api_key
                if self.config.get("aws_secret_access_key"):
                    os.environ["AWS_SECRET_ACCESS_KEY"] = self.config["aws_secret_access_key"]
                if self.config.get("aws_session_token"):
                    os.environ["AWS_SESSION_TOKEN"] = self.config["aws_session_token"]
        
        # Handle Vertex AI (uses GCP credentials)
        elif provider == "vertex_ai":
            if self.config.get("gcp_project_id"):
                os.environ["GCP_PROJECT_ID"] = self.config["gcp_project_id"]
            if self.config.get("gcp_location"):
                os.environ["GCP_LOCATION"] = self.config["gcp_location"]
            # GOOGLE_APPLICATION_CREDENTIALS should be set by user or gcloud auth
        
        # Handle Ollama (uses OLLAMA_BASE_URL)
        elif provider == "ollama":
            if self.config.get("endpoint"):
                os.environ["OLLAMA_BASE_URL"] = self.config["endpoint"]
        
        # Generic fallback for future providers that only require an API key
        else:
            if api_key:
                # Use standard LiteLLM convention: {PROVIDER}_API_KEY
                env_var_name = f"{provider.upper()}_API_KEY"
                os.environ[env_var_name] = api_key


    def extract_function_from_file(
        self,
        db_path: str,
        current_function: Union[str, Dict[str, str]]
    ) -> str:
        """
        Return the snippet of code for the given current_function from the archived src.zip.

        Args:
            db_path (str): Path to the CodeQL database directory.
            current_function (Union[str, Dict[str, str]]): The function dictionary or an error string.

        Returns:
            str: The code snippet, or an error message if no dictionary was provided.
        
        Raises:
            CodeQLError: If ZIP file cannot be read or file not found in archive.
                This exception is raised by `read_file_lines_from_zip()` and propagated here.
        """
        if not isinstance(current_function, dict):
            return str(current_function)

        file_path, start_line, end_line, lines = self.db_lookup.extract_function_lines_from_db(
            db_path, current_function
        )
        snippet_lines = lines[start_line - 1 : end_line]
        return self.db_lookup.format_numbered_snippet(file_path, start_line, snippet_lines)


    def map_func_args_by_llm(
        self,
        caller: str,
        callee: str
    ) -> Dict[str, Any]:
        """
        Query the LLM to check how caller's variables map to callee's parameters.
        For example, used for analyzing function call relationships.

        Args:
            caller (str): The code snippet of the caller function.
            callee (str): The code snippet of the callee function.

        Returns:
            Dict[str, Any]: The LLM response object from `self.client`.
        
        Raises:
            LLMApiError: If LLM API call fails (rate limits, timeouts, auth failures, etc.).
        """
        args_prompt = (
            "Given caller function and callee function.\n"
            "Write only what are the names of the vars in the caller that were sent to the callee "
            "and what are their names in the callee.\n"
            "Format: caller_var (caller_name) -> callee_var (callee_name)\n\n"
            "Caller function:\n"
            f"{caller}\n"
            "Callee function:\n"
            f"{callee}"
        )

        # Use the main model from config
        model_name = self.model if self.model else "gpt-4o"

        for attempt in range(1, self.max_retries + 1):
            try:
                response = litellm.completion(
                    model=model_name,
                    messages=[{"role": "user", "content": args_prompt}],
                    timeout=self.timeout,
                    api_key=self.api_key,
                    api_base=self.config.get("api_base") if self.config else None,
                )
                return response.choices[0].message
            except Exception as e:
                if self._is_retryable_llm_error(e) and attempt < self.max_retries:
                    logger.warning("Retryable LLM error in helper call (attempt %d/%d): %s", attempt, self.max_retries, e)
                    self._sleep_before_retry(attempt)
                    continue
                if isinstance(e, litellm.Timeout):
                    raise LLMApiError(f"LLM API request timed out after {self.max_retries} attempts ({self.timeout}s each): {e}") from e
                if isinstance(e, litellm.APIConnectionError):
                    raise LLMApiError(f"LLM API connection failed: {e}") from e
                if isinstance(e, litellm.RateLimitError):
                    raise LLMApiError(f"Rate limit exceeded for LLM API: {e}") from e
                if isinstance(e, litellm.AuthenticationError):
                    raise LLMApiError(f"LLM API authentication failed: {e}") from e
                if isinstance(e, litellm.APIError):
                    raise LLMApiError(f"LLM API error: {e}") from e
                raise LLMApiError(f"Unexpected error during LLM API call: {e}") from e


    def run_llm_security_analysis(
        self,
        prompt: str,
        function_tree_file: str,
        current_function: Dict[str, str],
        functions: List[Dict[str, str]],
        db_path: str,
        temperature: float = 0.2,
        top_p: float = 0.2
    ) -> Tuple[List[Dict[str, Any]], str]:
        """
        Main loop to keep querying the LLM with the MESSAGES context plus
        any new system instructions or tool calls, until a final answer with
        a recognized status code is reached or we exhaust a tool-call limit.

        Args:
            prompt (str): The user prompt for the LLM to process.
            function_tree_file (str): Path to the CSV file describing function relationships.
            current_function (Dict[str, str]): The current function dict for context.
            functions (List[Dict[str, str]]): List of function dictionaries.
            db_path (str): Path to the CodeQL DB folder.
            temperature (float, optional): Sampling temperature. Defaults to 0.2.
            top_p (float, optional): Nucleus sampling. Defaults to 0.2.

        Returns:
            Tuple[List[Dict[str, Any]], str]:
                - The final conversation messages,
                - The final content from the LLM's last message.
        
        Raises:
            RuntimeError: If LLM model not initialized.
            LLMApiError: If LLM API call fails (rate limits, timeouts, auth failures, etc.).
            CodeQLError: If CodeQL database files cannot be read (from tool calls).
        """
        if not self.model:
            raise RuntimeError("LLM model not initialized. Call init_llm_client() first.")
        
        got_answer = False
        db_path_clean = db_path.replace(" ", "")
        all_functions = functions

        messages: List[Dict[str, Any]] = self.MESSAGES[:]
        messages.append({"role": "user", "content": prompt})
        messages = self._collapse_system_messages_for_minimax(messages)

        amount_of_tools = 0
        final_content = ""

        while not got_answer:
            # Send the current messages + tools to the LLM endpoint
            try:
                # Build completion kwargs - Bedrock Claude doesn't allow both temperature and top_p
                completion_kwargs = {
                    "model": self.model,
                    "messages": messages,
                    "tools": self.tools,
                    "timeout": self.timeout,
                }

                if self.api_key:
                    completion_kwargs["api_key"] = self.api_key
                if self.config and self.config.get("api_base"):
                    completion_kwargs["api_base"] = self.config["api_base"]

                # Check if using Bedrock (model starts with "bedrock/" or contains "arn:aws:bedrock")
                is_bedrock = (
                    self.model and
                    (self.model.startswith("bedrock/") or "arn:aws:bedrock" in self.model)
                )

                if is_bedrock:
                    # Bedrock Claude only accepts temperature OR top_p, not both
                    completion_kwargs["temperature"] = temperature
                else:
                    completion_kwargs["temperature"] = temperature
                    completion_kwargs["top_p"] = top_p

                for attempt in range(1, self.max_retries + 1):
                    try:
                        response = litellm.completion(**completion_kwargs)
                        break  # success
                    except Exception as e:
                        if self._is_retryable_llm_error(e) and attempt < self.max_retries:
                            logger.warning("Retryable LLM error (attempt %d/%d): %s", attempt, self.max_retries, e)
                            self._sleep_before_retry(attempt)
                            continue
                        if isinstance(e, litellm.APIConnectionError):
                            raise LLMApiError(f"LLM API connection failed: {e}") from e
                        if isinstance(e, litellm.RateLimitError):
                            raise LLMApiError(f"Rate limit exceeded for LLM API: {e}") from e
                        if isinstance(e, litellm.AuthenticationError):
                            raise LLMApiError(f"LLM API authentication failed: {e}") from e
                        if isinstance(e, litellm.APIError):
                            raise LLMApiError(f"LLM API error: {e}") from e
                        raise
            except litellm.RateLimitError as e:
                raise LLMApiError(f"Rate limit exceeded for LLM API: {e}") from e
            except litellm.Timeout as e:
                raise LLMApiError(f"LLM API request timed out after {self.max_retries} attempts ({self.timeout}s each): {e}") from e
            except litellm.APIConnectionError as e:
                raise LLMApiError(f"LLM API connection failed: {e}") from e
            except litellm.AuthenticationError as e:
                raise LLMApiError(f"LLM API authentication failed: {e}") from e
            except litellm.APIError as e:
                raise LLMApiError(f"LLM API error: {e}") from e
            except LLMApiError:
                raise
            except Exception as e:
                raise LLMApiError(f"Unexpected error during LLM API call: {e}") from e
            
            if not response.choices:
                raise LLMApiError(f"LLM API response is empty: {response}")

            content_obj = response.choices[0].message
            messages.append({
                "role": content_obj.role,
                "content": content_obj.content,
                "tool_calls": content_obj.tool_calls
            })

            final_content = content_obj.content or ""
            tool_calls = content_obj.tool_calls

            if not tool_calls:
                # Check if we have a recognized status code
                if final_content and any(code in final_content for code in ["1337", "1007", "7331", "3713"]):
                    got_answer = True
                else:
                    self._append_system_message(messages, "Please follow all the instructions!")
            else:
                amount_of_tools += 1
                arg_messages: List[Dict[str, Any]] = []

                for tc in tool_calls:
                    tool_call_id = tc.id
                    tool_function_name = tc.function.name
                    tool_args = tc.function.arguments

                    # Convert tool_args to a dict if it's a JSON string
                    if not isinstance(tool_args, dict):
                        tool_args = json.loads(tool_args)
                    else:
                        # Ensure consistent string for role=tool message
                        tc.function.arguments = json.dumps(tool_args)

                    response_msg = ""

                    # Evaluate which tool to call
                    if tool_function_name == 'get_function_code' and "function_name" in tool_args:
                        child_function, parent_function = self.db_lookup.get_function_by_name(
                            function_tree_file, tool_args["function_name"], all_functions
                        )
                        if isinstance(child_function, dict):
                            all_functions.append(child_function)
                        child_code = self.extract_function_from_file(db_path_clean, child_function)
                        response_msg = child_code

                        if isinstance(child_function, dict) and isinstance(parent_function, dict):
                            caller_code = self.extract_function_from_file(db_path_clean, parent_function)
                            args_content = self.map_func_args_by_llm(caller_code, child_code)
                            arg_messages.append({
                                "role": args_content.role,
                                "content": args_content.content
                            })

                    elif tool_function_name == 'get_caller_function':
                        caller_function = self.db_lookup.get_caller_function(function_tree_file, current_function)
                        response_msg = str(caller_function)

                        if isinstance(caller_function, dict):
                            all_functions.append(caller_function)
                            caller_code = self.extract_function_from_file(db_path_clean, caller_function)
                            response_msg = (
                                f"Here is the caller function for '{current_function['function_name']}':\n"
                                + caller_code
                            )
                            args_content = self.map_func_args_by_llm(
                                caller_code,
                                self.extract_function_from_file(db_path_clean, current_function)
                            )
                            arg_messages.append({
                                "role": args_content.role,
                                "content": args_content.content
                            })
                            current_function = caller_function

                    elif tool_function_name == 'get_macro' and "macro_name" in tool_args:
                        macro = self.db_lookup.get_macro(db_path_clean, tool_args["macro_name"])
                        if isinstance(macro, dict):
                            response_msg = macro["body"]
                        else:
                            response_msg = macro

                    elif tool_function_name == 'get_global_var' and "global_var_name" in tool_args:
                        global_var = self.db_lookup.get_global_var(db_path_clean, tool_args["global_var_name"])
                        if isinstance(global_var, dict):
                            global_var_code = self.extract_function_from_file(db_path_clean, global_var)
                            response_msg = global_var_code
                        else:
                            response_msg = global_var

                    elif tool_function_name == 'get_class' and "object_name" in tool_args:
                        curr_class = self.db_lookup.get_class(db_path_clean, tool_args["object_name"])
                        if isinstance(curr_class, dict):
                            class_code = self.extract_function_from_file(db_path_clean, curr_class)
                            response_msg = class_code
                        else:
                            response_msg = curr_class

                    else:
                        response_msg = (
                            f"No matching tool '{tool_function_name}' or invalid args {tool_args}. "
                            "Try again."
                        )

                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call_id,
                        "name": tool_function_name,
                        "content": response_msg
                    })

                messages += arg_messages

                if amount_of_tools >= 6:
                    self._append_system_message(
                        messages,
                        (
                            "You called too many tools! If you still can't give a clear answer, "
                            "return the 'more data' status."
                        ),
                    )

        return messages, final_content
