import requests
import io
import re
import numpy as np
import faiss
import time
import gradio as gr
from pypdf import PdfReader
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from accelerate import Accelerator
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

def chunk_text(text, chunk_size=700, overlap_size=150):
    """Chunks text without breaking words in the middle (corrected overlap)."""
    chunks = []
    start = 0
    text_length = len(text)

    while start < text_length:
        end = min(start + chunk_size, text_length)

        # Ensure we do not split words
        if end < text_length and text[end].isalnum():
            last_space = text.rfind(" ", start, end)  # Find last space within the chunk
            if last_space != -1:  # If a space is found, adjust the end
                end = last_space

        chunk = text[start:end].strip()
        if chunk:  # Avoid empty chunks
            chunks.append(chunk)

        if end == text_length:
            break

        # Corrected overlap calculation
        overlap_start = max(0, end - overlap_size)
        if overlap_start < end: # Prevent infinite loop if overlap_start is equal to end.
            last_overlap_space = text.rfind(" ", 0, overlap_start)
            if last_overlap_space != -1 and last_overlap_space > start:
                start = last_overlap_space + 1
            else:
                start = end # If no space found, start at the last end.
        else:
            start = end

    return chunks

chunks = []
for data in all_data:
  chunks.extend(chunk_text(data))

embedding_model = SentenceTransformer("BAAI/bge-large-en")
# embedding_model = SentenceTransformer('multi-qa-mpnet-base-dot-v1')
embeddings = embedding_model.encode(chunks)

index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(embeddings)

def bm25_retrieval(query, documents, top_k=3):
    tokenized_docs = [doc.split() for doc in documents]
    bm25 = BM25Okapi(tokenized_docs)
    return [documents[i] for i in np.argsort(bm25.get_scores(query.split()))[::-1][:top_k]]

def adaptive_retrieval(query, index, chunks, top_k=3, bm25_weight=0.5):
    query_embedding = embedding_model.encode([query], convert_to_numpy=True, dtype=np.float16)
    _, indices = index.search(query_embedding, top_k)
    vector_results = [chunks[i] for i in indices[0]]
    bm25_results = bm25_retrieval(query, chunks, top_k)
    return list(set(vector_results + bm25_results))

def rerank(query, results):
    query_embedding = embedding_model.encode([query], convert_to_numpy=True)
    result_embeddings = embedding_model.encode(results, convert_to_numpy=True)
    similarities = np.dot(result_embeddings, query_embedding.T).flatten()
    return [results[i] for i in np.argsort(similarities)[::-1]], similarities

#Chunk merging.
def merge_chunks(retrieved_chunks, overlap_size=100):
    """Merges overlapping chunks properly by detecting the actual overlap."""
    merged_chunks = []
    buffer = retrieved_chunks[0] if retrieved_chunks else ""

    for i in range(1, len(retrieved_chunks)):
        chunk = retrieved_chunks[i]

        # Find actual overlap
        overlap_start = buffer[-overlap_size:]  # Get the last `overlap_size` chars of the previous chunk
        overlap_index = chunk.find(overlap_start)  # Find where this part appears in the new chunk

        if overlap_index != -1:
            # Merge only the non-overlapping part
            buffer += chunk[overlap_index + overlap_size:]
        else:
            # Store completed merged chunk and start a new one
            merged_chunks.append(buffer)
            buffer = chunk

    if buffer:
        merged_chunks.append(buffer)

    return merged_chunks

# def calculate_confidence(query, context, similarities):
#     return np.mean(similarities)  # Averaged similarity scores
def calculate_confidence(query, answer):
    P, R, F1 = score([answer], [query], lang="en", verbose=False)
    return F1.item()

# Load SLM
accelerator = Accelerator()
accelerator.free_memory()
MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", cache_dir="./my_models")
model = accelerator.prepare(model)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_response(query, context):
    prompt = f"""Your task is to analyze the given Context and answer the Question concisely in plain English. 
    **Guidelines:**
    - Do NOT include </think> tag, just provide the final answer only.
    - Provide a direct, factual answer based strictly on the Context.
    - Avoid generating Python code, solutions, or any irrelevant information.
    Context: {context}
    Question: {query}     
    Answer:
    """
    response = generator(prompt, max_new_tokens=150, num_return_sequences=1)[0]['generated_text']
    answer = response.split("Answer:")[1].strip()
    return answer

def process_query(pdf_urls_text, query):
    pdf_urls = [url.strip() for url in pdf_urls_text.split("\n") if url.strip()]
    if not pdf_urls:
        return "Please enter at least one PDF URL."

    index, chunks = load_and_index_data(pdf_urls)
    retrieved_chunks = adaptive_retrieval(query, index, chunks)
    merged_chunks = merge_chunks(retrieved_chunks, 150)
    reranked_chunks, similarities = rerank(query, merged_chunks)
    context = " ".join(reranked_chunks[:3])
    answer = generate_response(query, context)
    confidence = calculate_confidence(query, answer)
    full_response = f"{answer}\n\nConfidence: {confidence:.2f}"
    return full_response

iface = gr.Interface(
    fn=process_query,
    inputs=[gr.Textbox(lines=3, placeholder="Enter PDF URLs (one per line)"), gr.Textbox(placeholder="Enter your financial question")],
    outputs="text",
    title="Financial Document Q&A Chatbot",
    description="Enter PDF URLs and your question to get answers from the documents."
)

iface.launch()
accelerator.free_memory()