import os
import tempfile
import hashlib
import sqlite3
from datetime import datetime
from typing import List, Tuple, Any

import gradio as gr
from dotenv import load_dotenv

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq

load_dotenv()

CHROMA_DIR = "chroma_db"
DB_FILE = "uploaded_files.db"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GROQ_MODEL = "llama-3.1-8b-instant"
PROMPT_TEMPLATE = (
    "You are a helpful study assistant. Use the provided context to answer "
    "the student's question accurately. If the answer is not in the context, "
    "say that you don't know based on the available materials.\n\n"
    "Context:\n{context}\n\n"
    "Question: {question}"
)

embeddings = None
vectorstore = None
llm = None

# ==========================================
# SQLITE DATABASE FOR FILE TRACKING & CACHE
# ==========================================
def init_db():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS indexed_files (
            file_hash TEXT PRIMARY KEY,
            filename TEXT,
            upload_date TEXT,
            chunk_count INTEGER
        )
    ''')
    conn.commit()
    conn.close()

def get_file_hash(file_path: str) -> str:
    hasher = hashlib.md5()
    with open(file_path, 'rb') as f:
        buf = f.read()
        hasher.update(buf)
    return hasher.hexdigest()

def is_file_indexed(file_hash: str) -> bool:
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute("SELECT 1 FROM indexed_files WHERE file_hash = ?", (file_hash,))
    result = cursor.fetchone()
    conn.close()
    return result is not None

def add_file_to_db(file_hash: str, filename: str, chunk_count: int):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute(
        "INSERT INTO indexed_files (file_hash, filename, upload_date, chunk_count) VALUES (?, ?, ?, ?)",
        (file_hash, filename, datetime.now().strftime("%Y-%m-%d %H:%M:%S"), chunk_count)
    )
    conn.commit()
    conn.close()

def get_indexed_files() -> str:
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute("SELECT filename, upload_date, chunk_count FROM indexed_files ORDER BY upload_date DESC")
    results = cursor.fetchall()
    conn.close()
    
    if not results:
        return "No files indexed yet."
    
    formatted_list = []
    for f, d, c in results:
        formatted_list.append(f"📄 **{f}** (Chunks: {c} | Indexed on: {d})")
    return "\n\n".join(formatted_list)

init_db()

# ==========================================
# LANGCHAIN & RAG LOGIC
# ==========================================
def get_embeddings():
    global embeddings
    if embeddings is None:
        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    return embeddings

def get_vectorstore():
    global vectorstore
    if vectorstore is None:
        emb = get_embeddings()
        vectorstore = Chroma(persist_directory=CHROMA_DIR, embedding_function=emb)
    return vectorstore

def get_llm():
    global llm
    if llm is None:
        api_key = os.environ.get("GROQ_API_KEY")
        if not api_key:
            raise ValueError("GROQ_API_KEY environment variable not set.")
        llm = ChatGroq(model=GROQ_MODEL, groq_api_key=api_key, temperature=0.3)
    return llm

def process_file(file_path: str) -> int:
    ext = os.path.splitext(file_path)[1].lower()
    if ext == ".pdf":
        loader = PyPDFLoader(file_path)
    elif ext in (".txt", ".md"):
        loader = TextLoader(file_path)
    else:
        raise ValueError(f"Unsupported file type: {ext}")
        
    docs = loader.load()
    if not docs:
        return 0
        
    splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
    chunks = splitter.split_documents(docs)
    store = get_vectorstore()
    store.add_documents(chunks)
    return len(chunks)

def index_file(file) -> str:
    if file is None:
        return "No file uploaded."
    try:
        if isinstance(file, str):
            file_path = file
            original_name = os.path.basename(file_path)
        else:
            original_name = os.path.basename(file.name)
            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(original_name)[1]) as tmp:
                tmp.write(file.read())
                file_path = tmp.name
                
        file_hash = get_file_hash(file_path)
        if is_file_indexed(file_hash):
            return f"⚠️ **Skipped:** '{original_name}' is already in the database."

        chunks = process_file(file_path)
        
        if chunks > 0:
            add_file_to_db(file_hash, original_name, chunks)
            return f"✅ **Success:** Indexed {chunks} chunks from '{original_name}'."
        else:
            return f"⚠️ **Warning:** No readable text found in '{original_name}'."
            
    except Exception as e:
        return f"❌ **Error:** {str(e)}"

# ==========================================
# CHAT FUNCTION (FIXED: uses tuple-based history)
# ==========================================
def answer_question(question: str, history: List[Tuple[str, str]]) -> Tuple[str, List[Tuple[str, str]]]:
    """
    Gradio Chatbot (default mode) expects history as a list of (user, bot) tuples.
    We return the updated history with the new pair appended.
    """
    if not question.strip():
        return "", history
        
    store = get_vectorstore()
    llm_model = get_llm()
    docs = store.similarity_search(question, k=4)
    
    if not docs:
        answer = "I couldn't find any relevant information in the uploaded documents. Please upload some study materials first!"
    else:
        context = "\n\n".join(doc.page_content for doc in docs)
        prompt = PROMPT_TEMPLATE.format(context=context, question=question)
        response = llm_model.invoke(prompt)
        answer = response.content
        
    # Append the new exchange as a tuple
    history.append((question, answer))
    return "", history

# ==========================================
# GRADIO UI
# ==========================================
with gr.Blocks(title="Source.AI – RAG Study Assistant") as demo:
    gr.Markdown("# 🎓 Source.AI – RAG Study Assistant")
    gr.Markdown("Upload your study materials (PDF, TXT) and ask questions about them.")

    with gr.Tab("💬 Ask Questions"):
        # IMPORTANT: No type="messages" – using classic tuple‑based history
        chatbot = gr.Chatbot(label="Conversation")
        msg = gr.Textbox(label="Your question", placeholder="Type your doubt here...")
        clear = gr.Button("Clear Chat")
        
        msg.submit(answer_question, inputs=[msg, chatbot], outputs=[msg, chatbot])
        # Clears chatbot (empty list) and the input field
        clear.click(lambda: ([], ""), outputs=[chatbot, msg])

    with gr.Tab("📤 Upload Documents"):
        file_input = gr.File(label="Upload a file", file_types=[".pdf", ".txt", ".md"])
        upload_output = gr.Markdown(label="Status")
        upload_btn = gr.Button("Index Document", variant="primary")
        upload_btn.click(fn=index_file, inputs=[file_input], outputs=[upload_output])

    with gr.Tab("🗄️ Database Info"):
        gr.Markdown("### 📚 Currently Indexed Materials")
        db_list_output = gr.Markdown(get_indexed_files())
        refresh_btn = gr.Button("Refresh Database List")
        refresh_btn.click(fn=get_indexed_files, inputs=[], outputs=[db_list_output])

    gr.Markdown("---\n*Built with Groq, ChromaDB, and Hugging Face Spaces.*")

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(),
        share=False
    )