from typing import Any, Literal, Sequence, overload
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field

from moxn.base_models.blocks.context import MessageContext
from moxn.models import message as msg
from moxn.models.response import (
    LLMEvent,
    ParsedResponse,
    ResponseParser,
    ResponseTypeDetector,
)
from moxn_types import base
from moxn_types.content import Author, MessageRole, Provider
from moxn_types.request_config import RequestConfig, SchemaDefinition
from moxn_types.type_aliases.anthropic import AnthropicMessage, AnthropicMessagesParam
from moxn_types.type_aliases.google import (
    GoogleGenerateContentResponse,
    GoogleMessagesParam,
)
from moxn_types.type_aliases.openai_chat import (
    OpenAIChatCompletion,
    OpenAIChatMessagesParam,
)
from moxn_types.type_aliases.openai_responses import (
    OpenAIResponse,
    OpenAIResponsesMessagesParam,
)
from moxn_types.type_aliases.provider import (
    ProviderInvocationPayload,
    ProviderMessageParam,
    ProviderPayload,
)
from moxn_types.type_aliases.invocation import (
    AnthropicInvocationParam,
    GoogleInvocationParam,
    OpenAIChatInvocationParam,
    OpenAIResponsesInvocationParam,
)

from .content import PromptContent
from .conversion import MessageConverter
from .core import PromptTemplate
from .prompt_converter import PromptConverter
from .response_handler import ResponseHandler


class PromptSession(BaseModel):
    """Manages the runtime state and operations for a prompt execution."""

    id: UUID = Field(default_factory=uuid4)
    prompt: PromptTemplate
    content: PromptContent
    session_data: base.RenderableModel | None = None
    render_kwargs: dict[str, Any] = {}

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @property
    def prompt_id(self) -> UUID:
        return self.prompt.id

    @property
    def prompt_commit_id(self) -> UUID:
        """Deprecated: Use prompt.commit_id directly"""
        if not self.prompt.commit_id:
            raise ValueError("Prompt must have commit_id for telemetry")
        return self.prompt.commit_id

    @property
    def messages(self) -> list[msg.Message]:
        return self.content.messages

    @classmethod
    def from_prompt_template(
        cls,
        prompt: PromptTemplate,
        session_data: base.RenderableModel | None = None,
        render_kwargs: dict[str, Any] | None = None,
    ) -> "PromptSession":
        """Create a PromptSession from a base Prompt."""
        selected_messages = prompt.get_messages()
        return cls(
            prompt=prompt,
            content=PromptContent(messages=selected_messages),
            session_data=session_data,
            render_kwargs=render_kwargs or {},
        )

    def _create_context_from_session_data(self) -> MessageContext | None:
        """Create a MessageContext from session_data's rendered output.

        Returns:
            MessageContext with variables from session_data.render(), or None if no session_data.
        """
        if not self.session_data:
            return None

        # Render the session data to get variables
        rendered_data = self.session_data.render(**self.render_kwargs)

        # Create context from the rendered variables
        if rendered_data:
            return MessageContext.from_variables(rendered_data)

        return None

    def _normalize_context(
        self, user_context: MessageContext | dict | None
    ) -> MessageContext | None:
        """Normalize user context to MessageContext, merging with session_data if available.

        Args:
            user_context: User-provided context as MessageContext, dict, or None.

        Returns:
            Normalized MessageContext, or None if no context available.
        """
        base_context = self._create_context_from_session_data()

        if base_context:
            # We have session_data context - merge with user context
            if user_context is None:
                return base_context
            elif isinstance(user_context, dict):
                return base_context.merge(MessageContext.from_variables(user_context))
            else:
                return base_context.merge(user_context)
        else:
            # No session_data context - normalize user context only
            if user_context is None:
                return None
            elif isinstance(user_context, dict):
                return MessageContext.from_variables(user_context)
            else:
                return user_context

    def to_message_params(
        self,
        provider: Provider,
        context: MessageContext | dict | None = None,
    ) -> Sequence[ProviderMessageParam]:
        """Convert current state to provider-specific message-level blocks.

        Note: This method returns message-level blocks, not complete API payloads.
        For complete prompt-level payloads, use to_provider_payload() or the
        provider-specific methods (to_anthropic, to_openai_chat, etc.).

        Args:
            provider: The LLM provider to format messages for.
            context: Optional context for variable substitution and provider settings.
                     If not provided but session_data exists, context will be auto-generated
                     from session_data.render(). If provided, it will be merged with
                     session_data context (provided context takes precedence).

        Returns:
            Provider-specific message-level blocks.
        """
        final_context = self._normalize_context(context)
        return MessageConverter.to_message_params(
            self.messages,
            provider,
            final_context,
        )

    def to_provider_payload(
        self,
        provider: Provider,
        context: MessageContext | dict | None = None,
    ) -> ProviderPayload:
        """Convert session to provider-specific API payload.

        This is the generic method for converting to any provider's format.
        It produces the complete TypedDict structure expected by provider SDKs.

        Args:
            provider: The LLM provider to format for.
            context: Optional context for variable substitution.
                     If not provided but session_data exists, variables will be auto-populated
                     from session_data.render(). Provided context takes precedence.

        Returns:
            Provider-specific TypedDict ready for SDK usage.
            - Anthropic: {system: ..., messages: [...]}
            - OpenAI: {messages: [...]}
            - Google: {system_instruction: ..., content: [...]}
        """
        final_context = self._normalize_context(context)
        return PromptConverter.to_provider_payload(
            self.messages,
            provider,
            final_context,
        )

    def parse_provider_response(
        self,
        response: Any,
        provider: Provider,
    ) -> ParsedResponse:
        """Parse a provider response into a normalized format."""
        return ResponseHandler.parse_provider_response(response, provider)

    def create_llm_event_from_parsed_response(
        self,
        parsed_response: ParsedResponse,
        request_config: RequestConfig | None = None,
        schema_definition: SchemaDefinition | None = None,
        attributes: dict[str, Any] | None = None,
        validation_errors: list[str] | None = None,
    ) -> LLMEvent:
        """Create an LLM event from messages and parsed response.

        Args:
            parsed_response: The parsed LLM response
            request_config: Optional provider-specific request configuration
            schema_definition: Optional schema or tool definitions used
            attributes: Optional custom attributes
            validation_errors: Optional validation errors if schema validation failed

        Returns:
            LLMEvent with enhanced telemetry fields
        """
        # Detect response type using the detector
        response_type = ResponseTypeDetector.detect(parsed_response, request_config)

        # Count tool calls for telemetry
        tool_calls_count = ResponseTypeDetector.count_tool_calls(parsed_response)

        return LLMEvent(
            promptId=self.prompt_id,
            promptName=self.prompt.name,
            branchId=self.prompt.branch_id,
            commitId=self.prompt.commit_id,
            messages=[message.model_copy(deep=True) for message in self.messages],
            provider=parsed_response.provider,
            rawResponse=parsed_response.raw_response or {},
            parsedResponse=parsed_response,
            sessionData=self.session_data,
            renderedInput=(
                self.session_data.render(**self.render_kwargs)
                if self.session_data is not None
                else None
            ),
            attributes=attributes,
            isUncommitted=self.prompt.commit_id is None,
            # Enhanced telemetry fields
            responseType=response_type,
            requestConfig=request_config,
            schemaDefinition=schema_definition,
            toolCallsCount=tool_calls_count,
            validationErrors=validation_errors,
        )

    @overload
    def create_llm_event_from_response(
        self,
        response: AnthropicMessage,
        provider: Literal[Provider.ANTHROPIC],
    ) -> LLMEvent: ...

    @overload
    def create_llm_event_from_response(
        self,
        response: OpenAIChatCompletion,
        provider: Literal[Provider.OPENAI_CHAT],
    ) -> LLMEvent: ...

    @overload
    def create_llm_event_from_response(
        self,
        response: OpenAIResponse,
        provider: Literal[Provider.OPENAI_RESPONSES],
    ) -> LLMEvent: ...

    @overload
    def create_llm_event_from_response(
        self,
        response: GoogleGenerateContentResponse,
        provider: Literal[Provider.GOOGLE_GEMINI, Provider.GOOGLE_VERTEX],
    ) -> LLMEvent: ...

    def create_llm_event_from_response(
        self,
        response: (
            AnthropicMessage
            | OpenAIChatCompletion
            | OpenAIResponse
            | GoogleGenerateContentResponse
        ),
        provider: Provider,
    ) -> LLMEvent:
        match provider:
            case Provider.ANTHROPIC:
                parsed_response = ResponseParser.parse(response, Provider.ANTHROPIC)  # type: ignore
            case Provider.OPENAI_CHAT:
                parsed_response = ResponseParser.parse(response, Provider.OPENAI_CHAT)  # type: ignore
            case Provider.OPENAI_RESPONSES:
                parsed_response = ResponseParser.parse(
                    response, Provider.OPENAI_RESPONSES
                )  # type: ignore
            case Provider.GOOGLE_GEMINI:
                parsed_response = ResponseParser.parse(response, Provider.GOOGLE_GEMINI)  # type: ignore
            case Provider.GOOGLE_VERTEX:
                parsed_response = ResponseParser.parse(response, Provider.GOOGLE_VERTEX)  # type: ignore
            case _:
                raise ValueError(f"Unsupported provider: {provider}")

        return self.create_llm_event_from_parsed_response(parsed_response)

    def to_anthropic(
        self, context: MessageContext | dict | None = None
    ) -> AnthropicMessagesParam:
        """Convert session to Anthropic format with system/messages structure.

        Args:
            context: Optional context for variable substitution and provider settings.
                     If not provided but session_data exists, variables will be auto-populated
                     from session_data.render(). Provided context takes precedence.

        Returns:
            AnthropicMessagesParam TypedDict that can be unpacked with ** into the Anthropic SDK.
            The TypedDict has:
            - system: optional system prompt (string or list of blocks)
            - messages: list of user/assistant messages
        """
        final_context = self._normalize_context(context)
        return PromptConverter.to_provider_payload(
            self.messages,
            Provider.ANTHROPIC,
            final_context,
        )

    def to_openai_chat(
        self, context: MessageContext | dict | None = None
    ) -> OpenAIChatMessagesParam:
        """Convert session to OpenAI Chat format with messages array.

        Args:
            context: Optional context for variable substitution and provider settings.
                     If not provided but session_data exists, variables will be auto-populated
                     from session_data.render(). Provided context takes precedence.

        Returns:
            OpenAIChatMessagesParam TypedDict that can be unpacked with ** into the OpenAI SDK.
            The TypedDict has:
            - messages: list of all messages (system, user, assistant)
        """
        final_context = self._normalize_context(context)
        return PromptConverter.to_provider_payload(
            self.messages,
            Provider.OPENAI_CHAT,
            final_context,
        )

    def to_openai_responses(
        self, context: MessageContext | dict | None = None
    ) -> OpenAIResponsesMessagesParam:
        """Convert session to OpenAI Responses API format with input items and instructions.

        Args:
            context: Optional context for variable substitution and provider settings.
                     If not provided but session_data exists, variables will be auto-populated
                     from session_data.render(). Provided context takes precedence.

        Returns:
            OpenAIResponsesMessagesParam TypedDict that can be unpacked with ** into the OpenAI SDK.
            The TypedDict has:
            - input: list of message/tool items
            - instructions: system message text (if any system messages present)
        """
        final_context = self._normalize_context(context)
        return PromptConverter.to_provider_payload(
            self.messages,
            Provider.OPENAI_RESPONSES,
            final_context,
        )

    def to_google_gemini(
        self, context: MessageContext | dict | None = None
    ) -> GoogleMessagesParam:
        """Convert session to Google Gemini format.

        Args:
            context: Optional context for variable substitution and provider settings.
                     If not provided but session_data exists, variables will be auto-populated
                     from session_data.render(). Provided context takes precedence.

        Returns:
            GoogleMessagesParam TypedDict that can be unpacked with ** into the Google SDK.
            The TypedDict has:
            - system_instruction: optional system instruction string
            - content: list of content messages
        """
        final_context = self._normalize_context(context)
        return PromptConverter.to_provider_payload(
            self.messages,
            Provider.GOOGLE_GEMINI,
            final_context,
        )

    def to_google_vertex(
        self, context: MessageContext | dict | None = None
    ) -> GoogleMessagesParam:
        """Convert session to Google Vertex format.

        Args:
            context: Optional context for variable substitution and provider settings.
                     If not provided but session_data exists, variables will be auto-populated
                     from session_data.render(). Provided context takes precedence.

        Returns:
            GoogleMessagesParam TypedDict that can be unpacked with ** into the Google SDK.
            The TypedDict has:
            - system_instruction: optional system instruction string
            - content: list of content messages
        """
        final_context = self._normalize_context(context)
        return PromptConverter.to_provider_payload(
            self.messages,
            Provider.GOOGLE_VERTEX,
            final_context,
        )

    def append_user_text(
        self,
        text: str,
        name: str = "",
        description: str = "",
    ) -> None:
        """Append a user message with text content.

        Args:
            text: The text content to append
            name: Optional name for the message
            description: Optional description for the message
        """
        self.content.append_text(
            text=text,
            name=name,
            description=description,
            role=MessageRole.USER,
        )

    def append_assistant_text(
        self,
        text: str,
        name: str = "",
        description: str = "",
    ) -> None:
        """Append an assistant message with text content.

        Args:
            text: The text content to append
            name: Optional name for the message
            description: Optional description for the message
        """
        self.content.append_text(
            text=text,
            name=name,
            description=description,
            author=Author.MACHINE,
            role=MessageRole.ASSISTANT,
        )

    def append_assistant_response(
        self,
        parsed_response: ParsedResponse,
        candidate_idx: int = 0,
        name: str = "",
        description: str = "",
    ) -> None:
        """Append an assistant response from a parsed LLM response.

        Args:
            parsed_response: The parsed response from an LLM provider
            candidate_idx: Which candidate to use (default 0)
            name: Optional name for the message
            description: Optional description for the message
        """
        self.content.append_parsed_response(
            parsed_response=parsed_response,
            candidate_idx=candidate_idx,
            name=name,
            description=description,
        )

    # -------------------------------------------------------------------------
    # Invocation Methods - Combined message + model parameter payloads
    # -------------------------------------------------------------------------

    def _get_model_params(
        self,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> dict[str, Any]:
        """Build model parameters dict, merging prompt config with overrides.

        Priority: explicit kwargs > prompt.completion_config > omit

        Args:
            model: Model identifier override
            max_tokens: Max tokens override
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            Dictionary with model parameters (only includes non-None values)
        """
        config = self.prompt.completion_config
        params: dict[str, Any] = {}

        # Model
        if model is not None:
            params["model"] = model
        elif config and config.model:
            params["model"] = config.model

        # max_tokens
        if max_tokens is not None:
            params["max_tokens"] = max_tokens
        elif config and config.max_tokens:
            params["max_tokens"] = config.max_tokens

        # temperature
        if temperature is not None:
            params["temperature"] = temperature
        elif config and config.temperature is not None:
            params["temperature"] = config.temperature

        # top_p
        if top_p is not None:
            params["top_p"] = top_p
        elif config and config.top_p is not None:
            params["top_p"] = config.top_p

        return params

    def to_anthropic_invocation(
        self,
        context: MessageContext | dict | None = None,
        *,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> AnthropicInvocationParam:
        """Convert session to complete Anthropic invocation payload.

        Combines messages with model parameters for direct SDK usage:
            response = await anthropic.messages.create(
                **session.to_anthropic_invocation()
            )

        Args:
            context: Optional context for variable substitution.
            model: Model override (e.g., "claude-3-opus-20240229")
            max_tokens: Max tokens override (required by Anthropic)
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            AnthropicInvocationParam ready for SDK unpacking.
        """
        messages_payload = self.to_anthropic(context)
        model_params = self._get_model_params(model, max_tokens, temperature, top_p)
        result: dict[str, Any] = {**messages_payload, **model_params}
        return result  # type: ignore

    def to_openai_chat_invocation(
        self,
        context: MessageContext | dict | None = None,
        *,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> OpenAIChatInvocationParam:
        """Convert session to complete OpenAI Chat invocation payload.

        Combines messages with model parameters for direct SDK usage:
            response = await openai.chat.completions.create(
                **session.to_openai_chat_invocation()
            )

        Args:
            context: Optional context for variable substitution.
            model: Model override (e.g., "gpt-4")
            max_tokens: Max tokens override
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            OpenAIChatInvocationParam ready for SDK unpacking.
        """
        messages_payload = self.to_openai_chat(context)
        model_params = self._get_model_params(model, max_tokens, temperature, top_p)
        result: dict[str, Any] = {**messages_payload, **model_params}
        return result  # type: ignore

    def to_openai_responses_invocation(
        self,
        context: MessageContext | dict | None = None,
        *,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> OpenAIResponsesInvocationParam:
        """Convert session to complete OpenAI Responses API invocation payload.

        Combines input items with model parameters for direct SDK usage:
            response = await openai.responses.create(
                **session.to_openai_responses_invocation()
            )

        Args:
            context: Optional context for variable substitution.
            model: Model override
            max_tokens: Max tokens override
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            OpenAIResponsesInvocationParam ready for SDK unpacking.
        """
        messages_payload = self.to_openai_responses(context)
        model_params = self._get_model_params(model, max_tokens, temperature, top_p)
        result: dict[str, Any] = {**messages_payload, **model_params}
        return result  # type: ignore

    def to_google_gemini_invocation(
        self,
        context: MessageContext | dict | None = None,
        *,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> GoogleInvocationParam:
        """Convert session to complete Google Gemini invocation payload.

        Combines content with model parameters for direct SDK usage.
        Note: Renames max_tokens to max_output_tokens for Google API.

        Args:
            context: Optional context for variable substitution.
            model: Model override (e.g., "gemini-1.5-pro")
            max_tokens: Max tokens override (becomes max_output_tokens)
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            GoogleInvocationParam ready for SDK unpacking.
        """
        messages_payload = self.to_google_gemini(context)
        model_params = self._get_model_params(model, max_tokens, temperature, top_p)
        # Rename max_tokens to max_output_tokens for Google API
        if "max_tokens" in model_params:
            model_params["max_output_tokens"] = model_params.pop("max_tokens")
        result: dict[str, Any] = {**messages_payload, **model_params}
        return result  # type: ignore

    def to_google_vertex_invocation(
        self,
        context: MessageContext | dict | None = None,
        *,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> GoogleInvocationParam:
        """Convert session to complete Google Vertex invocation payload.

        Combines content with model parameters for direct SDK usage.
        Note: Renames max_tokens to max_output_tokens for Google API.

        Args:
            context: Optional context for variable substitution.
            model: Model override
            max_tokens: Max tokens override (becomes max_output_tokens)
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            GoogleInvocationParam ready for SDK unpacking.
        """
        messages_payload = self.to_google_vertex(context)
        model_params = self._get_model_params(model, max_tokens, temperature, top_p)
        # Rename max_tokens to max_output_tokens for Google API
        if "max_tokens" in model_params:
            model_params["max_output_tokens"] = model_params.pop("max_tokens")
        result: dict[str, Any] = {**messages_payload, **model_params}
        return result  # type: ignore

    def to_invocation(
        self,
        context: MessageContext | dict | None = None,
        *,
        provider: Provider | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
    ) -> ProviderInvocationPayload:
        """Convert session to provider-specific invocation payload.

        Generic method that auto-selects the correct provider method based on
        completion_config.provider or explicit provider override.

        Args:
            context: Optional context for variable substitution.
            provider: Provider override (uses completion_config.provider if not specified)
            model: Model override
            max_tokens: Max tokens override
            temperature: Temperature override
            top_p: Top-p override

        Returns:
            Provider-specific invocation payload ready for SDK unpacking.

        Raises:
            ValueError: If no provider available (neither in completion_config nor override)
        """
        effective_provider = provider or (
            self.prompt.completion_config.provider
            if self.prompt.completion_config
            else None
        )
        if not effective_provider:
            raise ValueError(
                "No provider specified: set completion_config.provider or pass provider="
            )

        match effective_provider:
            case Provider.ANTHROPIC:
                return self.to_anthropic_invocation(
                    context,
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
            case Provider.OPENAI_CHAT:
                return self.to_openai_chat_invocation(
                    context,
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
            case Provider.OPENAI_RESPONSES:
                return self.to_openai_responses_invocation(
                    context,
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
            case Provider.GOOGLE_GEMINI:
                return self.to_google_gemini_invocation(
                    context,
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
            case Provider.GOOGLE_VERTEX:
                return self.to_google_vertex_invocation(
                    context,
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
            case _:
                raise ValueError(f"Unsupported provider: {effective_provider}")
