import logging
import os
from dataclasses import dataclass
from enum import Enum
from fnmatch import fnmatch
from pathlib import Path
from typing import Any

import litellm
from jinja2 import (
    Environment,
    FileSystemLoader,
    select_autoescape,
)
from litellm import ModelResponse, completion_cost
from litellm.utils import supports_prompt_caching

from exaaiagnt.llm.config import LLMConfig
from exaaiagnt.llm.llm_traffic_controller import RequestPriority, get_traffic_controller
from exaaiagnt.llm.memory_compressor import MemoryCompressor
from exaaiagnt.llm.utils import parse_tool_invocations
from exaaiagnt.prompts import load_prompt_modules
from exaaiagnt.tools import get_tools_prompt


logger = logging.getLogger(__name__)

api_key = os.getenv("LLM_API_KEY")
if api_key:
    litellm.api_key = api_key

api_base = (
    os.getenv("LLM_API_BASE")
    or os.getenv("OPENAI_API_BASE")
    or os.getenv("LITELLM_BASE_URL")
    or os.getenv("OLLAMA_API_BASE")
)
if api_base:
    litellm.api_base = api_base


class LLMRequestFailedError(Exception):
    def __init__(self, message: str, details: str | None = None):
        super().__init__(message)
        self.message = message
        self.details = details


SUPPORTS_STOP_WORDS_FALSE_PATTERNS: list[str] = [
    "o1*",
    "grok-4-0709",
    "grok-code-fast-1",
    "deepseek-r1-0528*",
]

REASONING_EFFORT_PATTERNS: list[str] = [
    "o1-2024-12-17",
    "o1",
    "o3",
    "o3-2025-04-16",
    "o3-mini-2025-01-31",
    "o3-mini",
    "o4-mini",
    "o4-mini-2025-04-16",
    "gemini-2.5-flash",
    "gemini-2.5-pro",
    "gemini-2.0*",
    "gpt-5*",
    "gpt-4.5*",
    "deepseek-r1-0528*",
    "deepseek-v3*",
    "claude-sonnet-4-5*",
    "claude-haiku-4-5*",
    "claude-opus-4*",
    "llama-3.3*",
    "llama-4*",
    "mistral-large-2*",
    "qwen-2.5*",
    "qwen-3*",
]


def normalize_model_name(model: str) -> str:
    raw = (model or "").strip().lower()
    if "/" in raw:
        name = raw.split("/")[-1]
        if ":" in name:
            name = name.split(":", 1)[0]
    else:
        name = raw
    if name.endswith("-gguf"):
        name = name[: -len("-gguf")]
    return name


def model_matches(model: str, patterns: list[str]) -> bool:
    raw = (model or "").strip().lower()
    name = normalize_model_name(model)
    for pat in patterns:
        pat_l = pat.lower()
        if "/" in pat_l:
            if fnmatch(raw, pat_l):
                return True
        elif fnmatch(name, pat_l):
            return True
    return False


class StepRole(str, Enum):
    AGENT = "agent"
    USER = "user"
    SYSTEM = "system"


@dataclass
class LLMResponse:
    content: str
    tool_invocations: list[dict[str, Any]] | None = None
    scan_id: str | None = None
    step_number: int = 1
    role: StepRole = StepRole.AGENT


@dataclass
class RequestStats:
    input_tokens: int = 0
    output_tokens: int = 0
    cached_tokens: int = 0
    cache_creation_tokens: int = 0
    cost: float = 0.0
    requests: int = 0
    failed_requests: int = 0

    def to_dict(self) -> dict[str, int | float]:
        return {
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cached_tokens": self.cached_tokens,
            "cache_creation_tokens": self.cache_creation_tokens,
            "cost": round(self.cost, 4),
            "requests": self.requests,
            "failed_requests": self.failed_requests,
        }


class LLM:
    def __init__(
        self, config: LLMConfig, agent_name: str | None = None, agent_id: str | None = None
    ):
        self.config = config
        self.agent_name = agent_name
        self.agent_id = agent_id
        self._total_stats = RequestStats()
        self._last_request_stats = RequestStats()

        self.memory_compressor = MemoryCompressor(
            model_name=self.config.model_name,
            timeout=self.config.timeout,
        )

        if agent_name:
            prompt_dir = Path(__file__).parent.parent / "agents" / agent_name
            prompts_dir = Path(__file__).parent.parent / "prompts"

            loader = FileSystemLoader([prompt_dir, prompts_dir])
            self.jinja_env = Environment(
                loader=loader,
                autoescape=select_autoescape(enabled_extensions=(), default_for_string=False),
            )

            try:
                prompt_module_content = load_prompt_modules(
                    self.config.prompt_modules or [], self.jinja_env
                )

                def get_module(name: str) -> str:
                    return prompt_module_content.get(name, "")

                self.jinja_env.globals["get_module"] = get_module

                self.system_prompt = self.jinja_env.get_template("system_prompt.jinja").render(
                    get_tools_prompt=get_tools_prompt,
                    loaded_module_names=list(prompt_module_content.keys()),
                    **prompt_module_content,
                )
            except (FileNotFoundError, OSError, ValueError) as e:
                logger.warning(f"Failed to load system prompt for {agent_name}: {e}")
                self.system_prompt = "You are a helpful AI assistant."
        else:
            self.system_prompt = "You are a helpful AI assistant."

    def set_agent_identity(self, agent_name: str | None, agent_id: str | None) -> None:
        if agent_name:
            self.agent_name = agent_name
        if agent_id:
            self.agent_id = agent_id

    def _build_identity_message(self) -> dict[str, Any] | None:
        if not (self.agent_name and str(self.agent_name).strip()):
            return None
        identity_name = self.agent_name
        identity_id = self.agent_id
        content = (
            "\n\n"
            "<agent_identity>\n"
            "<meta>Internal metadata: do not echo or reference; "
            "not part of history or tool calls.</meta>\n"
            "<note>You are now assuming the role of this agent. "
            "Act strictly as this agent and maintain self-identity for this step. "
            "Now go answer the next needed step!</note>\n"
            f"<agent_name>{identity_name}</agent_name>\n"
            f"<agent_id>{identity_id}</agent_id>\n"
            "</agent_identity>\n\n"
        )
        return {"role": "user", "content": content}

    def _add_cache_control_to_content(
        self, content: str | list[dict[str, Any]]
    ) -> str | list[dict[str, Any]]:
        if isinstance(content, str):
            return [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]
        if isinstance(content, list) and content:
            last_item = content[-1]
            if isinstance(last_item, dict) and last_item.get("type") == "text":
                return content[:-1] + [{**last_item, "cache_control": {"type": "ephemeral"}}]
        return content

    def _is_anthropic_model(self) -> bool:
        if not self.config.model_name:
            return False
        model_lower = self.config.model_name.lower()
        return any(provider in model_lower for provider in ["anthropic/", "claude"])

    def _calculate_cache_interval(self, total_messages: int) -> int:
        if total_messages <= 1:
            return 10

        max_cached_messages = 3
        non_system_messages = total_messages - 1

        interval = 10
        while non_system_messages // interval > max_cached_messages:
            interval += 10

        return interval

    def _prepare_cached_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        if (
            not self.config.enable_prompt_caching
            or not supports_prompt_caching(self.config.model_name)
            or not messages
        ):
            return messages

        if not self._is_anthropic_model():
            return messages

        cached_messages = list(messages)

        if cached_messages and cached_messages[0].get("role") == "system":
            system_message = cached_messages[0].copy()
            system_message["content"] = self._add_cache_control_to_content(
                system_message["content"]
            )
            cached_messages[0] = system_message

        total_messages = len(cached_messages)
        if total_messages > 1:
            interval = self._calculate_cache_interval(total_messages)

            cached_count = 0
            for i in range(interval, total_messages, interval):
                if cached_count >= 3:
                    break

                if i < len(cached_messages):
                    message = cached_messages[i].copy()
                    message["content"] = self._add_cache_control_to_content(message["content"])
                    cached_messages[i] = message
                    cached_count += 1

        return cached_messages

    async def generate(  # noqa: PLR0912, PLR0915
        self,
        conversation_history: list[dict[str, Any]],
        scan_id: str | None = None,
        step_number: int = 1,
    ) -> LLMResponse:
        messages = [{"role": "system", "content": self.system_prompt}]

        identity_message = self._build_identity_message()
        if identity_message:
            messages.append(identity_message)

        compressed_history = list(self.memory_compressor.compress_history(conversation_history))

        conversation_history.clear()
        conversation_history.extend(compressed_history)
        messages.extend(compressed_history)

        cached_messages = self._prepare_cached_messages(messages)

        try:
            response = await self._make_request(cached_messages)
            self._update_usage_stats(response)

            content = ""
            if (
                response.choices
                and hasattr(response.choices[0], "message")
                and response.choices[0].message
            ):
                content = getattr(response.choices[0].message, "content", "") or ""

            tool_invocations = parse_tool_invocations(content)

            return LLMResponse(
                scan_id=scan_id,
                step_number=step_number,
                role=StepRole.AGENT,
                content=content,
                tool_invocations=tool_invocations if tool_invocations else None,
            )

        except litellm.RateLimitError as e:
            raise LLMRequestFailedError("LLM request failed: Rate limit exceeded", str(e)) from e
        except litellm.AuthenticationError as e:
            raise LLMRequestFailedError("LLM request failed: Invalid API key", str(e)) from e
        except litellm.NotFoundError as e:
            raise LLMRequestFailedError("LLM request failed: Model not found", str(e)) from e
        except litellm.ContextWindowExceededError as e:
            raise LLMRequestFailedError("LLM request failed: Context too long", str(e)) from e
        except litellm.ContentPolicyViolationError as e:
            raise LLMRequestFailedError(
                "LLM request failed: Content policy violation", str(e)
            ) from e
        except litellm.ServiceUnavailableError as e:
            raise LLMRequestFailedError("LLM request failed: Service unavailable", str(e)) from e
        except litellm.Timeout as e:
            raise LLMRequestFailedError("LLM request failed: Request timed out", str(e)) from e
        except litellm.UnprocessableEntityError as e:
            raise LLMRequestFailedError("LLM request failed: Unprocessable entity", str(e)) from e
        except litellm.InternalServerError as e:
            raise LLMRequestFailedError("LLM request failed: Internal server error", str(e)) from e
        except litellm.APIConnectionError as e:
            raise LLMRequestFailedError("LLM request failed: Connection error", str(e)) from e
        except litellm.UnsupportedParamsError as e:
            raise LLMRequestFailedError("LLM request failed: Unsupported parameters", str(e)) from e
        except litellm.BudgetExceededError as e:
            raise LLMRequestFailedError("LLM request failed: Budget exceeded", str(e)) from e
        except litellm.APIResponseValidationError as e:
            raise LLMRequestFailedError(
                "LLM request failed: Response validation error", str(e)
            ) from e
        except litellm.JSONSchemaValidationError as e:
            raise LLMRequestFailedError(
                "LLM request failed: JSON schema validation error", str(e)
            ) from e
        except litellm.InvalidRequestError as e:
            raise LLMRequestFailedError("LLM request failed: Invalid request", str(e)) from e
        except litellm.BadRequestError as e:
            raise LLMRequestFailedError("LLM request failed: Bad request", str(e)) from e
        except litellm.APIError as e:
            raise LLMRequestFailedError("LLM request failed: API error", str(e)) from e
        except litellm.OpenAIError as e:
            raise LLMRequestFailedError("LLM request failed: OpenAI error", str(e)) from e
        except Exception as e:
            raise LLMRequestFailedError(f"LLM request failed: {type(e).__name__}", str(e)) from e

    @property
    def usage_stats(self) -> dict[str, dict[str, int | float]]:
        return {
            "total": self._total_stats.to_dict(),
            "last_request": self._last_request_stats.to_dict(),
        }

    def get_cache_config(self) -> dict[str, bool]:
        return {
            "enabled": self.config.enable_prompt_caching,
            "supported": supports_prompt_caching(self.config.model_name),
        }

    def _should_include_stop_param(self) -> bool:
        if not self.config.model_name:
            return True

        return not model_matches(self.config.model_name, SUPPORTS_STOP_WORDS_FALSE_PATTERNS)

    def _should_include_reasoning_effort(self) -> bool:
        if not self.config.model_name:
            return False

        return model_matches(self.config.model_name, REASONING_EFFORT_PATTERNS)

    async def _make_request(
        self,
        messages: list[dict[str, Any]],
    ) -> ModelResponse:
        completion_args: dict[str, Any] = {
            "model": self.config.model_name,
            "messages": messages,
            "timeout": self.config.timeout,
        }

        # Add max_tokens to limit output and reduce consumption
        if hasattr(self.config, "max_tokens_per_request") and self.config.max_tokens_per_request:
            completion_args["max_tokens"] = self.config.max_tokens_per_request

        if self._should_include_stop_param():
            completion_args["stop"] = ["</function>"]

        # Use medium reasoning effort in lightweight mode, high otherwise
        if self._should_include_reasoning_effort():
            if hasattr(self.config, "lightweight_mode") and self.config.lightweight_mode:
                completion_args["reasoning_effort"] = "medium"
            else:
                completion_args["reasoning_effort"] = "high"

        # Use Adaptive Traffic Controller for intelligent rate limiting with retries
        controller = get_traffic_controller()
        agent_id = self.agent_id or "unknown_agent"

        async def do_request():
            try:
                from litellm import completion
                return await completion(**completion_args, stream=False)
            except litellm.RateLimitError:
                # Let tenacity (in controller) handle retry
                raise
            except Exception as e:
                # Log other transient errors for potential retry
                logger.warning(f"Transient LLM error: {e}")
                raise

        # Wrap in retry logic
        import litellm
        from tenacity import (
            retry,
            retry_if_exception_type,
            stop_after_attempt,
            wait_exponential,
        )

        @retry(
            retry=retry_if_exception_type((
                litellm.RateLimitError,
                litellm.ServiceUnavailableError,
                litellm.APIConnectionError,
                litellm.Timeout
            )),
            wait=wait_exponential(multiplier=1, min=4, max=60),
            stop=stop_after_attempt(5),
            reraise=True
        )
        async def execute_with_retries():
            return await controller.queue_request(
                do_request,
                agent_id=agent_id,
                priority=RequestPriority.NORMAL
            )

        try:
            response = await execute_with_retries()
        except Exception:
            self._total_stats.failed_requests += 1
            raise

        self._total_stats.requests += 1
        self._last_request_stats = RequestStats(requests=1)

        return response

    def _update_usage_stats(self, response: ModelResponse) -> None:
        try:
            if hasattr(response, "usage") and response.usage:
                input_tokens = getattr(response.usage, "prompt_tokens", 0)
                output_tokens = getattr(response.usage, "completion_tokens", 0)

                cached_tokens = 0
                cache_creation_tokens = 0

                if hasattr(response.usage, "prompt_tokens_details"):
                    prompt_details = response.usage.prompt_tokens_details
                    if hasattr(prompt_details, "cached_tokens"):
                        cached_tokens = prompt_details.cached_tokens or 0

                if hasattr(response.usage, "cache_creation_input_tokens"):
                    cache_creation_tokens = response.usage.cache_creation_input_tokens or 0

            else:
                input_tokens = 0
                output_tokens = 0
                cached_tokens = 0
                cache_creation_tokens = 0

            try:
                cost = completion_cost(response) or 0.0
            except Exception as e:  # noqa: BLE001
                logger.warning(f"Failed to calculate cost: {e}")
                cost = 0.0

            self._total_stats.input_tokens += input_tokens
            self._total_stats.output_tokens += output_tokens
            self._total_stats.cached_tokens += cached_tokens
            self._total_stats.cache_creation_tokens += cache_creation_tokens
            self._total_stats.cost += cost

            self._last_request_stats.input_tokens = input_tokens
            self._last_request_stats.output_tokens = output_tokens
            self._last_request_stats.cached_tokens = cached_tokens
            self._last_request_stats.cache_creation_tokens = cache_creation_tokens
            self._last_request_stats.cost = cost

            if cached_tokens > 0:
                logger.info(f"Cache hit: {cached_tokens} cached tokens, {input_tokens} new tokens")
            if cache_creation_tokens > 0:
                logger.info(f"Cache creation: {cache_creation_tokens} tokens written to cache")

            logger.info(f"Usage stats: {self.usage_stats}")
        except Exception as e:  # noqa: BLE001
            logger.warning(f"Failed to update usage stats: {e}")
