import chainlit as cl
import pydantic
from datetime import datetime, timezone
from enum import Enum
from typing import Optional
from uuid import UUID, uuid4

class MessageType(Enum):
    system = "system"
    user = "user"
    ai = "assistant"
    tool = "tool"

class Message(pydantic.BaseModel):
    """Domain model for chat messages"""
    id: Optional[UUID] = None
    type: MessageType
    interview_id: UUID
    content: Optional[str] = None
    participant_id: UUID
    timestamp: datetime = datetime.now(timezone.utc)

class State(pydantic.BaseModel):
    messages: list[Message]

@cl.on_chat_start
async def main():

    print("Chat session has started.")

    wsgi_env = cl.context.session.environ
    # print(f"WSGI Environment: {wsgi_env}")
    cookies = wsgi_env.get("HTTP_COOKIE", "No cookies found")
    print(f"Cookies: {cookies}")
    # Parse the cookie string to extract X-Chainlit-Session-id
    session_id = "unknown"
    if cookies != "No cookies found":
        cookie_parts = [part.strip().split("=") for part in cookies.split(";")]
        session_cookie = next((p[1] for p in cookie_parts if len(p) == 2 and p[0] == "X-Chainlit-Session-id"), None)
        if session_cookie:
            session_id = session_cookie
            cl.user_session.set("session_id", session_id)
    print(f"X-Chainlit-Session-id: {session_id}")

    state = State(messages=[])
    cl.user_session.set("state", state)

    await cl.Message(content="Hello! This is a basic Chainlit app.").send()

    # add a message to the state
    state.messages.append(Message(
        id=uuid4(),
        type=MessageType.system,
        interview_id=UUID("12345678-1234-5678-1234-567812345678"),
        content="Hello! This is a basic Chainlit app.",
        participant_id=UUID("87654321-4321-8765-4321-876543218765"),
        timestamp=datetime.now(timezone.utc)
    ))

    actions = [
        cl.Action(name="yes", label="yes", payload={"value": "yes"}),
        cl.Action(name="no", label="no", payload={"value": "no"})
    ]


    response = await cl.AskActionMessage(
        content="yes or no?",
        actions=actions,
        raise_on_timeout=False,
        timeout=3600
    ).send()

    print(f"User response: {response}")

    user_name = await cl.AskUserMessage(content="What is your name?", timeout=3600).send()

    state.messages.append(Message(
        id=uuid4(),
        type=MessageType.user,
        interview_id=UUID("12345678-1234-5678-1234-567812345678"),
        content=user_name['output'] if user_name else None,
        participant_id=UUID("87654321-4321-8765-4321-876543218765"),
        timestamp=datetime.now(timezone.utc)
    ))

    if user_name:

        state.messages.append(Message(
            id=uuid4(),
            type=MessageType.user,
            interview_id=UUID("12345678-1234-5678-1234-567812345678"),
            content=user_name['output'],
            participant_id=UUID("87654321-4321-8765-4321-876543218765"),
            timestamp=datetime.now(timezone.utc)
        ))

        await cl.Message(
            content=f"Your name is: {user_name['output']}.\nChainlit installation is working!\nYou can now start building your own chainlit apps!",
        ).send()

        state.messages.append(Message(
            id=uuid4(),
            type=MessageType.system,
            interview_id=UUID("12345678-1234-5678-1234-567812345678"),
            content=f"Your name is: {user_name['output']}.\nChainlit installation is working!\nYou can now start building your own chainlit apps!",
            participant_id=UUID("87654321-4321-8765-4321-876543218765"),
            timestamp=datetime.now(timezone.utc)
        ))

@cl.on_chat_end
async def on_chat_end():
    print("Chat session has ended.")

    print("Final state:", cl.user_session.get("state"))


@cl.on_chat_resume
async def on_chat_resume():
    print("Chat session has resumed.")

    print("Restored state:", cl.user_session.get("state"))

@cl.on_message
async def on_message(message: cl.Message):
    print("New message received:", message.content)

    state: State = cl.user_session.get("state")

    print("Current state:", state)
    if state is None:
        state = State(messages=[])

    state.messages.append(Message(
        id=uuid4(),
        type=MessageType.user,
        interview_id=UUID("12345678-1234-5678-1234-567812345678"),
        content=message.content,
        participant_id=UUID("87654321-4321-8765-4321-876543218765"),
        timestamp=datetime.now(timezone.utc)
    ))

    await cl.Message(content=f"You said: {message.content}").send()

    state.messages.append(Message(
        id=uuid4(),
        type=MessageType.system,
        interview_id=UUID("12345678-1234-5678-1234-567812345678"),
        content=f"You said: {message}",
        participant_id=UUID("87654321-4321-8765-4321-876543218765"),
        timestamp=datetime.now(timezone.utc)
    ))

    cl.user_session.set("state", state)