import uuid
import time
import requests
from typing import List, Dict, Optional, Any
from dataclasses import dataclass

u/dataclass
class Message:
    content: str
    role: str
    id: str = None
    parent_id: Optional[str] = None
    children_ids: List[str] = None
    timestamp: int = None
    models: List[str] = None

    def __post_init__(self):
        if not self.id:
            self.id = str(uuid.uuid4())
        if self.children_ids is None:
            self.children_ids = []
        if not self.timestamp:
            self.timestamp = int(time.time())
        if not self.models:
            self.models = []

class OpenWebUIClient:
    def __init__(self, base_url: str, model_name: str, auth_token: Optional[str] = None):
        """Initialize the Open WebUI client.

        Args:
            base_url: The base URL of your Open WebUI instance
            model_name: The name of the model to use (e.g., 'llama2:latest')
            auth_token: Optional authentication token for protected instances
        """
        self.base_url = base_url.rstrip('/')
        self.model_name = model_name
        self.session = requests.Session()

        if auth_token:
            self.set_auth_token(auth_token)

    def set_auth_token(self, token: str):
        """Set or update the authentication token.

        Args:
            token: The authentication token to use for requests
        """
        self.session.headers.update({"Authorization": f"Bearer {token}"})

    def create_chat(self) -> str:
        """Create a new chat and return its ID."""
        url = f"{self.base_url}/api/v1/chats/new"
        current_time = int(time.time())

        payload = {
            "chat": {
                "id": "",
                "title": "New Chat",
                "models": [self.model_name],
                "params": {},
                "history": {
                    "messages": {},
                    "currentId": None
                },
                "messages": [],
                "tags": [],
                "timestamp": current_time
            }
        }

        response = self.session.post(url, json=payload)
        response.raise_for_status()
        return response.json()['id']

    def add_message(self, chat_id: str, message: Message) -> Dict:
        """Add a message to a chat and get the assistant's response."""
        # First, update the chat with the user's message
        chat_url = f"{self.base_url}/api/v1/chats/{chat_id}"

        # Prepare the message format
        message_dict = {
            "id": message.id,
            "parentId": message.parent_id,
            "childrenIds": message.children_ids,
            "role": message.role,
            "content": message.content,
            "timestamp": message.timestamp,
            "models": [self.model_name]
        }

        # Update chat with user message
        history_messages = {message.id: message_dict}

        chat_payload = {
            "chat": {
                "models": [self.model_name],
                "messages": [message_dict],
                "history": {
                    "messages": history_messages,
                    "currentId": message.id
                }
            }
        }

        response = self.session.post(chat_url, json=chat_payload)
        response.raise_for_status()

        # If it's a user message, get the assistant's response
        if message.role == "user":
            # Generate assistant response
            completion_url = f"{self.base_url}/api/chat/completions"
            assistant_message_id = str(uuid.uuid4())

            completion_payload = {
                "stream": False,  # Set to True if you want to handle streaming
                "model": self.model_name,
                "messages": [{"role": message.role, "content": message.content}],
                "params": {"num_ctx": 8192},
                "chat_id": chat_id,
                "id": assistant_message_id
            }

            completion_response = self.session.post(completion_url, json=completion_payload)
            completion_response.raise_for_status()
            assistant_content = completion_response.json()['choices'][0]['message']['content']

            # Create assistant message
            assistant_message = Message(
                id=assistant_message_id,
                content=assistant_content,
                role="assistant",
                parent_id=message.id,
                timestamp=int(time.time()),
                models=[self.model_name]
            )

            # Update user message with assistant's message ID
            message_dict["childrenIds"] = [assistant_message_id]

            # Create updated chat payload with both messages
            assistant_dict = {
                "id": assistant_message.id,
                "parentId": assistant_message.parent_id,
                "childrenIds": assistant_message.children_ids,
                "role": assistant_message.role,
                "content": assistant_message.content,
                "timestamp": assistant_message.timestamp,
                "model": self.model_name
            }

            history_messages[assistant_message_id] = assistant_dict

            final_payload = {
                "chat": {
                    "models": [self.model_name],
                    "messages": [message_dict, assistant_dict],
                    "history": {
                        "messages": history_messages,
                        "currentId": assistant_message_id
                    }
                }
            }

            response = self.session.post(chat_url, json=final_payload)
            response.raise_for_status()

            return assistant_dict

        return message_dict

    def get_chat_history(self, chat_id: str) -> Dict[str, Any]:
        """Get the full history of a chat."""
        url = f"{self.base_url}/api/v1/chats/{chat_id}"
        response = self.session.get(url)
        response.raise_for_status()
        return response.json()

    def list_chats(self) -> List[Dict[str, Any]]:
        """Get a list of all chats."""
        url = f"{self.base_url}/api/v1/chats/list"
        response = self.session.get(url)
        response.raise_for_status()
        return response.json()

# Example usage:
def main():
    # Initialize client with auth token
    client = OpenWebUIClient(
        "http://localhost:8080",
        "your-model-here",
        auth_token="sk-123456"  # Replace with your actual token
    )

    # Alternatively, set the token after initialization:
    # client = OpenWebUIClient("http://localhost:3000", "llama2:latest")
    # client.set_auth_token("your_auth_token_here")

    # Create a new chat
    chat_id = client.create_chat()
    print(f"Created chat with ID: {chat_id}")

    # Send a message and get response
    message = Message(
        content="Tell me a joke about programming",
        role="user"
    )

    response = client.add_message(chat_id, message)
    print(f"Assistant's response: {response['content']}")

    # Get chat history
    history = client.get_chat_history(chat_id)
    print("\\nFull chat history:")
    for msg in history['chat']['messages']:
        print(f"{msg['role']}: {msg['content']}")

if __name__ == "__main__":
    main()