print("DEBUG: Starting Agent.py")

try:
    from Retrieve import retrieve_chunks_from_multiple_vdbs
    print("DEBUG: Imported Retrieve successfully")
except Exception as e:
    print(f"DEBUG: Error importing Retrieve: {e}")
    raise

try:
    from Ask_llm import query_llm, ask_user_for_query_refinement, query_llm_rewrite_only
    print("DEBUG: Imported Ask_llm successfully")
except Exception as e:
    print(f"DEBUG: Error importing Ask_llm: {e}")
    raise

import hashlib
import re
from config import SIMILARITY_THRESHOLD, MAX_CHUNKS, EXTRA_CHUNKS, MAX_HISTORY_TURNS
#Change the summaryse history function to eliminate the pbag with the history

print("DEBUG: Agent.py imports completed")

def clean_text(text):
    text = text.lower().strip()
    text = re.sub(r'\s+', ' ', text)
    return text

def hash_chunk(text):
    return hashlib.md5(clean_text(text).encode()).hexdigest()

def llm_reconstruct_with_context(user_query, chat_history):
    if not chat_history:
        return user_query.strip()

    history_text = "\n".join([f"User: {q}\nAssistant: {a}" for q, a, _ in chat_history])
    prompt = (
        "Rewrite the new user question below as a complete, standalone version using context from the conversation.\n"
        "Avoid vague references like 'these tests' or 'this result'. Only return the rewritten question. Do not explain, summarize, or suggest anything else.\n\n"
        f"Conversation:\n{history_text}\n\n"
        f"New user question:\n{user_query}"
    )

    return query_llm_rewrite_only(prompt, model="gpt-4.1-mini")

def summarize_history(chat_history):
    history_text = "\n".join([f"User: {q}\nAssistant: {a}" for q, a, _ in chat_history])
    summary_prompt = (
        "Summarize the following conversation history in a few concise sentences to preserve context for follow-up questions:\n\n"
        f"{history_text}\n\nSummary:"
    )
    return query_llm("", summary_prompt)

def iterative_rag_agent(user_query, selected_dbs, chat_history=None):
    if chat_history is None:
        chat_history = []

    # Step 1: Retrieve more chunks than needed											  
    gathered_chunks = []
    gathered_references = []
    seen_chunk_hashes = set()

    corrected_query = llm_reconstruct_with_context(user_query, chat_history)

    new_chunks, new_references = retrieve_chunks_from_multiple_vdbs(
        corrected_query, selected_dbs, SIMILARITY_THRESHOLD, EXTRA_CHUNKS
    )

    for chunk, ref in zip(new_chunks, new_references):
        chunk_hash = hash_chunk(chunk)
        if chunk_hash not in seen_chunk_hashes:
            seen_chunk_hashes.add(chunk_hash)
            gathered_chunks.append(chunk)
            gathered_references.append(ref)
        if len(gathered_chunks) >= MAX_CHUNKS:
            break

    # Step 2: Build history context (with summarization if needed)																  
    if len(chat_history) > MAX_HISTORY_TURNS:
        summary = summarize_history(chat_history[:-MAX_HISTORY_TURNS])
        limited_history = [(f"Summary of previous conversation: {summary}", "","")]
        limited_history += chat_history[-MAX_HISTORY_TURNS:]
    else:
        limited_history = chat_history

    history_text = "\n".join([f"User: {q}\nAssistant: {a}" for q, a, _ in limited_history])

    # Step 3: Query the LLM						   
    context = "\n\n".join(gathered_chunks)
    full_prompt = (
        f"Conversation history:\n{history_text.strip()}\n\n"
        f"Current user query:\n{corrected_query.strip()}\n\n"
        "Instructions:\n"
        "Answer the question using primarily the retrieved context above.\n"
        "If more information is needed, end your answer with:\n"
        '"NEED MORE INFO: YES"\n'
        "and suggest a refined query in the following format:\n"
        '"Suggested Query: <more specific query here>"\n\n'
        "Otherwise, end with: NEED MORE INFO: NO"
    )
    response = query_llm(context, full_prompt)
    # Step 4: Append current turn to history											
    chat_history.append((user_query, response, corrected_query))

    # Step 5: Check for refinement logic										
    if "NEED MORE INFO: YES" in response:
        match = re.search(r"Suggested Query:\s*(.+)", response)
        suggested_query = match.group(1).strip() if match else None
        return {
            "needs_refinement": True,
            "answer": response,
            "references": gathered_references,
            "chunks": gathered_chunks,
            "chat_history": chat_history,
            "corrected_query": corrected_query,
            "suggested_query": suggested_query
        }

    # Step 6: Return result						   
    return {
        "needs_refinement": False,
        "answer": response,
        "references": gathered_references,
        "chunks": gathered_chunks,
        "chat_history": chat_history,
        "corrected_query": corrected_query
    }