import logging
from typing import Any, Callable, Literal, Optional, Union

import httpx
from agentsts.adk import ADKTokenPropagationPlugin
from google.adk.agents import Agent
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import ToolUnion
from google.adk.agents.readonly_context import ReadonlyContext
from google.adk.agents.remote_a2a_agent import AGENT_CARD_WELL_KNOWN_PATH, DEFAULT_TIMEOUT
from google.adk.models.anthropic_llm import Claude as ClaudeLLM
from google.adk.models.google_llm import Gemini as GeminiLLM
from google.adk.tools.mcp_tool import SseConnectionParams, StreamableHTTPConnectionParams
from pydantic import BaseModel, Field

from kagent.adk._approval import make_approval_callback
from kagent.adk._mcp_toolset import KAgentMcpToolset
from kagent.adk._remote_a2a_tool import KAgentRemoteA2AToolset
from kagent.adk.models._anthropic import KAgentAnthropicLlm
from kagent.adk.models._bedrock import KAgentBedrockLlm
from kagent.adk.models._ollama import create_ollama_llm
from kagent.adk.sandbox_code_executer import SandboxedLocalCodeExecutor
from kagent.adk.tools.ask_user_tool import AskUserTool

from .models import AzureOpenAI as OpenAIAzure
from .models import OpenAI as OpenAINative

logger = logging.getLogger(__name__)

# Proxy host header used for Gateway API routing when using a proxy
PROXY_HOST_HEADER = "x-kagent-host"

# Key used to store headers in session state
HEADERS_STATE_KEY = "headers"


def create_header_provider(
    allowed_headers: list[str] | None = None,
    sts_header_provider: Callable[[Optional[ReadonlyContext]], dict[str, str]] | None = None,
) -> Callable[[Optional[ReadonlyContext]], dict[str, str]] | None:
    """Create a header provider that combines STS tokens and allowed headers.

    Args:
        allowed_headers: List of header names to propagate from A2A request.
                        Header names are case-insensitive.
                        Authorization headers CAN be propagated if explicitly listed.
        sts_header_provider: Optional STS header provider for token propagation.
                            When provided, STS-generated headers (e.g., Authorization)
                            will take precedence over any matching allowed headers.
                            This is a security measure to prevent request headers from
                            overwriting authentication tokens generated by STS.

    Returns:
        A header provider function, or None if no headers need to be propagated.
    """
    if not allowed_headers and not sts_header_provider:
        return None

    # Normalize header names to lowercase for case-insensitive matching
    normalized_allowed = [h.lower() for h in allowed_headers] if allowed_headers else []

    def header_provider(readonly_context: Optional[ReadonlyContext]) -> dict[str, str]:
        headers: dict[str, str] = {}

        # Add allowed headers from session state first
        if normalized_allowed and readonly_context:
            request_headers = readonly_context.state.get(HEADERS_STATE_KEY, {})
            for header_name, header_value in request_headers.items():
                if header_name.lower() in normalized_allowed:
                    headers[header_name] = header_value

        # Add STS headers last so they take precedence (security: prevent
        # allowed headers from overwriting authentication tokens)
        # Use case-insensitive replacement to handle header name case variations
        if sts_header_provider:
            sts_headers = sts_header_provider(readonly_context)
            if sts_headers:
                for sts_key, sts_value in sts_headers.items():
                    # Remove any existing header with same name (case-insensitive)
                    keys_to_remove = [k for k in headers if k.lower() == sts_key.lower()]
                    for k in keys_to_remove:
                        del headers[k]
                    headers[sts_key] = sts_value

        return headers

    return header_provider


def _convert_ollama_options(options: dict[str, str] | None) -> dict[str, Any]:
    """Convert Ollama options from string values to their correct types.

    Uses type annotations from the official Ollama Python SDK Options class
    to determine the correct type for each option.

    Since the options come from YAML/JSON config as strings, we need to convert them
    to the types expected by the Ollama API.
    """
    if not options:
        return {}

    # Import Ollama SDK Options class for type introspection
    from typing import get_args, get_origin

    from ollama import Options as OllamaOptions

    # Get type hints from the Ollama SDK Options class
    type_hints = OllamaOptions.model_fields

    converted: dict[str, Any] = {}
    for key, value in options.items():
        if key in type_hints:
            field_info = type_hints[key]
            # Get the annotation and unwrap Optional[T] -> T
            annotation = field_info.annotation
            origin = get_origin(annotation)
            if origin is Union:
                # Optional[T] is Union[T, None], get the non-None type
                args = [arg for arg in get_args(annotation) if arg is not type(None)]
                if args:
                    target_type = args[0]
                else:
                    target_type = str
            else:
                target_type = annotation

            # Convert based on the target type
            if target_type is int:
                converted[key] = int(value)
            elif target_type is float:
                converted[key] = float(value)
            elif target_type is bool:
                converted[key] = value.lower() == "true"
            else:
                # Keep as string for other types (e.g., Sequence[str])
                converted[key] = value
        else:
            # Unknown option - keep as string
            converted[key] = value

    return converted


class HttpMcpServerConfig(BaseModel):
    params: StreamableHTTPConnectionParams
    tools: list[str] = Field(default_factory=list)
    allowed_headers: list[str] | None = None  # Headers to forward from A2A request to MCP calls
    require_approval: list[str] | None = None  # Tools requiring human approval before execution


class SseMcpServerConfig(BaseModel):
    params: SseConnectionParams
    tools: list[str] = Field(default_factory=list)
    allowed_headers: list[str] | None = None  # Headers to forward from A2A request to MCP calls
    require_approval: list[str] | None = None  # Tools requiring human approval before execution


class RemoteAgentConfig(BaseModel):
    name: str
    url: str
    headers: dict[str, Any] | None = None
    timeout: float = DEFAULT_TIMEOUT
    description: str = ""


class BaseLLM(BaseModel):
    model: str
    headers: dict[str, str] | None = None

    # TLS/SSL configuration (applies to all model types)
    tls_disable_verify: bool | None = None
    tls_ca_cert_path: str | None = None
    tls_disable_system_cas: bool | None = None

    # API key passthrough: forward the Bearer token from incoming requests as the LLM API key
    api_key_passthrough: bool | None = None


class OpenAI(BaseLLM):
    base_url: str | None = None
    frequency_penalty: float | None = None
    max_tokens: int | None = None
    n: int | None = None
    presence_penalty: float | None = None
    reasoning_effort: str | None = None
    seed: int | None = None
    temperature: float | None = None
    timeout: int | None = None
    top_p: float | None = None

    type: Literal["openai"]


class AzureOpenAI(BaseLLM):
    type: Literal["azure_openai"]


class Anthropic(BaseLLM):
    base_url: str | None = None

    type: Literal["anthropic"]


class GeminiVertexAI(BaseLLM):
    type: Literal["gemini_vertex_ai"]


class GeminiAnthropic(BaseLLM):
    type: Literal["gemini_anthropic"]


class Ollama(BaseLLM):
    options: dict[str, str] | None = None
    type: Literal["ollama"]


class Gemini(BaseLLM):
    type: Literal["gemini"]


class Bedrock(BaseLLM):
    region: str | None = None
    type: Literal["bedrock"]


ModelUnion = Union[OpenAI, Anthropic, GeminiVertexAI, GeminiAnthropic, Ollama, AzureOpenAI, Gemini, Bedrock]


class ContextCompressionSettings(BaseModel):
    compaction_interval: int
    overlap_size: int
    summarizer_model: ModelUnion | None = Field(default=None, discriminator="type")
    prompt_template: str | None = None
    token_threshold: int | None = None
    event_retention_size: int | None = None


class ContextConfig(BaseModel):
    """Context management configuration containing compaction settings."""

    compaction: ContextCompressionSettings | None = None


class EmbeddingConfig(BaseModel):
    model: str
    provider: str
    base_url: str | None = None


class MemoryConfig(BaseModel):
    """Memory configuration. Its presence signals that memory is enabled."""

    ttl_days: int = 0  # TTL for memory entries in days. 0 means use the server default.
    embedding: EmbeddingConfig | None = None  # Embedding model config for memory tools.


class AgentConfig(BaseModel):
    model: ModelUnion = Field(discriminator="type")
    description: str
    instruction: str
    http_tools: list[HttpMcpServerConfig] | None = None  # Streamable HTTP MCP tools
    sse_tools: list[SseMcpServerConfig] | None = None  # SSE MCP tools
    remote_agents: list[RemoteAgentConfig] | None = None  # remote agents
    execute_code: bool | None = None
    stream: bool | None = None  # Refers to LLM response streaming, not A2A streaming
    memory: MemoryConfig | None = None  # Memory configuration
    context_config: ContextConfig | None = None

    def to_agent(self, name: str, sts_integration: Optional[ADKTokenPropagationPlugin] = None) -> Agent:
        if name is None or not str(name).strip():
            raise ValueError("Agent name must be a non-empty string.")
        tools: list[ToolUnion] = []
        tools_requiring_approval: set[str] = set()
        sts_header_provider = None
        if sts_integration:
            sts_header_provider = sts_integration.header_provider
        if self.http_tools:
            for http_tool in self.http_tools:  # add http tools
                # Create header provider combining STS and allowed headers for this tool
                tool_header_provider = create_header_provider(
                    allowed_headers=http_tool.allowed_headers,
                    sts_header_provider=sts_header_provider,
                )
                tools.append(
                    KAgentMcpToolset(
                        connection_params=http_tool.params,
                        tool_filter=http_tool.tools,
                        header_provider=tool_header_provider,
                    )
                )
                if http_tool.require_approval:
                    tools_requiring_approval.update(http_tool.require_approval)
        if self.sse_tools:
            for sse_tool in self.sse_tools:  # add sse tools
                # Create header provider combining STS and allowed headers for this tool
                tool_header_provider = create_header_provider(
                    allowed_headers=sse_tool.allowed_headers,
                    sts_header_provider=sts_header_provider,
                )
                tools.append(
                    KAgentMcpToolset(
                        connection_params=sse_tool.params,
                        tool_filter=sse_tool.tools,
                        header_provider=tool_header_provider,
                    )
                )
                if sse_tool.require_approval:
                    tools_requiring_approval.update(sse_tool.require_approval)
        if self.remote_agents:
            for remote_agent in self.remote_agents:  # Add remote agents as tools
                # Prepare httpx client parameters
                timeout = httpx.Timeout(timeout=remote_agent.timeout)
                headers: dict[str, str] | None = remote_agent.headers
                base_url: str | None = None
                event_hooks: dict[str, list[Callable[[httpx.Request], None]]] | None = None

                # If headers includes the proxy host header, it means we're using a proxy
                # RemoteA2aAgent may use URLs from agent card response, so we need to
                # rewrite all request URLs to use the proxy URL while preserving the proxy host header
                if remote_agent.headers and PROXY_HOST_HEADER in remote_agent.headers:
                    # Parse the proxy URL to extract base URL
                    from urllib.parse import urlparse as parse_url

                    parsed_proxy = parse_url(remote_agent.url)
                    proxy_base = f"{parsed_proxy.scheme}://{parsed_proxy.netloc}"
                    target_host = remote_agent.headers[PROXY_HOST_HEADER]

                    # Event hook to rewrite request URLs to use proxy while preserving the proxy host header
                    # Note: Relative paths are handled by base_url below, so they'll already point to proxy_base
                    def make_rewrite_url_to_proxy(proxy_base: str, target_host: str) -> Callable[[httpx.Request], None]:
                        async def rewrite_url_to_proxy(request: httpx.Request) -> None:
                            parsed = parse_url(str(request.url))
                            proxy_netloc = parse_url(proxy_base).netloc

                            # If URL is absolute and points to a different host, rewrite to the proxy base URL
                            if parsed.netloc and parsed.netloc != proxy_netloc:
                                # This is an absolute URL pointing to the target service, rewrite it
                                new_url = f"{proxy_base}{parsed.path}"
                                if parsed.query:
                                    new_url += f"?{parsed.query}"
                                request.url = httpx.URL(new_url)

                            # Always set proxy host header for Gateway API routing
                            request.headers[PROXY_HOST_HEADER] = target_host

                        return rewrite_url_to_proxy

                    # Set base_url so relative paths work correctly with httpx
                    # httpx requires either base_url or absolute URLs - relative paths will fail without base_url
                    base_url = proxy_base
                    event_hooks = {"request": [make_rewrite_url_to_proxy(proxy_base, target_host)]}

                # Note: httpx doesn't accept None for base_url/event_hooks, so we only pass the parameters if set
                if base_url and event_hooks:
                    client = httpx.AsyncClient(
                        timeout=timeout,
                        headers=headers,
                        base_url=base_url,
                        event_hooks=event_hooks,
                    )
                elif headers:
                    client = httpx.AsyncClient(
                        timeout=timeout,
                        headers=headers,
                    )
                else:
                    client = httpx.AsyncClient(
                        timeout=timeout,
                    )

                tools.append(
                    KAgentRemoteA2AToolset(
                        name=remote_agent.name,
                        description=remote_agent.description,
                        agent_card_url=f"{remote_agent.url}{AGENT_CARD_WELL_KNOWN_PATH}",
                        httpx_client=client,
                    )
                )

        code_executor = SandboxedLocalCodeExecutor() if self.execute_code else None
        model = _create_llm_from_model_config(self.model)

        # Add built-in ask_user tool unconditionally — every agent can ask the user questions.
        tools.append(AskUserTool())

        # Build before_tool_callback if any tools require approval
        before_tool_callback = make_approval_callback(tools_requiring_approval) if tools_requiring_approval else None

        # static_instruction is sent directly to the model without any placeholder processing
        agent = Agent(
            name=name,
            model=model,
            description=self.description,
            static_instruction=self.instruction,
            tools=tools,
            code_executor=code_executor,
            before_tool_callback=before_tool_callback,
        )

        # Configure memory if enabled
        if self.memory is not None:
            self._configure_memory(agent)

        return agent

    def _configure_memory(self, agent: Agent) -> None:
        """Configures the agent to properly use memory.

        1. Adds memory tools to the agent
        2. Adds a callback to auto-save the session to memory
        3. Adds instructions to the agent to use memory
        """
        try:
            from kagent.adk.tools.memory_tools import LoadMemoryTool, SaveMemoryTool
            from kagent.adk.tools.prefetch_memory_tool import PrefetchMemoryTool

            agent.tools.append(PrefetchMemoryTool())
            agent.tools.append(LoadMemoryTool())
            agent.tools.append(SaveMemoryTool())

            memory_suffix = (
                "\n\nYou have long-term memory: use save_memory to store important findings, learnings, and user preferences. "
                "When you need more context or are unsure, use load_memory to search past conversations for relevant information."
            )
            if agent.static_instruction:
                agent.static_instruction += memory_suffix
            else:
                agent.instruction += memory_suffix

            # Define auto-save callback
            async def auto_save_session_to_memory_callback(callback_context: CallbackContext):
                try:
                    session = callback_context._invocation_context.session
                    # Count user messages from events list
                    user_msg_count = sum(1 for e in session.events if e.author == "user")

                    # Save every 5 turns (skip 0)
                    if user_msg_count > 0 and user_msg_count % 5 == 0:
                        logger.info("Auto-saving session %s to memory (turn %d)", session.id, user_msg_count)

                        # Pass the agent's model directly for summarization
                        await callback_context._invocation_context.memory_service.add_session_to_memory(
                            session,
                            model=agent.model,
                        )
                    else:
                        logger.debug("Skipping auto-save for session %s (turn %d)", session.id, user_msg_count)
                except Exception as e:
                    logger.error("Failed to auto-save session to memory: %s", e)

            # Append to after agent callback list
            if not hasattr(agent, "after_agent_callback") or agent.after_agent_callback is None:
                agent.after_agent_callback = []
            agent.after_agent_callback.append(auto_save_session_to_memory_callback)

        except ImportError as e:
            logger.warning("Failed to import memory tools (google-adk update may be needed): %s", e)
        except Exception as e:
            logger.error("Failed to inject memory configuration: %s", e)


def _create_llm_from_model_config(model_config: ModelUnion):
    extra_headers = model_config.headers or {}
    base_url = getattr(model_config, "base_url", None)

    if model_config.type == "openai":
        return OpenAINative(
            type="openai",
            base_url=base_url,
            default_headers=extra_headers,
            frequency_penalty=model_config.frequency_penalty,
            max_tokens=model_config.max_tokens,
            model=model_config.model,
            n=model_config.n,
            presence_penalty=model_config.presence_penalty,
            reasoning_effort=model_config.reasoning_effort,
            seed=model_config.seed,
            temperature=model_config.temperature,
            timeout=model_config.timeout,
            top_p=model_config.top_p,
            tls_disable_verify=model_config.tls_disable_verify,
            tls_ca_cert_path=model_config.tls_ca_cert_path,
            tls_disable_system_cas=model_config.tls_disable_system_cas,
            api_key_passthrough=model_config.api_key_passthrough,
        )
    if model_config.type == "anthropic":
        return KAgentAnthropicLlm(
            model=model_config.model,
            base_url=base_url,
            extra_headers=extra_headers,
            api_key_passthrough=model_config.api_key_passthrough,
        )
    if model_config.type == "gemini_vertex_ai":
        return GeminiLLM(model=model_config.model)
    if model_config.type == "gemini_anthropic":
        return ClaudeLLM(model=model_config.model)
    if model_config.type == "ollama":
        ollama_options = _convert_ollama_options(getattr(model_config, "options", None))
        # api key passthrough is not applicable for ollama
        return create_ollama_llm(
            model=model_config.model,
            options=ollama_options,
            extra_headers=extra_headers,
        )
    if model_config.type == "azure_openai":
        return OpenAIAzure(
            model=model_config.model,
            type="azure_openai",
            default_headers=extra_headers,
            tls_disable_verify=model_config.tls_disable_verify,
            tls_ca_cert_path=model_config.tls_ca_cert_path,
            tls_disable_system_cas=model_config.tls_disable_system_cas,
            api_key_passthrough=model_config.api_key_passthrough,
        )
    if model_config.type == "gemini":
        return model_config.model
    if model_config.type == "bedrock":
        # api key passthrough is not applicable for bedrock
        return KAgentBedrockLlm(
            model=model_config.model,
            extra_headers=extra_headers,
        )
    raise ValueError(f"Invalid model type: {model_config.type}")


def build_adk_context_configs(
    context_config: ContextConfig,
) -> tuple:
    """Build ADK-native context configuration objects from kagent ContextConfig.

    Returns:
        Tuple of (EventsCompactionConfig | None, None).
    """
    from google.adk.apps.app import EventsCompactionConfig

    events_compaction_config = None

    if context_config.compaction is not None:
        comp = context_config.compaction
        summarizer = None

        if comp.summarizer_model is not None:
            from google.adk.apps.llm_event_summarizer import LlmEventSummarizer

            summarizer = LlmEventSummarizer(
                llm=_create_llm_from_model_config(comp.summarizer_model),
                prompt_template=comp.prompt_template,
            )

        events_compaction_config = EventsCompactionConfig(
            compaction_interval=comp.compaction_interval,
            overlap_size=comp.overlap_size,
            summarizer=summarizer,
        )
        if comp.token_threshold is not None and hasattr(EventsCompactionConfig, "token_threshold"):
            events_compaction_config.token_threshold = comp.token_threshold
        if comp.event_retention_size is not None and hasattr(EventsCompactionConfig, "event_retention_size"):
            events_compaction_config.event_retention_size = comp.event_retention_size

    return events_compaction_config, None
