chat_database.py

import pickle
import os
from datetime import datetime, timezone

FILE_PATH = "chat_history.pkl"

if not os.path.exists(FILE_PATH):
    with open(FILE_PATH, "wb") as file:
        pickle.dump({}, file)


async def save_context_detail(session_id, name, title, summary, categories):
    try:
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        now = datetime.now(timezone.utc).isoformat()

        if session_id not in data:
            print("Session id not in data")
            data[session_id] = {
                "title": "New Chat",
                "createdAt": now,
                "lastUpdatedAt": now,
                "chat": [],
                "context": [],
                "prompt": "",
            }

        session = data.get(session_id)
        contexts = session.get("context", [])

        contexts.append({"name": name, "title": title,
                        "summary": summary, "categories": categories})

        data[session_id]["lastUpdatedAt"] = now

        with open(FILE_PATH, "wb") as file:
            pickle.dump(data, file)

    except Exception as e:
        print(f"Error saving context entry: {e}")


def clear_context_detail(session_id):
    try:
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        now = datetime.now(timezone.utc).isoformat()

        if session_id not in data:
            print("Session id not in data")
            return False

        data[session_id]["context"] = []

        data[session_id]["lastUpdatedAt"] = now

        with open(FILE_PATH, "wb") as file:
            pickle.dump(data, file)

    except Exception as e:
        print(f"Error saving context entry: {e}")


def save_chat_entry(session_id, role, transcript):
    try:
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        now = datetime.now(timezone.utc).isoformat()

        if session_id not in data:
            print("Session id not in data")
            data[session_id] = {
                "title": "New Chat",
                "createdAt": now,
                "lastUpdatedAt": now,
                "chat": [],
                "context": [],
                "prompt": "",
            }

        messages = data[session_id]["chat"]

        if role == "user":
            messages.append({
                "role": role,
                "transcript": transcript
            })
        else:
            if messages and messages[-1]["role"] == "assistant":
                messages[-1]["transcript"] += " " + transcript
            else:
                messages.append({
                    "role": role,
                    "transcript": transcript
                })

        data[session_id]["lastUpdatedAt"] = now

        with open(FILE_PATH, "wb") as file:
            pickle.dump(data, file)

    except Exception as e:
        print(f"Error saving chat entry: {e}")


def get_chat_history(session_id, limit=15):
    try:
        with open(FILE_PATH, "rb") as file:
            data = pickle.load(file)

        session = data.get(session_id)
        if not session or not isinstance(session, dict):
            return []

        # or "messages" if you’ve standardized on that
        # messages = session.get("chat", [])

        # message_history = []
        # for entry in messages:
        #     role = entry.get('role', '')
        #     transcript = entry.get('transcript', '')
        #     if role and transcript:
        #         message_history.append({"role": role, "content": transcript})

        # return message_history[-15:]

        tail = session.get("chat", [])[-limit:]
        chat_history = [
            {"role": msg["role"], "content": msg["transcript"]}
            for msg in tail
            if msg.get("role") and msg.get("transcript")
        ]

        user_prompt = session.get("prompt", "")
        return chat_history, user_prompt

    except (FileNotFoundError, pickle.UnpicklingError) as e:
        print(f"Error reading or parsing the file: {e}")
        return []
    except Exception as e:
        print(f"Unexpected error: {e}")
        return []


def get_all_chat_details():
    try:
        with open(FILE_PATH, "rb") as file:
            data = pickle.load(file)

        chat_list = []

        for session_id, chat in data.items():
            if not isinstance(chat, dict):
                continue

            messages = []
            for entry in chat.get("chat", []):
                role = entry.get("role", "")
                transcript = entry.get("transcript", "")
                if role and transcript:
                    messages.append({
                        "role": role,
                        "content": transcript
                    })

            chat_list.append({
                "id": session_id,
                "title": chat.get("title", "Untitled"),
                "createdAt": chat.get("createdAt"),
                "lastUpdatedAt": chat.get("lastUpdatedAt"),
                "chat": messages,
                "context": chat.get("context", []),
                "prompt": chat.get("prompt", ""),
            })

        return chat_list

    except (FileNotFoundError, EOFError):
        return []

    except Exception as e:
        print(f"Error reading chats: {e}")
        return []


def create_chat_entry(session_id):
    try:
        # Load existing data or initialize an empty dict
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        now = datetime.now(timezone.utc).isoformat()

        if session_id not in data:
            data[session_id] = {
                "title": "New Chat",
                "createdAt": now,
                "lastUpdatedAt": now,
                "chat": [],
                "context": [],
            }

            # Save the updated data back to file
            with open(FILE_PATH, "wb") as file:
                pickle.dump(data, file)

        return True

    except Exception as e:
        print(f"Error create chat entry : {e}")
        return False


def rename_chat_title(session_id, title):
    try:
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        if session_id not in data:
            return False

        data[session_id]["title"] = title
        data[session_id]["lastUpdatedAt"] = datetime.now(
            timezone.utc).isoformat()

        with open(FILE_PATH, "wb") as file:
            pickle.dump(data, file)

        print(f"Renamed chat: {data[session_id]}")
        return True

    except Exception as e:
        print(f"Error renaming chat title: {e}")
        return False


def save_system_prompt(session_id, prompt):
    try:
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        now = datetime.now(timezone.utc).isoformat()

        if session_id not in data:
            return False

        data[session_id]["prompt"] = prompt
        data[session_id]["lastUpdatedAt"] = now

        with open(FILE_PATH, "wb") as file:
            pickle.dump(data, file)

        print(f"Saved Prompt : {data[session_id]}")
        return True

    except Exception as e:
        print(f"Error saving context entry: {e}")
        return False


def delete_chat(session_id):
    try:
        try:
            with open(FILE_PATH, "rb") as file:
                data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            data = {}

        if session_id not in data:
            return True

        data.pop(session_id)

        with open(FILE_PATH, "wb") as file:
            pickle.dump(data, file)

        if session_id not in data:
            return True

        return False

    except Exception as e:
        print(f"Error deleting chat: {e}")
        return False