from pydantic import BaseModel, Field  # type: ignore
from typing import Dict, List, Optional, Any, Union, Mapping
from datetime import datetime, timezone
from . rag import ChromaDBGetResponse
from ollama._types import Options # type: ignore

class Tunables(BaseModel):
    enable_rag: bool = True      # Enable RAG collection chromadb matching
    enable_tools: bool = True    # Enable LLM to use tools
    enable_context: bool = True  # Add <|context|> field to message

class MessageMetaData(BaseModel):
    rag: List[ChromaDBGetResponse] = Field(default_factory=list)
    eval_count: int = 0
    eval_duration: int = 0
    prompt_eval_count: int = 0
    prompt_eval_duration: int = 0
    context_size: int = 0
    resume_generation: Optional[Dict[str, Any]] = None
    options: Optional[Union[Mapping[str, Any], Options]] = None
    tools: Optional[Dict[str, Any]] = None
    timers: Optional[Dict[str, float]] = None

    #resume : str = ""
    #match_stats: Optional[Dict[str, Dict[str, Any]]] = Field(default=None)

class Message(BaseModel):
    model_config = {"arbitrary_types_allowed": True}  # Allow Event
    # Required
    prompt: str  # Query to be answered

    # Tunables
    tunables: Tunables = Field(default_factory=Tunables)

    # Generated while processing message
    status: str = ""  # Status of the message
    preamble: Dict[str, Any] = Field(default_factory=dict)  # Preamble to be prepended to the prompt
    system_prompt: str = ""  # System prompt provided to the LLM
    context_prompt: str = ""  # Full content of the message (preamble + prompt)
    response: str = ""  # LLM response to the preamble + query
    metadata: MessageMetaData = Field(default_factory=MessageMetaData)
    network_packets: int = 0  # Total number of streaming packets
    network_bytes: int = 0  # Total bytes sent while streaming packets
    actions: List[str] = (
        []
    )  # Other session modifying actions performed while processing the message
    timestamp: str = str(datetime.now(timezone.utc))
    chunk: str = ""
    partial_response: str = ""
    title: str = ""
    context_size: int = 0

    def add_action(self, action: str | list[str]) -> None:
        """Add a actions(s) to the message."""
        if isinstance(action, str):
            self.actions.append(action)
        else:
            self.actions.extend(action)

    def get_summary(self) -> str:
        """Return a summary of the message."""
        response_summary = (
            f"Response: {self.response} (Actions: {', '.join(self.actions)})"
            if self.response
            else "No response yet"
        )
        return (
            f"Message at {self.timestamp}:\n"
            f"Query: {self.preamble}{self.content}\n"
            f"{response_summary}"
        )
