import pickle
import os
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Dict, Any

app = FastAPI()

FILE_PATH = "rag_database.pkl"

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


def save_document_to_entry(session_id, filename, file_path):
    try:
        try:
            with open(file_path, "rb") as file:
                stored_data = pickle.load(file)
        except (FileNotFoundError, EOFError):
            stored_data = {}

        if session_id not in stored_data:
            stored_data[session_id] = {
                "files": []
            }

        stored_data[session_id]["files"].append(filename)

        with open(file_path, "wb") as file:
            pickle.dump(stored_data, file)

    except Exception as e:
        print(f"Error saving file data: {e}")


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

        return data.get(session_id, [])

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


async def save_document():
    save_document_to_file(doc.session_id, doc.data, doc.filename)
    return {"message": "Document saved successfully!"}


async def get_document():
    documents = get_document_from_file(sessionId)
    return {"session_id": sessionId, "documents": documents}