# --- IMPORTS ---
import gradio as gr
import os
import re
import requests
import numpy as np
import torch
from sklearn.neighbors import NearestNeighbors
from transformers import AutoTokenizer, AutoModel

# --- CONFIGURATION ---
HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
HF_MODEL = "HuggingFaceH4/zephyr-7b-beta"  # Change if you want
HF_API_URL = f"https://api-inference.huggingface.co/models/{HF_MODEL}"
headers = {"Authorization": f"Bearer {HF_TOKEN}"}

FILES = ["main1.txt", "main2.txt", "main3.txt", "main4.txt", "main5.txt", "main6.txt"]  # Your text files
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  # Light and fast

EMBEDDING_CACHE_FILE = "embeddings.npy"
CHUNKS_CACHE_FILE = "chunks.npy"

# --- FUNCTIONS ---

def load_text_files(file_list):
    knowledge = ""
    for file_name in file_list:
        try:
            with open(file_name, "r", encoding="utf-8") as f:
                knowledge += "\n" + f.read()
        except Exception as e:
            print(f"Error reading {file_name}: {e}")
    return knowledge.strip()

def chunk_text(text, max_chunk_length=500):
    sentences = re.split(r'(?<=[.!?])\s+', text)
    chunks = []
    current_chunk = ""
    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= max_chunk_length:
            current_chunk += " " + sentence
        else:
            chunks.append(current_chunk.strip())
            current_chunk = sentence
    if current_chunk:
        chunks.append(current_chunk.strip())
    return chunks

def embed_texts(texts):
    encoded = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        model_output = model(**encoded)
    embeddings = model_output.last_hidden_state.mean(dim=1)
    return embeddings.cpu().numpy()

def save_cache(embeddings, chunks):
    np.save(EMBEDDING_CACHE_FILE, embeddings)
    np.save(CHUNKS_CACHE_FILE, np.array(chunks))

def load_cache():
    if os.path.exists(EMBEDDING_CACHE_FILE) and os.path.exists(CHUNKS_CACHE_FILE):
        embeddings = np.load(EMBEDDING_CACHE_FILE, allow_pickle=True)
        chunks = np.load(CHUNKS_CACHE_FILE, allow_pickle=True).tolist()
        print("✅ Loaded cached embeddings and chunks.")
        return embeddings, chunks
    return None, None

def retrieve_chunks(query, top_k=5):
    query_embedding = embed_texts([query])
    distances, indices = nn_model.kneighbors(query_embedding, n_neighbors=top_k)
    retrieved = [chunks[i] for i in indices[0]]
    return retrieved

def build_prompt(question):
    relevant_chunks = retrieve_chunks(question)
    context = "\n".join(relevant_chunks)
    system_instruction = """ You are an AI-supported financial expert. You answer questions **exclusively in the context of the "Financial Markets" lecture** at the University of Duisburg-Essen. Your answers are **clear, fact-based, and clearly formulated.**
Observe the following rules:
1. Use the provided lecture excerpts ("lecture_slides") primarily as a source of information.
2. If an answer is **not** covered by the lecture content, you can add to it – but only if you are **absolutely certain**. No hallucinations!
3. If you are unsure, answer politely:
_"Sorry. Unfortunately, I don't know the answer to this question."_
4. If a formula is relevant, **show the exact formula** and explain it in **simple terms.**
5. Avoid vague statements. It's better not to give an answer at all than to give an uncertain one. 
6. Only answer in german! """
    
    prompt = f"""{system_instruction}

Knowledge Base:
{context}

User Question: {question}
Answer:"""
    return prompt

def respond(message, history):
    try:
        prompt = build_prompt(message)

        payload = {
            "inputs": prompt,
            "parameters": {"temperature": 0.2, "max_new_tokens": 400},
        }

        response = requests.post(HF_API_URL, headers=headers, json=payload, timeout=30)
        response.raise_for_status()
        output = response.json()
        generated_text = output[0]["generated_text"]
        answer = generated_text.split("Answer:")[-1].strip()

    except Exception as e:
        print("API Error:", e)
        answer = "❌ Error contacting the model. Please try again later."

    if history is None:
        history = []

    history.append({"role": "assistant", "content": answer})
    return answer

# --- INIT SECTION ---

# Load tokenizer and model for embeddings
tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL)
model = AutoModel.from_pretrained(EMBEDDING_MODEL)

# Try to load cached embeddings and chunks
chunk_embeddings, chunks = load_cache()

if chunk_embeddings is None or chunks is None:
    print("🔄 No cache found. Processing...")
    knowledge_base = load_text_files(FILES)
    chunks = chunk_text(knowledge_base)
    chunk_embeddings = embed_texts(chunks)
    save_cache(chunk_embeddings, chunks)
    print("✅ Embeddings and chunks cached.")

# Build the search model
nn_model = NearestNeighbors(metric="cosine")
nn_model.fit(chunk_embeddings)

# --- GRADIO INTERFACE ---

demo = gr.ChatInterface(
    fn=respond,
    title="📚 Text Knowledge RAG Chatbot",
    description="Ask questions based on the provided text files.",
    chatbot=gr.Chatbot(type="messages"),
)

demo.launch()