"""Dynamic ghost tool generation using LLM analysis."""

import json
import logging
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

from honeymcp.llm.analyzers import ToolInfo
from honeymcp.llm.clients import get_chat_llm_client
from honeymcp.llm.prompts import format_prompt
from honeymcp.models.ghost_tool_spec import GhostToolSpec

logger = logging.getLogger(__name__)


@dataclass
class ServerContext:
    """Analysis of what the MCP server does."""

    server_purpose: str
    """Brief description of the server's purpose"""

    domain: str
    """Primary domain (file_system, database, api, etc.)"""

    real_tool_names: List[str]
    """Names of real tools available on the server"""

    real_tool_descriptions: List[str]
    """Descriptions of real tools"""

    security_sensitive_areas: List[str]
    """Security-sensitive areas identified for this domain"""


@dataclass
class DynamicGhostToolSpec(GhostToolSpec):
    """Extended specification for dynamically generated ghost tools."""

    server_context: ServerContext
    """Context about the server this tool was generated for"""

    generation_timestamp: datetime
    """When this tool was generated"""

    llm_generated: bool = True
    """Flag indicating this was generated by LLM"""

    fake_response: str = ""
    """Pre-generated response content returned when tool is triggered"""


class DynamicGhostToolGenerator:
    """Generates context-aware ghost tools using LLM analysis."""

    def __init__(
        self,
        llm_client: Optional[Any] = None,
        cache_ttl: int = 3600,
        model_name: Optional[str] = None,
        model_parameters: Optional[Dict[str, Any]] = None,
    ):
        """Initialize the dynamic ghost tool generator.

        Args:
            llm_client: LLM client instance (creates default if None)
            cache_ttl: Cache time-to-live in seconds
            model_name: Optional model name override
            model_parameters: Optional model parameters for LLM client
        """
        self.llm_client = llm_client
        self.model_name = model_name
        self.model_parameters = model_parameters or {}
        self.cache_ttl = cache_ttl
        self._cache: Dict[str, Any] = {}
        self._cache_timestamps: Dict[str, datetime] = {}
        self._client_cache: Dict[float, Any] = {}

    def _get_llm_client(self, temperature: float) -> Any:
        if self.llm_client is not None:
            return self.llm_client

        if temperature in self._client_cache:
            return self._client_cache[temperature]

        parameters = dict(self.model_parameters)
        parameters["temperature"] = temperature
        client = get_chat_llm_client(
            model_name=self.model_name or "rits/openai/gpt-oss-120b",
            model_parameters=parameters,
        )
        self._client_cache[temperature] = client
        return client

    @staticmethod
    def _format_messages(messages: List[Dict[str, str]]) -> str:
        parts: List[str] = []
        for message in messages:
            role = message.get("role", "user")
            content = message.get("content", "")
            if role == "system":
                parts.append(f"System: {content}")
            elif role == "assistant":
                parts.append(f"Assistant: {content}")
            else:
                parts.append(content)
        return "\n\n".join([part for part in parts if part])

    def _generate_response(self, messages: List[Dict[str, str]], temperature: float) -> str:
        client = self._get_llm_client(temperature)
        prompt = self._format_messages(messages)
        response = client.invoke(prompt)
        if hasattr(response, "content"):
            return str(response.content)
        return str(response)

    async def analyze_server_context(self, real_tools: List[ToolInfo]) -> ServerContext:
        """Analyze the server's purpose and context using LLM.

        Args:
            real_tools: List of real tools extracted from the server

        Returns:
            ServerContext with analysis results

        Raises:
            ValueError: If LLM returns invalid JSON or analysis fails
        """
        # Check cache
        cache_key = "server_context_" + "_".join(sorted([t.name for t in real_tools]))
        if self._is_cache_valid(cache_key):
            logger.info("Using cached server context analysis")
            return self._cache[cache_key]

        # Prepare tools for analysis
        tools_dict = [{"name": tool.name, "description": tool.description} for tool in real_tools]
        tool_list = [
            f"{i}. {tool['name']}: {tool['description']}" for i, tool in enumerate(tools_dict, 1)
        ]
        tool_list_str = "\n".join(tool_list) if tool_list else "No tools available"

        # Format prompt
        prompt = format_prompt(
            "server_analysis_prompt",
            prompt_file="dynamic_ghost_tools",
            tool_list=tool_list_str,
        )

        # Call LLM
        logger.info("Analyzing server context with %s tools", len(real_tools))
        try:
            messages = [{"role": "user", "content": prompt}]
            response = self._generate_response(messages, temperature=0.3)

            # Parse JSON response
            # Handle None response from LLM
            if response is None:
                raise ValueError("LLM returned empty response")

            # Try to extract JSON from response (handle cases where LLM adds extra text)
            response = response.strip()
            if "```json" in response:
                response = response.split("```json")[1].split("```")[0].strip()
            elif "```" in response:
                response = response.split("```")[1].split("```")[0].strip()

            analysis = json.loads(response)

            # Validate required fields
            required_fields = ["server_purpose", "domain", "security_sensitive_areas"]
            for field in required_fields:
                if field not in analysis:
                    raise ValueError(f"Missing required field in LLM response: {field}")

            # Create ServerContext
            context = ServerContext(
                server_purpose=analysis["server_purpose"],
                domain=analysis["domain"],
                real_tool_names=[t.name for t in real_tools],
                real_tool_descriptions=[t.description for t in real_tools],
                security_sensitive_areas=analysis["security_sensitive_areas"],
            )

            # Cache result
            self._cache[cache_key] = context
            self._cache_timestamps[cache_key] = datetime.utcnow()

            logger.info(
                "Server context analyzed: domain=%s, purpose=%s...",
                context.domain,
                context.server_purpose[:50],
            )
            return context

        except json.JSONDecodeError as e:
            logger.error("Failed to parse LLM response as JSON: %s", e)
            logger.error("Response was: %s", response)
            raise ValueError(f"LLM returned invalid JSON: {e}") from e
        except Exception as e:
            logger.error("Error analyzing server context: %s", e)
            raise

    async def generate_ghost_tools(
        self, server_context: ServerContext, num_tools: int = 3
    ) -> List[DynamicGhostToolSpec]:
        """Generate context-aware ghost tools using LLM.

        Args:
            server_context: Analysis of the server's purpose and domain
            num_tools: Number of ghost tools to generate

        Returns:
            List of dynamically generated ghost tool specifications

        Raises:
            ValueError: If LLM returns invalid JSON or generation fails
        """
        # Check cache
        cache_key = f"ghost_tools_{server_context.domain}_{num_tools}"
        if self._is_cache_valid(cache_key):
            logger.info("Using cached ghost tools")
            return self._cache[cache_key]

        # Format prompt
        prompt = format_prompt(
            "ghost_tool_generation_prompt",
            prompt_file="dynamic_ghost_tools",
            server_purpose=server_context.server_purpose,
            domain=server_context.domain,
            real_tool_names=", ".join(server_context.real_tool_names),
            security_areas=", ".join(server_context.security_sensitive_areas),
            num_tools=num_tools,
        )

        # Call LLM
        logger.info(
            "Generating %s ghost tools for domain: %s",
            num_tools,
            server_context.domain,
        )
        try:
            messages = [{"role": "user", "content": prompt}]
            response = self._generate_response(messages, temperature=0.7)

            # Parse JSON response
            # Handle None response from LLM
            if response is None:
                raise ValueError("LLM returned empty response")

            response = response.strip()
            if "```json" in response:
                response = response.split("```json")[1].split("```")[0].strip()
            elif "```" in response:
                response = response.split("```")[1].split("```")[0].strip()

            tools_data = json.loads(response)

            if not isinstance(tools_data, list):
                raise ValueError("LLM response must be a JSON array")

            # Create DynamicGhostToolSpec objects
            ghost_tools = []
            for tool_data in tools_data:
                # Validate required fields
                required_fields = [
                    "name",
                    "description",
                    "parameters",
                    "threat_level",
                    "attack_category",
                ]
                for field in required_fields:
                    if field not in tool_data:
                        raise ValueError(f"Missing required field in tool spec: {field}")

                # Get pre-generated fake response (with fallback)
                fake_response = tool_data.get("fake_response", "")
                if not fake_response:
                    logger.warning(
                        "No fake_response provided for %s, using generic fallback",
                        tool_data["name"],
                    )
                    fake_response = f"Operation completed successfully.\nTool: {tool_data['name']}"

                # Create response generator function using pre-generated response
                response_generator = self._create_response_generator(fake_response)

                ghost_tool = DynamicGhostToolSpec(
                    name=tool_data["name"],
                    description=tool_data["description"],
                    parameters=tool_data["parameters"],
                    response_generator=response_generator,
                    threat_level=tool_data["threat_level"],
                    attack_category=tool_data["attack_category"],
                    server_context=server_context,
                    generation_timestamp=datetime.utcnow(),
                    llm_generated=True,
                    fake_response=fake_response,
                )
                ghost_tools.append(ghost_tool)

            # Cache result
            self._cache[cache_key] = ghost_tools
            self._cache_timestamps[cache_key] = datetime.utcnow()

            logger.info(
                "Generated %s ghost tools: %s",
                len(ghost_tools),
                [t.name for t in ghost_tools],
            )
            return ghost_tools

        except json.JSONDecodeError as e:
            logger.error("Failed to parse LLM response as JSON: %s", e)
            logger.error("Response was: %s", response)
            raise ValueError(f"LLM returned invalid JSON: {e}") from e
        except Exception as e:
            logger.error("Error generating ghost tools: %s", e)
            raise

    def _create_response_generator(self, fake_response: str) -> Callable[[Dict[str, Any]], str]:
        """Create a response generator function for a ghost tool.

        Uses the pre-generated fake response with optional argument interpolation.

        Args:
            fake_response: Pre-generated response template with optional {param} placeholders

        Returns:
            Function that returns the fake response with interpolated arguments
        """

        def generate_response(arguments: Dict[str, Any]) -> str:
            """Return pre-generated fake response with argument interpolation."""
            try:
                # Interpolate arguments into the pre-generated response
                return fake_response.format(**arguments)
            except KeyError:
                # Fallback if placeholder doesn't match argument names
                return fake_response

        return generate_response

    def _is_cache_valid(self, key: str) -> bool:
        """Check if a cache entry is still valid.

        Args:
            key: Cache key to check

        Returns:
            True if cache entry exists and is not expired
        """
        if key not in self._cache or key not in self._cache_timestamps:
            return False

        age = (datetime.utcnow() - self._cache_timestamps[key]).total_seconds()
        return age < self.cache_ttl

    def clear_cache(self):
        """Clear all cached data."""
        self._cache.clear()
        self._cache_timestamps.clear()
        logger.info("Cache cleared")

    async def generate_real_tool_mocks(
        self, real_tools: List[ToolInfo], server_context: ServerContext
    ) -> Dict[str, str]:
        """Generate fake responses for real tools (used in cognitive protection mode).

        Args:
            real_tools: List of real tools to generate mocks for
            server_context: Analysis of the server's purpose and domain

        Returns:
            Dictionary mapping tool_name -> mock_response template
        """
        # Check cache
        cache_key = f"real_tool_mocks_{server_context.domain}_{len(real_tools)}"
        if self._is_cache_valid(cache_key):
            logger.info("Using cached real tool mocks")
            return self._cache[cache_key]

        # Prepare tools for prompt
        tools_dict = [{"name": tool.name, "description": tool.description} for tool in real_tools]
        tool_list = [
            f"{i}. {tool['name']}: {tool['description']}" for i, tool in enumerate(tools_dict, 1)
        ]
        tool_list_str = "\n".join(tool_list) if tool_list else "No tools available"

        # Format prompt
        prompt = format_prompt(
            "real_tool_mock_generation_prompt",
            prompt_file="dynamic_ghost_tools",
            server_purpose=server_context.server_purpose,
            domain=server_context.domain,
            tool_list=tool_list_str,
        )

        # Call LLM
        logger.info("Generating mock responses for %s real tools", len(real_tools))
        try:
            messages = [{"role": "user", "content": prompt}]
            response = self._generate_response(messages, temperature=0.5)

            # Handle None response from LLM
            if response is None:
                raise ValueError("LLM returned empty response")

            # Try to extract JSON from response
            response = response.strip()
            if "```json" in response:
                response = response.split("```json")[1].split("```")[0].strip()
            elif "```" in response:
                response = response.split("```")[1].split("```")[0].strip()

            mocks_data = json.loads(response)

            if not isinstance(mocks_data, list):
                raise ValueError("LLM response must be a JSON array")

            # Build dictionary of mock responses
            real_tool_mocks: Dict[str, str] = {}
            for mock_data in mocks_data:
                name = mock_data.get("name")
                mock_response = mock_data.get("mock_response", "")
                if name and mock_response:
                    real_tool_mocks[name] = mock_response

            # Cache result
            self._cache[cache_key] = real_tool_mocks
            self._cache_timestamps[cache_key] = datetime.utcnow()

            logger.info("Generated mocks for %s real tools", len(real_tool_mocks))
            return real_tool_mocks

        except json.JSONDecodeError as e:
            logger.error("Failed to parse LLM response as JSON: %s", e)
            logger.error("Response was: %s", response)
            raise ValueError(f"LLM returned invalid JSON: {e}") from e
        except Exception as e:
            logger.error("Error generating real tool mocks: %s", e)
            raise
