import os
import streamlit as st
from git import Repo
import shutil

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    ServiceContext,
    SentenceSplitter,
)
from llama_index.llms.llama_cpp import LlamaCPP
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

st.set_page_config(page_title="GitHub Repo Explainer", layout="wide")
st.title("📘 GitHub Repository Explainer (100% Free)")

github_url = st.text_input("GitHub URL", placeholder="https://github.com/user/repo")

if st.button("Load and Analyze"):
    if github_url:
        try:
            # Clean previous repo if exists
            if os.path.exists("repo"):
                shutil.rmtree("repo")

            with st.spinner("📥 Cloning GitHub repository..."):
                Repo.clone_from(github_url, "repo")
                st.success("✅ Repo cloned successfully.")

            with st.spinner("🔧 Loading LLM and embeddings..."):
                llm = LlamaCPP(
                    model_path="tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
                    temperature=0.7,
                    max_new_tokens=512,
                    context_window=2048,
                    generate_kwargs={"top_p": 0.95, "top_k": 50},
                    model_kwargs={"n_gpu_layers": 20, "n_batch": 512},
                    verbose=True,
                )

                embed_model = HuggingFaceEmbedding(
                    model_name="sentence-transformers/all-MiniLM-L6-v2"
                )

                service_context = ServiceContext.from_defaults(
                    llm=llm,
                    embed_model=embed_model,
                    node_parser=SentenceSplitter(chunk_size=512, chunk_overlap=50),
                )

            with st.spinner("📄 Reading and parsing files..."):
                docs = SimpleDirectoryReader("repo").load_data()
                st.write(f"📚 {len(docs)} documents loaded.")

            with st.spinner("🔍 Building index..."):
                index = VectorStoreIndex.from_documents(docs, service_context=service_context)
                query_engine = index.as_query_engine()

            with st.spinner("🧠 Querying the model..."):
                query = "Explain the purpose, structure, and setup steps of this GitHub repository."
                response = query_engine.query(query)

            st.subheader("🧾 Repository Summary")
            st.write(str(response))

        except Exception as e:
            st.error(f"❌ Something went wrong:\n\n{e}")
    else:
        st.warning("⚠️ Please enter a GitHub repo URL.")

if st.button("Reset"):
    shutil.rmtree("repo", ignore_errors=True)
    st.experimental_rerun()