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

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

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

from digitalkin.modules._base_module import BaseModule
from digitalkin.services.services_models import ServicesStrategy
from digitalkin.services.setup.setup_strategy import SetupData

# 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 CPU 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 CPU 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 CPU Archetype module."""

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


class CPUInput(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 CPU 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 CPU 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 CPU Archetype module."""

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


class CPUOutput(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 CPUConfigSetup(BaseModel):
    """Config Setup model definining data that will be pre-computed for each setup and module instance."""

    files: list[str] = Field(
        ...,
        title="Files to embed",
        description="List of files to embed in the setup lifecycle.",
    )


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

    model_name: str = Field(
        ...,
        title="Model Name",
        description="The name of the CPU 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 CPUToolSecret(BaseModel):
    """Secret model defining module configuration parameters."""


server_config = ServerConfig(
    host="[::]",
    port=50151,
    mode=ServerMode.ASYNC,
    security=SecurityMode.INSECURE,
    max_workers=10,
    credentials=None,
)


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


class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret, None]):
    """A CPU endpoint tool module module."""

    name = "CPUIntensiveModule"
    description = "A module that interacts with CPU API to process text"

    # Define the schema formats for the module
    input_format = CPUInput
    output_format = CPUOutput
    setup_format = CPUSetup
    secret_format = CPUToolSecret

    # Define module metadata for discovery
    metadata: ClassVar[dict[str, Any]] = {
        "name": "CPUIntensiveModule",
        "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": {"chat_history": None},
            "client_config": client_config,
        },
        "filesystem": {
            "client_config": client_config,
        },
        "cost": {
            "config": {},
            "client_config": client_config,
        },
    }

    async def initialize(self, setup_data: SetupData) -> 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.
        """

    async def run(
        self,
        input_data: CPUInput,
        setup_data: CPUSetup,
        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)
        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)

        total = 0
        input = MessageInputPayload.model_validate(input_model.payload).user_prompt

        for i in range(int(input)):
            total += i * i
            if i % 100 == 0 or i == int(input) - 1:
                message_output_payload = MessageOutputPayload(
                    payload_type="message",
                    user_response=f"result iteration {i}: {total}",
                )
                output_model = self.output_format.model_validate({"payload": message_output_payload})
                await callback(output_data=output_model)
        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.
