# ———————————————————————————————————————————————————————————————————————————>
# Comment: various test-blocks to prevent session-outage with more recent UI;
# chainlit==1.1.404
# new imports: user_session, init_ws_context, chainlit.session/ WebsocketSession
# ref--'Session is disconnected' rt7E1rI_uhficQm8AAAA ..
# <———————————————————————————————————————————————————————————————————————————

import os
import re
import uuid
import json
import asyncio
import requests

from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv

import chainlit as cl

from concurrent.futures import ThreadPoolExecutor

from chainlit import user_session
from chainlit.session import WebsocketSession

from langchain_openai import OpenAI
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from langchain.memory.buffer import ConversationBufferMemory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory

from langchain_core.messages import BaseMessage
from typing import List
from pydantic import BaseModel, Field


# -------------------------------=== class action ===------------------------------
class InMemoryHistory(BaseChatMessageHistory, BaseModel):
    messages: List[BaseMessage] = Field(default_factory=list)

    def add_messages(self, messages: List[BaseMessage]) -> None:
        self.messages.extend(messages)

    def clear(self) -> None:
        self.messages = []

chat_histories = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in chat_histories:
        chat_histories[session_id] = InMemoryHistory()
    return chat_histories[session_id]

# -------------------------------=== environment ===------------------------------
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
auth_token = os.getenv("DAYSOFF_API_TOKEN")
API_URL = "https://aivisions.no/data/daysoff/api/v1/booking/"

# ----------------------------=== system-instruct ===------------------------------
daysoff_assistant_template = """
You are a customer support assistant for Daysoff ('Daysoff Kundeservice AI Support') who helps users retrieve booking information based on their bookingnummer. 
You should concisely use the term ’bookingnummer’. Maintain a friendly and professional tone, **reflecting the warmth of a female customer support 
representative archetype.** By default, you answer in **Norwegian**.
============================
Chat History: {chat_history}
Question: {question}
============================
Answer:
"""

daysoff_assistant_prompt = PromptTemplate(
    input_variables=["chat_history", "question"],
    template=daysoff_assistant_template,
)

# --------------------------------=== API Call ===---------------------------------
async def async_post_request(url, headers, data):
    return await asyncio.to_thread(requests.post, url, headers=headers, json=data)

# ----------------------------=== @cl.set_starters ===------------------------------
@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Booking ID request",
            message="Kan du gi meg info om en reservasjon?",
            icon="/public/booking_id.svg",
            ),
        cl.Starter(
            label="Metric Space Self-Identity Framework",
            message="Explain the Metric Space Self-Identity Framework like I'm six years old.",
            icon="/public/learn.svg",
            ),
        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="/public/terminal.svg",
            ),
        cl.Starter(
            label="Morning routine ideation",
            message="Can you help me create a personalized Yoga/pranayama/meditation 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="/public/idea.svg",
            )
        ]
# ----------------------------=== @cl.on_chat_start ===------------------------------
@cl.on_chat_start
async def setup():
    try:
        cl.user_session.set("socket_auth", True)
        cl.user_session.set("max_retries", 3)
        cl.user_session.set("connection_attempts", 0)

        llm = OpenAI(
            model="gpt-3.5-turbo-instruct",
            temperature=0.7,
            openai_api_key=OPENAI_API_KEY,
            max_tokens=2048,
            top_p=0.9,
            frequency_penalty=0.1,
            presence_penalty=0.1,
        )

        conversation_memory = ConversationBufferMemory(
            memory_key="chat_history",
            max_len=30,
            return_messages=True
        )

        llm_chain = LLMChain(
            llm=llm,
            prompt=daysoff_assistant_prompt,
            memory=conversation_memory,
        )

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

    except Exception as e:
        await cl.Message(content=f"Errorisme in init chat session: {str(e)}").send()

# ----------------------------=== long_running_task ===------------------------------
async def long_running_task(message_content: str):
    loop = asyncio.get_running_loop()
    with ThreadPoolExecutor() as pool:
        llm_chain = cl.user_session.get("llm_chain")
        if llm_chain:
            return await loop.run_in_executor(
                pool,
                lambda: llm_chain.invoke({
                    "question": message_content,
                    "chat_history": ""
                })
            )
        else:
            return {"text": "Errorism: LLM Chain is not init."}

# ----------------------------=== @cl.on_message ===------------------------------
@cl.on_message
async def handle_message(message: cl.Message):
    user_message = message.content
    llm_chain = cl.user_session.get("llm_chain")

    if not llm_chain:
        await cl.Message(content="Errorism: Chat session not properly init. Consider a restart.").send()
        return

    booking_pattern = r'\b[A-Z]{6}\d{6}\b'
    match = re.search(booking_pattern, user_message)

    if match:
        bestillingskode = match.group()
        headers = {
            "Authorization": auth_token,
            "Content-Type": "application/json"
        }
        payload = {"booking_id": bestillingskode}

        try:
            response = await async_post_request(API_URL, headers, payload)
            response.raise_for_status()
            booking_data = response.json()

            #if "order_number" in booking_data.get("data", {}):
            if "booking_id" in booking_data:
                try:
                    data = booking_data["data"]
                    
                    table = (
                        "| 𝑭𝒊𝒆𝒍𝒅      | 𝗜𝗻𝗳𝗼                |\n"
                        "|:-----------|:---------------------|\n"
                        f"| 𝙱𝚎𝚜𝚝𝚒𝚕𝚕𝚒𝚗𝚐𝚜𝚔𝚘𝚍𝚎 | {booking_data.get('booking_id', 'N/A')} |\n"
                        f"| 𝙁𝙪𝙡𝙡 𝙉𝙖𝙢𝙚  | {booking_data.get('full_name', 'N/A')} |\n"
                        f"| 𝘼𝙢𝙤𝙪𝙣𝙩     | {booking_data.get('amount', 0)} kr |\n"
                        f"| 𝘾𝙝𝙚𝙘𝙠-𝙞𝙣   | {booking_data.get('checkin', 'N/A')} |\n"
                        f"| 𝘾𝙝𝙚𝙘𝙠-𝙤𝙪𝙩  | {booking_data.get('checkout', 'N/A')} |\n"
                        f"| 𝘼𝙙𝙙𝙧𝙚𝙨𝙨    | {booking_data.get('address', 'N/A')} |\n"
                        f"| 𝙐𝙨𝙚𝙧 𝙄𝘿    | {booking_data.get('user_id', 0)} |\n"
                        f"| 𝙄𝙣𝙛𝙤 𝙏𝙚𝙭𝙩  | {booking_data.get('infotext', 'N/A')} |\n"
                        f"| 𝙄𝙣𝙘𝙡𝙪𝙙𝙚𝙙   | {booking_data.get('included', 'N/A')} |"
                    )

                    combined_message = f"### Her er informasjon for {bestillingskode}:\n\n{table}"
                    await cl.Message(content=combined_message).send()
                    

                except Exception as e:
                    await cl.Message(content=f"En uventet feil oppsto: {str(e)}").send()
            else:
                await cl.Message(content="Ingen bookinginformasjon funnet.").send()

        except requests.exceptions.RequestException as e:
            await cl.Message(content=f"API tilkoblingen ble ikke utført: {str(e)}").send()

    else:
        try:
            response = await long_running_task(message.content)
            await cl.Message(content=response["text"]).send()
            
        except Exception as e:
            await cl.Message(content=f"Errorism!: {str(e)}").send()