import spaces  

import nest_asyncio
nest_asyncio.apply()

import os
import subprocess
import gradio as gr
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.llms.groq import Groq
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.vector_stores.qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams

@spaces.GPU
def zero_gpu_dummy():
    pass

Settings.llm = Groq(model="openai/gpt-oss-20b", api_key=os.environ.get("GROQ_API_KEY"))
Settings.embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-large-en-v1.5",
    device="cpu"
)
Settings.chunk_size = 512
Settings.chunk_overlap = 50

client = QdrantClient(location=":memory:")
client.create_collection(
    collection_name="tech_docs",
    vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
)
vector_store = QdrantVectorStore(client=client, collection_name="tech_docs")
index = VectorStoreIndex.from_vector_store(vector_store=vector_store)

TARGET_REPO = "https://github.com/tiangolo/fastapi.git"
clone_path = "/tmp/fastapi_docs"

if not os.path.exists(clone_path):
    print(f"Cloning {TARGET_REPO}...")
    subprocess.run(["git", "clone", "--depth", "1", TARGET_REPO, clone_path])

print("Indexing repository documentation...")
try:
    documents = SimpleDirectoryReader(
        input_dir=clone_path, 
        required_exts=[".md", ".mdx"], 
        recursive=True
    ).load_data()

    for doc in documents[:200]:
        index.insert(doc)
    print("GitHub repository indexed successfully!")
except Exception as e:
    print(f"Error indexing GitHub repo: {e}")

query_engine = index.as_query_engine(streaming=False)

def chat_with_docs(message, history, file_obj):
    global index, query_engine
    
    if file_obj is not None:
        try:
            pdf_docs = SimpleDirectoryReader(input_files=[file_obj.name]).load_data()
            for doc in pdf_docs:
                index.insert(doc)
            query_engine = index.as_query_engine(streaming=False)
        except Exception as e:
            return f"Error processing uploaded PDF: {str(e)}"
            
    response = query_engine.query(message)
    return str(response)

demo = gr.ChatInterface(
    fn=chat_with_docs,
    title="📚 TechDocs-Knowledge-Engine",
    description="Ask technical questions about the indexed codebase or upload your own technical PDF using the file attachment icon below.",
    additional_inputs=[
        gr.File(label="Upload Technical PDF (Optional)", file_types=[".pdf"])
    ],
    textbox=gr.Textbox(placeholder="Ask a technical question about code or architecture...", container=False, scale=7)
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860, 
        ssr_mode=False, 
        auth=("user", os.environ.get("ACCESS_CODE"))
    )