"""Base Agent Implementation

Provides core LLM interaction, prompt templating, and retry logic.
"""

from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypeVar

from jinja2 import TemplateNotFound
from pydantic import BaseModel, ValidationError
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from ..errors import ReviewateError, ReviewateValidationError
from .types import (
    AgentConfig,
    AgentResponse,
    ReasoningEffort,
    TokenUsage,
    ToolHandler,
)
from .utils import (
    TEMPERATURE,
    build_messages,
    build_model_name,
    format_cached_diff,
    get_output_format,
    strip_markdown_json,
)

# Suppress LiteLLM verbose logging BEFORE importing
os.environ["LITELLM_LOG"] = "ERROR"

import httpx
import litellm
from jinja2 import Environment, FileSystemLoader, StrictUndefined
from litellm import acompletion

# SSL verification handling for self-hosted endpoints
# Only bypass SSL for explicitly whitelisted hosts (comma-separated list)
# Example: INSECURE_SSL_HOSTS=ollama.internal.company.com,llm.local:8080
_insecure_hosts_env = os.environ.get("INSECURE_SSL_HOSTS", "")
_insecure_hosts = frozenset(h.strip().lower() for h in _insecure_hosts_env.split(",") if h.strip())

if _insecure_hosts:
    # Only create insecure client if hosts are explicitly whitelisted
    # WARNING: This disables SSL verification - only use for trusted self-hosted endpoints
    import warnings

    warnings.warn(
        f"SSL verification disabled for hosts: {', '.join(_insecure_hosts)}. "
        "This is a security risk if these hosts are not trusted.",
        UserWarning,
        stacklevel=1,
    )
    _client = httpx.AsyncClient(verify=False)
    litellm.aclient_session = _client
# Default: full SSL verification (secure)

logger = logging.getLogger("base_agent")

# Disable LiteLLM logging after import too
logging.getLogger("LiteLLM").setLevel(logging.ERROR)
logging.getLogger("LiteLLM").propagate = False

if TYPE_CHECKING:
    from .output_parser import OutputParserAgent

T = TypeVar("T", bound=BaseModel)


class BaseAgentImpl:
    """Base implementation for all agents.

    Provides:
    - LLM calling with multiple providers via LiteLLM
    - Jinja2 template formatting
    - JSON response parsing with markdown stripping
    - Platform-specific output format instructions
    - Automatic retry with OutputParserAgent on parse failures
    """

    # Re-export utility functions as static methods for backward compatibility
    format_cached_diff = staticmethod(format_cached_diff)
    strip_markdown_json = staticmethod(strip_markdown_json)

    def __init__(
        self,
        name: str,
        description: str,
        config: AgentConfig,
        reasoning_effort: ReasoningEffort = None,
    ):
        """Initialize base agent.

        Args:
            name: Agent name for logging
            description: Agent description
            config: LLM configuration
            reasoning_effort: Optional reasoning level for thinking models.
                             None = no reasoning (default), "low"/"medium"/"high" = increasing depth.
        """
        self.name = name
        self.description = description
        self.config = config
        self.reasoning_effort = reasoning_effort

        # Setup Jinja2 environment for prompt templating
        prompts_dir = Path(__file__).parent.parent / "prompts"
        self.jinja_env = Environment(
            loader=FileSystemLoader(prompts_dir),
            autoescape=False,
            trim_blocks=True,
            lstrip_blocks=True,
            undefined=StrictUndefined,
        )

    def _build_messages(
        self,
        system_prompt: str,
        user_prompt: str,
        cached_content: str | None = None,
    ) -> list[dict[str, Any]]:
        """Build LLM message list with optional cached content."""
        return build_messages(system_prompt, user_prompt, cached_content)

    def _strip_cache_control(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Remove cache_control from messages.

        Used for Gemini when tools are present, since Gemini's CachedContent
        API doesn't support tools being passed in GenerateContent separately.
        """
        result = []
        for msg in messages:
            new_msg = msg.copy()
            content = new_msg.get("content")
            if isinstance(content, list):
                new_content = []
                for block in content:
                    if isinstance(block, dict) and "cache_control" in block:
                        new_block = {k: v for k, v in block.items() if k != "cache_control"}
                        new_content.append(new_block)
                    else:
                        new_content.append(block)
                new_msg["content"] = new_content
            result.append(new_msg)
        return result

    def _try_parse_response(self, content: str, response_model: type[T]) -> T | None:
        """Try to parse response content into Pydantic model.

        Handles markdown stripping and validation errors gracefully.

        Args:
            content: Raw response content
            response_model: Pydantic model to parse into

        Returns:
            Parsed model or None if parsing failed
        """
        if not content:
            return None

        cleaned = strip_markdown_json(content)

        try:
            return response_model.model_validate_json(cleaned)
        except (ValidationError, ValueError, TypeError):
            return None

    def get_output_format(self, response_model: type[T]) -> str:
        """Generate output format instructions for LLM."""
        return get_output_format(response_model)

    async def _make_request(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        response_format: type[T] | None = None,
        temperature: float = TEMPERATURE,
    ) -> tuple[Any, TokenUsage]:
        """Make a single LLM request.

        Low-level method that handles all provider-specific logic.
        Does not parse responses - returns raw LiteLLM response.

        Args:
            messages: Pre-built message list (use _build_messages for caching)
            tools: Optional tool definitions for tool-calling mode
            response_format: Optional Pydantic model for structured output
            temperature: LLM temperature

        Returns:
            Tuple of (raw LiteLLM response, TokenUsage)

        Raises:
            ReviewateError: If the LLM call fails
        """
        model_name = build_model_name(self.config.provider, self.config.model)

        # Build provider-specific kwargs
        extra_kwargs: dict[str, Any] = {}

        if self.config.api_base:
            extra_kwargs["api_base"] = self.config.api_base

        # OpenRouter provider order
        if self.config.provider == "openrouter":
            provider_order_env = None
            if self.config.tier:
                tier_env_var = f"{self.config.tier.upper()}_OPENROUTER_PROVIDER_ORDER"
                provider_order_env = os.getenv(tier_env_var)
            if not provider_order_env:
                provider_order_env = os.getenv("OPENROUTER_PROVIDER_ORDER")
            if not provider_order_env and self.config.tier:
                from ..config_file import get_config_openrouter_order, load_config_file

                file_config = load_config_file()
                if file_config:
                    provider_order_env = get_config_openrouter_order(self.config.tier, file_config)
            if provider_order_env:
                provider_order = [p.strip() for p in provider_order_env.split(",") if p.strip()]
                if provider_order:
                    extra_kwargs["extra_body"] = {
                        "provider": {"order": provider_order, "allow_fallbacks": False}
                    }

        # Reasoning effort for thinking models
        thinking_disabled = os.getenv("DISABLE_THINKING") == "1"
        reasoning = None if thinking_disabled else self.reasoning_effort

        # Claude requires temperature=1.0 for reasoning
        effective_temp = temperature
        if "claude" in self.config.model and reasoning is not None:
            effective_temp = 1.0

        # Gemini's explicit CachedContent API requires system_instruction/tools
        # to be in the cache, not in GenerateContent. Strip cache_control from
        # messages when using Gemini with tools to avoid this conflict.
        # Note: Gemini's implicit caching (enabled by default) still works.
        effective_messages = messages
        if self.config.provider == "gemini" and tools:
            effective_messages = self._strip_cache_control(messages)

        call_kwargs: dict[str, Any] = {
            "model": model_name,
            "messages": effective_messages,
            "temperature": effective_temp,
            "api_key": self.config.api_key,
            "caching": True,
            "drop_params": True,
            **extra_kwargs,
        }

        if tools:
            call_kwargs["tools"] = tools
            call_kwargs["tool_choice"] = "auto"

        if response_format:
            call_kwargs["response_format"] = response_format

        if reasoning:
            call_kwargs["reasoning_effort"] = reasoning
            if self.config.provider == "openrouter":
                call_kwargs["allowed_openai_params"] = ["reasoning_effort"]

        try:
            response = await acompletion(**call_kwargs)
        except Exception as e:
            raise ReviewateError(f"LLM call failed: {e}") from e

        # Extract token usage
        usage = TokenUsage.from_litellm_usage(response.usage, model_name)

        return response, usage

    async def _call_with_tools(
        self,
        system_prompt: str,
        user_prompt: str,
        response_model: type[T],
        tools: list[dict[str, Any]],
        tool_handler: ToolHandler,
        cached_content: str | None = None,
        max_iterations: int = 5,
        temperature: float = 0.1,
    ) -> tuple[T, TokenUsage]:
        """Call LLM with tool support in an iteration loop.

        Iterates up to max_iterations, allowing the LLM to call tools.
        On the last iteration, forces a structured response without tools.
        Only the first tool call per iteration is processed (sequential mode)
        so the model sees each result before making the next call.

        Args:
            system_prompt: System prompt for the agent
            user_prompt: User prompt (typically output format instructions)
            response_model: Pydantic model for final structured response
            tools: List of tool definitions (OpenAI format)
            tool_handler: Async callback to process tool calls
            cached_content: Optional content to cache (e.g., diff). Uses proper
                           cache_control blocks for cache hits.
            max_iterations: Maximum iterations before forcing final response
            temperature: LLM temperature

        Returns:
            Tuple of (parsed response, aggregated token usage)

        Raises:
            ReviewateError: If max iterations reached without valid response
        """
        model_name = build_model_name(self.config.provider, self.config.model)

        # Build initial messages with proper cache_control
        messages = self._build_messages(system_prompt, user_prompt, cached_content)

        total_usage = TokenUsage.zero(model=model_name)
        tool_usage = TokenUsage.zero(model=model_name)

        for iteration in range(max_iterations):
            is_last = iteration == max_iterations - 1
            logger.debug(f"{self.name} tool iteration {iteration + 1}/{max_iterations}")

            # Prompt guard: on the last iteration, tell the model to stop
            # calling tools and return its final JSON response.
            if is_last:
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "This is your final turn. Do not call any tools. "
                            "Provide your response as valid JSON matching the required schema."
                        ),
                    }
                )

            # Always pass tools so Anthropic doesn't reject requests that
            # contain tool history (role:"tool", tool_calls) without tools=.
            response, usage = await self._make_request(
                messages=messages,
                tools=tools,
                response_format=response_model if is_last else None,
                temperature=temperature,
            )
            total_usage = total_usage + usage

            logger.debug(
                f"{self.name} iteration {iteration + 1}: "
                f"{usage.total_tokens} tokens (cached={usage.cached_tokens})"
            )

            # Handle empty response
            if not response.choices:
                raise ReviewateError(
                    f"{self.name} returned empty choices at iteration {iteration + 1}"
                )

            message = response.choices[0].message

            # Handle tool calls — skip on last iteration even if the model
            # returns them; parse message.content as the final answer instead.
            if message.tool_calls and not is_last:
                # Only process the first tool call per iteration so the model sees
                # each result before making the next call. This also caps the
                # assistant message to match, avoiding "missing tool result" errors.
                tool_calls_to_process = message.tool_calls[:1]

                logger.debug(f"{self.name} processing {len(tool_calls_to_process)} tool call(s)")

                # Log the model's reasoning alongside tool calls
                if message.content:
                    logger.debug(f"{self.name} reasoning: {message.content}")
                if getattr(message, "reasoning_content", None):
                    logger.debug(f"{self.name} thinking: {message.reasoning_content}")

                assistant_msg: dict[str, Any] = {
                    "role": "assistant",
                    "content": message.content,
                    "tool_calls": [tc.model_dump() for tc in tool_calls_to_process],
                }
                # Preserve reasoning_content for thinking models (e.g., Kimi K2.5)
                if getattr(message, "reasoning_content", None):
                    assistant_msg["reasoning_content"] = message.reasoning_content
                messages.append(assistant_msg)

                for tool_call in tool_calls_to_process:
                    result = await tool_handler(tool_call)

                    if result.usage:
                        total_usage = total_usage + result.usage
                        tool_usage = tool_usage + result.usage

                    messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": result.content,
                        }
                    )
                continue

            # No tool calls - try to parse response
            content = message.content or ""
            parsed = self._try_parse_response(content, response_model)

            if parsed is not None:
                # Attach tool usage as exploration_cost if any
                if tool_usage.total_tokens > 0:
                    total_usage = TokenUsage(
                        input_tokens=total_usage.input_tokens,
                        output_tokens=total_usage.output_tokens,
                        total_tokens=total_usage.total_tokens,
                        cached_tokens=total_usage.cached_tokens,
                        model=total_usage.model,
                        exploration_cost=tool_usage,
                    )
                return parsed, total_usage

            # Parse failed - add retry message if not last iteration
            if not is_last:
                logger.warning(f"{self.name} failed to parse response, retrying...")
                messages.append({"role": "assistant", "content": content})
                messages.append(
                    {
                        "role": "user",
                        "content": "Invalid JSON. Respond with valid JSON matching the schema.",
                    }
                )

        raise ReviewateError(f"{self.name} failed after {max_iterations} iterations")

    async def call_llm(
        self,
        system_prompt: str,
        user_prompt: str,
        response_model: type[T],
        cached_content: str | None = None,
    ) -> tuple[T, TokenUsage]:
        """Call LLM and parse response into Pydantic model.

        Args:
            system_prompt: The system prompt (agent-specific instructions)
            user_prompt: The user prompt (task/question)
            response_model: Pydantic model class to parse response into
            cached_content: Optional content to cache (e.g., diff). Uses proper
                           cache_control blocks for cache hits.

        Returns:
            Tuple of (parsed output, token usage)

        Raises:
            ReviewateError: If LLM call fails or response cannot be parsed
        """
        # Build messages with proper cache_control
        messages = self._build_messages(system_prompt, user_prompt, cached_content)

        # Make request
        response, usage = await self._make_request(
            messages=messages,
            response_format=response_model,
        )

        # Extract content
        if not response.choices or not response.choices[0].message:
            raise ReviewateError("LLM returned empty response")

        message = response.choices[0].message
        content = message.content
        if not content:
            raise ReviewateError("LLM returned empty content")

        # Parse JSON into Pydantic model
        try:
            # Strip markdown code blocks if present (some models wrap JSON in ```json ... ```)
            cleaned_content = strip_markdown_json(content)
            parsed = response_model.model_validate_json(cleaned_content)
        except (ValidationError, ValueError) as e:
            raise ReviewateValidationError(
                f"Failed to parse LLM response into {response_model}: {e}",
                invalid_response=content,
            ) from e

        return parsed, usage

    def load_prompt(self, template_path: str) -> str:
        """Load a prompt template file.

        Args:
            template_path: Relative path to template (e.g., "business_agent.txt")

        Returns:
            Template content as string

        Raises:
            ReviewateError: If template cannot be loaded
        """
        try:
            # Get template source using the loader
            source, _, _ = self.jinja_env.loader.get_source(self.jinja_env, template_path)
            return source
        except TemplateNotFound as e:
            raise ReviewateError(f"Failed to load prompt template {template_path}: {e}") from e

    @retry(
        retry=retry_if_exception_type((ReviewateError, Exception)),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=10, min=1, max=100),
        reraise=True,
    )
    async def analyze(
        self,
        system_prompt: str,
        user_prompt: str,
        response_model: type[T],
        output_parser: OutputParserAgent | None = None,
        cached_content: str | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_handler: ToolHandler | None = None,
        max_iterations: int = 5,
    ) -> AgentResponse[T]:
        """Analyze with automatic retry on JSON parse failures.

        This is the main entry point for agents. It calls the LLM and automatically
        retries with OutputParserAgent if the response cannot be parsed.

        Supports optional tool-calling: when tools and tool_handler are provided,
        runs an iteration loop allowing the LLM to call tools before returning.

        Uses tenacity for exponential backoff retry:
        - Max 3 attempts
        - Exponential backoff: 1s, 10s, 100s
        - Retries on ReviewateError and general exceptions

        Args:
            system_prompt: The system prompt (agent instructions)
            user_prompt: The user prompt (task/question)
            response_model: Pydantic model class to parse response into
            output_parser: Optional OutputParserAgent for retry logic
            cached_content: Optional content to cache (e.g., diff). For Anthropic/OpenRouter,
                           this is sent as a separate user message and cached across calls.
            tools: Optional list of tool definitions (OpenAI format). When provided with
                  tool_handler, enables tool-calling mode.
            tool_handler: Optional async callback to process tool calls. Required if tools
                         is provided.
            max_iterations: Maximum iterations for tool-calling mode (default: 5)

        Returns:
            AgentResponse with parsed output and token usage

        Raises:
            ReviewateError: If analysis fails after retries
        """
        try:
            # Tool-calling mode: use _call_with_tools
            if tools is not None and tool_handler is not None:
                output, token_usage = await self._call_with_tools(
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    response_model=response_model,
                    tools=tools,
                    tool_handler=tool_handler,
                    cached_content=cached_content,
                    max_iterations=max_iterations,
                )
                return AgentResponse(output=output, metadata=token_usage)

            # Standard mode: single LLM call
            output, token_usage = await self.call_llm(
                system_prompt, user_prompt, response_model, cached_content=cached_content
            )
            return AgentResponse(output=output, metadata=token_usage)

        except ReviewateValidationError as e:
            # If we have an output parser, try to fix the response
            if output_parser and "Failed to parse LLM response" in str(e):
                logger.warning(f"Validation error, attempting fix with output parser: {e}")
                # Extract the JSON schema for the output parser
                schema = response_model.model_json_schema()
                json_schema = json.dumps(schema, indent=2)

                # Get the invalid response from the error message
                # (we'll need to capture it from call_llm - for now, retry)
                try:
                    output, retry_usage = await output_parser.parse_output(
                        invalid_response=e.invalid_response,
                        json_schema=json_schema,
                        response_model=response_model,
                    )
                    logger.info("Successfully recovered from validation error using output parser")
                    return AgentResponse(output=output, metadata=retry_usage)
                except Exception as retry_error:
                    raise ReviewateError(
                        f"Failed to parse response even after retry: {retry_error}"
                    ) from retry_error

            # No output parser or different error, re-raise
            raise
        except Exception as e:
            logger.error(f"Agent {self.name} analyze failed: {e}")
            raise e
