import chromadb
from sklearn.cluster import KMeans
import pandas as pd
import json
import ast
from itertools import groupby
import uuid
from datetime import datetime
from pathlib import Path

root_dir = Path(__file__).parent.parent
data_dir = root_dir / "datasets"

embed_faq_df = pd.read_csv(data_dir / "processed" / "bank_faq" / "embed_faq.csv")

# chroma db
client = chromadb.PersistentClient(path="./chroma_db")

# client.delete_collection("bank_faq")
# client.delete_collection("chat_history")

# create collection
collection = client.get_or_create_collection(
    name="bank_faq",
    metadata={"hnsw:space": "cosine"}  # cosine similarity for sentence transformers
)

embed_faq_df["embedding"] = embed_faq_df["embedding"].apply(ast.literal_eval)
# prepare data
# ids        = [str(i) for i in range(len(embed_faq_df))]
ids = [
    f"doc_{row['source_id']}_chunk_{row['chunk_num']}"
    for _, row in embed_faq_df.iterrows()
]
embeddings = embed_faq_df["embedding"].tolist()
documents  = embed_faq_df["text"].tolist()
metadatas  = [
    {
        "doc_id": f"doc_{row['source_id']}_chunk_{row['chunk_num']}",
        "domain": row["domain"],
        "source_id": row["source_id"],
        "chunk_num": row["chunk_num"],
        "question": row["question"]
    }
    for _, row in embed_faq_df.iterrows()
]

# print(ids)
# print(type(embeddings[0]))
# print(documents)
# print(metadatas)

# store
collection.add(
    ids=ids,
    embeddings=embeddings,
    documents=documents,
    metadatas=metadatas
)

print(f"Stored {collection.count()} documents")

faq_collection = client.get_collection(name="bank_faq")
results = faq_collection.get(
    include=["documents", "metadatas"]
)

source_id_to_doc_ids = {}
for doc_id, meta in zip(ids, metadatas):
    source_id = meta["source_id"]
    if source_id not in source_id_to_doc_ids:
        source_id_to_doc_ids[source_id] = []
    source_id_to_doc_ids[source_id].append(doc_id)

# build a dict grouped by source_id
groups = {}
for doc_id, meta, text in zip(results["ids"], results["metadatas"], results["documents"]):
    source_id = meta["source_id"]
    if source_id not in groups:
        groups[source_id] = []
    groups[source_id].append({
        "doc_id":    doc_id,
        "chunk_num": int(meta["chunk_num"]),
        "domain":    meta["domain"],
        "question":  meta["question"],
        "text":      text.split("A:", 1)[1].strip() if "A:" in text else text,
    })

# build viewable
viewable = []
for source_id, chunks in groups.items():
    chunks = sorted(chunks, key=lambda x: x["chunk_num"])  # sort by chunk order
    
    entry = {
        "id":       source_id,
        "domain":   chunks[0]["domain"],
        "question": chunks[0]["question"],
    }

    if len(chunks) == 1:
        entry["answer"] = chunks[0]["text"]
    else:
        entry["answer"] = {"chunks": [{"chunk": c["chunk_num"] + 1, "doc_id":   c["doc_id"], "text": c["text"]} for c in chunks]}

    viewable.append(entry)


with open("chroma_preview.json", "w") as f:
    json.dump(viewable, f, indent=2)

print("Saved to chroma_preview.json")

with open("chroma_preview.json", "r") as f:        # or test.json / val.json
    faq_db = json.load(f)

for item in faq_db:
    source_id = item["id"]
    item["relevant_doc_ids"] = source_id_to_doc_ids.get(source_id, [])

with open("chroma_preview.json", "w") as f:
    json.dump(faq_db, f, indent=2)

# save chats in db
chat_collection = client.get_or_create_collection(name="chat_history")

def create_new_chat():
    # creates a new chat session and returns its ID
    chat_id = str(uuid.uuid4())
    print(f"New chat created: {chat_id}")
    return chat_id

def get_chat_history(chat_id):
    """Retrieve full history of a specific chat."""
    results = chat_collection.get(
        where={"chat_id": chat_id}
    )

    # Pair up queries and answers
    history = []
    for doc, meta in zip(results["documents"], results["metadatas"]):
        history.append({
            "query": doc,
            "answer": meta["answer"],
            "timestamp": meta["timestamp"]
        })

    # Sort by timestamp
    history.sort(key=lambda x: x["timestamp"])
    return history

def load_chats_from_db():
    # rebuild the chats dict from ChromaDB on page load
    try:
        results = chat_collection.get(include=["metadatas", "documents"])
    except Exception:
        return {}

    chats = {}
    for doc, meta in zip(results["documents"], results["metadatas"]):
        chat_id   = meta.get("chat_id")
        chat_name = meta.get("chat_name", "Chat 1")
        timestamp = meta.get("timestamp", "")
        query     = doc
        answer    = meta.get("answer", "")

        if chat_name not in chats:
            chats[chat_name] = {"messages": [], "chat_id": chat_id}

        chats[chat_name]["messages"].append({"role": "user",      "content": query, "timestamp": timestamp})
        chats[chat_name]["messages"].append({"role": "assistant", "content": answer, "timestamp": timestamp})

    # sort messages within each chat by timestamp
    for chat_name in chats:
        chats[chat_name]["messages"].sort(key=lambda x: x.get("timestamp", ""))

    return chats

def delete_chat(chat_id: str):
    """Remove a chat from ChromaDB and the JSON log."""
    # Remove from ChromaDB
    results = chat_collection.get(where={"chat_id": chat_id})
    if results["ids"]:
        chat_collection.delete(ids=results["ids"])

    print(f"Chat deleted: {chat_id}")