"""Simple module calling an LLM."""

import logging
import os
from collections.abc import Callable
from typing import Any, ClassVar, Literal

import openai
from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
from pydantic import BaseModel, Field

from digitalkin.modules._base_module import BaseModule
from digitalkin.services.services_models import ServicesStrategy

# Configure logging with clear formatting
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


class MessageInputPayload(BaseModel):
    """Message trigger model for the OpenAI Archetype module."""

    payload_type: Literal["message"] = "message"
    user_prompt: str = Field(
        ...,
        title="User Prompt",
        description="The prompt provided by the user for processing.",
    )


class InputFile(BaseModel):
    """File model for the OpenAI Archetype module."""

    name: str = Field(
        ...,
        title="File Name",
        description="The name of the file to be processed.",
    )
    content: bytes = Field(
        ...,
        title="File Content",
        description="The content of the file to be processed.",
    )

    file_type: str = Field(
        ...,
        title="File Type",
        description="The type of the file to be processed.",
    )


class FileInputPayload(BaseModel):
    """File input model for the OpenAI Archetype module."""

    payload_type: Literal["file"] = "file"
    files: list[InputFile] = Field(
        ...,
        title="Files",
        description="List of files to be processed.",
    )


class OpenAIInput(BaseModel):
    """Input model defining what data the module expects."""

    payload: MessageInputPayload | FileInputPayload = Field(
        ...,
        discriminator="payload_type",
        title="Payload",
        description="Either a message or list of file input.",
    )


class MessageOutputPayload(BaseModel):
    """Message output model for the OpenAI Archetype module."""

    payload_type: Literal["message"] = "message"
    user_response: str = Field(
        ...,
        title="User Response",
        description="The response generated by the assistant based on the user prompt.",
    )


class OutputFile(BaseModel):
    """File model for the OpenAI Archetype module."""

    name: str = Field(
        ...,
        title="File Name",
        description="The name of the file to be processed.",
    )
    url: str | None = Field(
        ...,
        title="File URL",
        description="The URL of the file to be processed.",
    )

    message: str | None = Field(
        None,
        title="Message",
        description="Optional message associated with the file.",
    )


class FileOutputPayload(BaseModel):
    """File output model for the OpenAI Archetype module."""

    payload_type: Literal["file"] = "file"
    files: list[OutputFile] = Field(
        ...,
        title="Files",
        description="List of files generated by the assistant.",
    )


class OpenAIOutput(BaseModel):
    """Output model defining what data the module produces."""

    payload: MessageOutputPayload | FileOutputPayload = Field(
        ...,
        discriminator="payload_type",
        title="Payload",
        description="Either a message or file response.",
    )


class OpenAISetup(BaseModel):
    """Setup model defining module configuration parameters."""

    model_name: str = Field(
        ...,
        title="Model Name",
        description="The name of the OpenAI model to use for processing.",
    )
    developer_prompt: str = Field(
        ...,
        title="Developer Prompt",
        description="The developer prompt new versions of system prompt, it defines the behavior of the assistant.",
    )
    temperature: float = Field(
        0.7,
        title="Temperature",
        description="Controls the randomness of the model's output. Higher values make output more random.",
    )
    max_tokens: int = Field(
        100,
        title="Max Tokens",
        description="The maximum number of tokens to generate in the response.",
    )


class OpenAIConfigSetup(BaseModel):
    """Setup model defining module configuration parameters."""

    rag_files: list[bytes] = Field(
        ...,
        title="RAG Files",
        description="Files used for retrieval-augmented generation (RAG) with the OpenAI module.",
    )


class OpenAIToolSecret(BaseModel):
    """Secret model defining module configuration parameters."""


client_config = ClientConfig(
    host="[::]",
    port=50151,
    mode=ServerMode.ASYNC,
    security=SecurityMode.INSECURE,
    credentials=None,
)


class OpenAIToolModule(
    BaseModule[
        OpenAIInput,
        OpenAIOutput,
        OpenAISetup,
        OpenAIToolSecret,
        OpenAIConfigSetup,
    ]
):
    """A openAI endpoint tool module module."""

    name = "OpenAIToolModule"
    description = "A module that interacts with OpenAI API to process text"

    # Define the schema formats for the module
    config_setup_format = OpenAIConfigSetup
    input_format = OpenAIInput
    output_format = OpenAIOutput
    setup_format = OpenAISetup
    secret_format = OpenAIToolSecret

    openai_client: openai.OpenAI

    # Define module metadata for discovery
    metadata: ClassVar[dict[str, Any]] = {
        "name": "OpenAIToolModule",
        "description": "Transforms input text using a streaming LLM response.",
        "version": "1.0.0",
        "tags": ["text", "transformation", "encryption", "streaming"],
    }
    # Define services_config_params with default values
    services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {
        "storage": {
            "config": {"setups": OpenAISetup},
            "client_config": client_config,
        },
        "filesystem": {
            "config": {},
            "client_config": client_config,
        },
        "cost": {
            "config": {},
            "client_config": client_config,
        },
    }

    async def run_config_setup(
        self,
        config_setup_data: OpenAIConfigSetup,
        setup_data: OpenAISetup,
        callback: Callable,
    ) -> None:
        """Configure the module with additional setup data.

        Args:
            config_setup_data: Additional configuration content.
            setup_data: Initial setup data for the module.
            callback: Function to send output data back to the client.
        """
        logger.info("Configuring OpenAIToolModule with additional setup data. %s", config_setup_data)

        # Here you can process config_content and update setup_data as needed
        # For now, we just return the original setup_data
        setup_data.developer_prompt = "| + |".join(f.decode("utf-8") for f in config_setup_data.rag_files)
        await callback(setup_data)

    async def initialize(self, setup_data: OpenAISetup) -> None:
        """Initialize the module capabilities.

        This method is called when the module is loaded by the server.
        Use it to set up module-specific resources or configurations.
        """
        self.client: openai.AsyncOpenAI = openai.AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
        # Define what capabilities this module provides
        self.capabilities = ["text-processing", "streaming", "transformation"]
        logger.info(
            "Module %s initialized with capabilities: %s",
            self.metadata["name"],
            self.capabilities,
        )

    async def run_message(
        self,
        input_model: MessageInputPayload,
        setup_model: OpenAISetup,
        callback: Callable,
    ) -> None:
        """Process input text and stream LLM responses.

        Args:
            input_data: Contains the text to process.
            setup_data: Contains model configuration and development prompt.
            callback: Function to send output data back to the client.

        Raises:
            grpc.RpcError: If gRPC communication fails.
            openai.AuthenticationError: If authentication with OpenAI fails.
            openai.APIConnectionError: If an API connection error occurs.
            Exception: For any unexpected runtime errors.
        """
        # response = await self.client.responses.create(
        #     model=setup_model.model_name,
        #     instructions=setup_model.developer_prompt,
        #     temperature=setup_model.temperature,
        #     max_output_tokens=setup_model.max_tokens,
        #     input=input_model.user_prompt,
        # )
        # logger.info("Recieved answer from OpenAI: %s", response)

        # Get and save the output data
        message_output_payload = MessageOutputPayload(
            payload_type="message",
            user_response="Mock data",
            # user_response=response.output_text,
        )
        output_model = self.output_format.model_validate({"payload": message_output_payload})
        await callback(output_data=output_model)

    async def run(
        self,
        input_data: OpenAIInput,
        setup_data: OpenAISetup,
        callback: Callable,
    ) -> None:
        """Run the module.

        Args:
            input_data: Input data for the module
            setup_data: Setup data for the module
            callback: Callback function to report progress

        Raises:
            ValueError: If the payload type is unknown
        """
        # Validate the input data
        input_model = self.input_format.model_validate(input_data)
        setup_model = self.setup_format.model_validate(setup_data)
        logger.debug("Running with input data: %s", input_model)

        if not hasattr(input_model, "payload"):
            error_msg = "Input data is missing 'payload' field"
            raise ValueError(error_msg)

        if not hasattr(input_model.payload, "payload_type"):
            error_msg = "Input payload is missing 'type' field"
            raise ValueError(error_msg)

        if input_model.payload.payload_type == "message":
            # Validate against MessageInputPayload
            message_payload = MessageInputPayload.model_validate(input_model.payload)
            await self.run_message(message_payload, setup_model, callback)
        elif input_model.payload.payload_type == "file":
            # Validate against FileInputPayload
            file_payload = FileInputPayload.model_validate(input_model.payload)
            await self.run_file(file_payload, setup_model, callback)
        else:
            error_msg = f"Unknown input type '{input_model.payload.payload_type}'. Expected 'message' or 'file'."
            raise ValueError(error_msg)
        logger.info("Job %s completed", self.job_id)

    async def cleanup(self) -> None:
        """Clean up any resources when the module is stopped.

        This method is called when the module is being shut down.
        Use it to close connections, free resources, etc.
        """
        logger.info("Cleaning up module %s", self.metadata["name"])
        # Release any resources here if needed.
