import chainlit as cl
import aiohttp
import json
import uuid
from typing import Dict, Any
from dotenv import load_dotenv

# Load environment variables from a .env file
load_dotenv()

# API URL for the external AI service
API_URL = "https://startrz-ai-agent.hf.space/api/v1/prediction/0f46681a-ca8f-416d-acc2-22feb80da8c2"

# Load user credentials from a JSON file
def load_user_credentials():
    try:
        with open('user.json', 'r') as file:
            return json.load(file)
    except Exception as e:
        print(f"Error loading user credentials: {e}")
        return {}

USER_CREDENTIALS = load_user_credentials()

# Class to manage chat sessions
class ChatSession:
    def __init__(self, session_id=None):
        self.session_id = session_id or str(uuid.uuid4())
        self.messages = []

    def add_message(self, role: str, content: str):
        self.messages.append({"role": role, "content": content})

    def get_messages(self):
        return self.messages

# Function to query the external API
async def query(payload: Dict[str, Any]) -> Dict[str, Any]:
    async with aiohttp.ClientSession() as session:
        try:
            async with session.post(API_URL, json=payload) as response:
                if response.status == 200:
                    return await response.json()
                else:
                    return {"error": f"API request failed with status {response.status}"}
        except aiohttp.ClientError as e:
            return {"error": str(e)}

# Authentication callback to verify user credentials
@cl.password_auth_callback
def auth_callback(username: str, password: str):
    try:
        stored_password = USER_CREDENTIALS.get(username)
        if stored_password and stored_password == password:
            return cl.User(
                identifier=username,
                metadata={"role": "admin" if username == "admin" else "user", "provider": "credentials"}
            )
    except Exception as e:
        print(f"Authentication error: {e}")

    return None

# Start a new chat session
@cl.on_chat_start
async def start():
    session_id = cl.user_session.get("session_id", str(uuid.uuid4()))
    cl.user_session.set("session_id", session_id)
    chat_session = ChatSession(session_id)
    cl.user_session.set("chat_session", chat_session)

    if chat_session.messages:
        for msg in chat_session.messages:
            await cl.Message(content=msg["content"], author=msg["role"].capitalize()).send()

# Provide predefined starter messages for the chat
@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Morning routine ideation",
            message="Can you help me create a personalized morning routine that would help increase my productivity throughout the day? Start by asking me about my current habits and what activities energize me in the morning.",
            icon="🌅",
        ),
        cl.Starter(
            label="Explain superconductors",
            message="Explain superconductors like I'm five years old.",
            icon="⚡",
        ),
        cl.Starter(
            label="Python script for daily email reports",
            message="Write a script to automate sending daily email reports in Python, and walk me through how I would set it up.",
            icon="🐍",
        ),
        cl.Starter(
            label="Text inviting friend to wedding",
            message="Write a text asking a friend to be my plus-one at a wedding next month. I want to keep it super short and casual, and offer an out.",
            icon="💌",
        )
    ]

# Handle incoming messages from the user
@cl.on_message
async def main(message: cl.Message):
    chat_session = cl.user_session.get("chat_session")
    chat_session.add_message("user", message.content)

    payload = {
        "question": message.content,
        "overrideConfig": {
            "sessionId": chat_session.session_id
        }
    }

    response = await query(payload)

    if "error" in response:
        await cl.Message(content=f"Error: {response['error']}", author="System").send()
        return

    ai_response = response.get("text", "Sorry, I didn't understand that.")

    chat_session.add_message("assistant", ai_response)

    await cl.Message(content=ai_response, author="Assistant").send()

# Handle when the chat is stopped
@cl.on_stop
async def stop():
    chat_session = cl.user_session.get("chat_session")
    if chat_session:
        chat_session.add_message("system", "Chat stopped by user.")

# Handle when the chat is resumed
@cl.on_chat_resume
async def resume():
    session_id = cl.user_session.get("session_id")
    chat_session = ChatSession(session_id)
    cl.user_session.set("chat_session", chat_session)

    if chat_session.messages:
        for msg in chat_session.messages:
            await cl.Message(content=msg["content"], author=msg["role"].capitalize()).send()

# Handle when the chat ends
@cl.on_chat_end
async def end():
    chat_session = cl.user_session.get("chat_session")
    if chat_session:
        chat_session.add_message("system", "Chat ended by user.")
        cl.user_session.delete("session_id")
        cl.user_session.delete("chat_session")