app.py

import gradio as gr
import logging
from datetime import datetime
from pathlib import Path
from scripts.RepositoryHandler import RepositoryHandler
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "7"

# --- Setup Logging ---
def setup_logger():
    log_dir = Path("logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M")
    log_file = log_dir / f"{timestamp}_code_compass.log"

    logger = logging.getLogger("code_compass")
    logger.setLevel(logging.DEBUG)

    # Console handler
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)

    # File handler
    fh = logging.FileHandler(log_file)
    fh.setLevel(logging.DEBUG)

    # Formatter
    formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
    ch.setFormatter(formatter)
    fh.setFormatter(formatter)

    logger.addHandler(ch)
    logger.addHandler(fh)
    return logger

setup_logger()
logger = logging.getLogger("code_compass")
# Global repository handler instance
logger.info("Checking for model...")
if not os.path.exists("models/Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf"):
    logger.info("Downloading model...")
    Path("models").mkdir(parents=True, exist_ok=True)
    os.system("wget -q https://huggingface.co/bartowski/Qwen2.5-Coder-7B-Instruct-GGUF/resolve/main/Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf -O models/Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf")



repo_handler = RepositoryHandler()


def process_repository(input_type, github_url, zip_file):
    """Process repository based on input type"""
    
    # Clean up any previous repository
    repo_handler.cleanup()
    
    if input_type == "GitHub URL":
        if not github_url or not github_url.strip():
            return "❌ Please enter a GitHub repository URL", "", "disabled", "disabled"
        
        if not repo_handler.validate_github_url(github_url.strip()):
            return "❌ Invalid GitHub URL format. Please use: https://github.com/username/repository", "", "disabled", "disabled"
        
        success, message = repo_handler.download_github_repo(github_url.strip())
        
    else:  # ZIP File
        if zip_file is None:
            return "❌ Please upload a ZIP file", "", "disabled", "disabled"
        
        is_valid, validation_msg = repo_handler.validate_zip_file(zip_file)
        if not is_valid:
            return f"❌ {validation_msg}", "", "disabled", "disabled"
        
        success, message = repo_handler.extract_zip_file(zip_file)
    
    if success:
        structure = repo_handler.get_repo_structure()
        return message, structure, "🚀 Process Repository", "disabled"  # Enable process button, keep query disabled
    else:
        return message, "", "disabled", "disabled"

def process_chunks():
    """Process repository into chunks and store in vector database"""
    if not repo_handler.is_loaded:
        return "❌ No repository loaded", "disabled"
    
    # Run processing in background thread to avoid blocking UI
    def background_processing():
        return repo_handler.process_and_store_chunks()
    
    try:
        success, message = background_processing()
        if success:
            return message, "Ask AI"  # Enable query functionality
        else:
            return message, "disabled"
    except Exception as e:
        return f"❌ Error processing chunks: {str(e)}", "disabled"

def handle_query(query):
    """Handle user queries about the repository"""
    if not repo_handler.is_loaded:
        return "❌ No repository loaded. Please load a repository first."
    
    if not repo_handler.chunks:
        return "❌ Repository not processed yet. Please click 'Process Repository' first."
    
    if not query or not query.strip():
        return "Please enter a query about the repository."
    
    return repo_handler.query_repository(query.strip())

def get_repo_stats():
    """Get repository statistics for display"""
    if not repo_handler.is_loaded:
        return "No repository loaded"
    
    if repo_handler.vector_store and repo_handler.chunks:
        try:
            # Get repository overview from vector store
            overview = repo_handler.vector_store.get_repository_overview(repo_handler.repo_name)
            logger.debug(f"Repository overview: {overview}")
            if "error" not in overview:
                stats = f"""📊 **Repository Statistics**

🏷️ **Repository:** {overview['repo_name']}
📦 **Total Chunks:** {overview['total_chunks']}
📁 **Files:** {overview['files_count']}
🏛️ **Classes:** {overview['classes_count']}  
⚙️ **Functions:** {overview['functions_count']}
💻 **Languages:** {', '.join(overview['languages'])}

📋 **Chunk Distribution:**
"""
                for chunk_type, count in overview['chunk_distribution'].items():
                    stats += f"- {chunk_type.title()}: {count}\n"
                
                return stats
            else:
                return f"Error getting stats: {overview['error']}"
        except Exception as e:
            return f"Error getting repository stats: {str(e)}"
    
    return "Repository loaded but not processed yet"
# Additional handler functions for LLM integration
def initialize_llm():
    """Initialize LLM model loading"""
    return repo_handler.initialize_llm()

def handle_query_with_llm(query, use_llm):
    """Handle user queries with optional LLM processing"""
    if not repo_handler.is_loaded:
        return "❌ No repository loaded. Please load a repository first."
    
    if not repo_handler.chunks:
        return "❌ Repository not processed yet. Please click 'Process Repository' first."
    
    if not query or not query.strip():
        return "Please enter a query about the repository."
    
    return repo_handler.query_repository(query.strip(), use_llm=use_llm)

def clear_conversation():
    """Clear LLM conversation history"""
    if repo_handler.llm:
        repo_handler.llm.clear_conversation()
        return "🗑️ Conversation history cleared!"
    return "❌ LLM not initialized"

def export_conversation():
    """Export conversation history"""
    if repo_handler.llm and repo_handler.llm.is_model_ready():
        conversation = repo_handler.llm.export_conversation()
        if conversation:
            # Format for display
            export_text = "# Conversation Export\n\n"
            for msg in conversation:
                role_emoji = {"system": "⚙️", "user": "👤", "assistant": "🤖"}.get(msg["role"], "💬")
                export_text += f"## {role_emoji} {msg['role'].title()}\n"
                export_text += f"**Time:** {msg['timestamp']}\n\n"
                export_text += f"{msg['content']}\n\n---\n\n"
            return export_text
        else:
            return "No conversation to export"
    return "❌ LLM not ready or no conversation history"

def get_llm_status():
    """Get current LLM status"""
    if not repo_handler.llm_loading_started:
        return "🔄 LLM not initialized"
    elif repo_handler.llm.is_model_ready():
        model_info = repo_handler.llm.get_model_info()
        conversation_summary = repo_handler.llm.get_conversation_summary()
        return f"""✅ **LLM Ready!**
        
**Model:** Qwen2.5-Coder-7B-Instruct (Q4_K_M)
**Context Window:** {model_info['context_window']} tokens
**Temperature:** {model_info['temperature']}
**Status:** {conversation_summary}

🤖 Ready for intelligent code analysis!"""
    else:
        return "⏳ **LLM Loading...** Please wait for model initialization to complete."

def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(title="Code Compass", theme=gr.themes.Soft()) as demo:
        
        gr.Markdown("""
        # 🔍 Code Compass
        
        Upload your repository via GitHub URL or ZIP file, process it with AI-powered chunking, and query your codebase using semantic search!
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                
                # Input section
                with gr.Group():
                    gr.Markdown("### 📥 Repository Input")
                    
                    input_type = gr.Dropdown(
                        choices=["GitHub URL", "ZIP File"], 
                        value="GitHub URL",
                        label="Input Method",
                        info="Choose how you want to provide your repository"
                    )
                    
                    github_url = gr.Textbox(
                        label="GitHub Repository URL",
                        placeholder="https://github.com/username/repository",
                        visible=True
                    )
                    
                    zip_file = gr.File(
                        label="Upload ZIP File",
                        file_types=[".zip"],
                        visible=False
                    )
                    
                    load_btn = gr.Button("📁 Load Repository", variant="primary")
                
                # Processing section
                with gr.Group():
                    gr.Markdown("### ⚙️ Repository Processing")
                    gr.Markdown("After loading, process your repository to enable AI-powered search")
                    
                    process_btn = gr.Button("🚀 Process Repository", interactive=False, variant="secondary")
                    
                # Status section
                with gr.Group():
                    gr.Markdown("### 📊 Status")
                    status_output = gr.Textbox(
                        label="Status",
                        placeholder="Ready to load repository...",
                        interactive=False,
                        lines=3
                    )
            
            with gr.Column(scale=1):
                with gr.Group():
                    gr.Markdown("### 📁 Repository Structure")
                    structure_output = gr.Code(
                        label="Directory Structure",
                        # language="text",
                        interactive=False,
                        lines=10
                    )
                
                with gr.Group():
                    gr.Markdown("### 📊 Repository Stats")
                    stats_output = gr.Markdown(
                        value="Load and process a repository to see statistics"
                    )
                with gr.Group():
                    gr.Markdown("### 🤖 LLM Status")
                    llm_status = gr.Markdown(
                        value="🔄 LLM not initialized"
                    )
                    init_llm_btn = gr.Button("🚀 Initialize LLM", variant="secondary")
        # Query section
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 💬 Query Repository")
                gr.Markdown("Ask questions about your code using natural language. The AI will search through your processed code chunks to find relevant information.")
                
                with gr.Row():
                    query_input = gr.Textbox(
                        label="Ask about your code",
                        placeholder="e.g., 'What does this repository do?', 'Show me authentication functions', 'How is error handling implemented?'",
                        lines=2,
                        scale=4
                    )
                    query_btn = gr.Button("🔍 Ask Question", interactive=False, scale=1)
                    use_llm_toggle = gr.Checkbox(
                            label="Use AI Analysis",
                            value=True,
                            info="Get intelligent responses using LLM"
                        )
                    # Conversation controls
                with gr.Row():
                    clear_chat_btn = gr.Button("🗑️ Clear Chat History", variant="secondary", interactive=False)
                    export_chat_btn = gr.Button("📥 Export Chat", variant="secondary", interactive=False)
                query_output = gr.Markdown(
                    value="Load and process a repository first to start querying...",
                    height=400
                )
        
        # Advanced options (collapsible)
        # with gr.Accordion("🛠️ Advanced Options", open=False):
        #     with gr.Row():
        #         with gr.Column():
        #             gr.Markdown("### 🔧 Pinecone Configuration")
        #             api_key_input = gr.Textbox(
        #                 label="Pinecone API Key",
        #                 placeholder="Enter your Pinecone API key (or set PINECONE_API_KEY env var)",
        #                 type="password"
        #             )
        #             environment_input = gr.Textbox(
        #                 label="Pinecone Environment",
        #                 value="us-west1-gcp-free",
        #                 placeholder="e.g., us-west1-gcp-free"
        #             )
                
        #         with gr.Column():
        #             gr.Markdown("### 📈 Processing Options")
        #             complexity_threshold = gr.Slider(
        #                 minimum=5,
        #                 maximum=50,
        #                 value=20,
        #                 step=5,
        #                 label="Complexity Threshold",
        #                 info="Functions above this complexity will be sub-chunked"
        #             )
        
        # Event handlers
        def toggle_inputs(choice):
            return (
                gr.update(visible=(choice == "GitHub URL")),
                gr.update(visible=(choice == "ZIP File"))
            )
        
        def update_buttons_after_load(status_text):
            # Enable process button if repository is successfully loaded
            is_loaded = "✅" in status_text and "successfully" in status_text.lower()
            return gr.update(interactive=is_loaded)
        
        def update_query_button_after_process(status_text):
            # Enable query button if processing is successful
            is_processed = "✅" in status_text and "complete" in status_text.lower()
            return gr.update(interactive=is_processed)
        
        def update_buttons_after_process(status_text):
            # Enable query button if processing is successful
            is_processed = "✅" in status_text and "complete" in status_text.lower()
            return (
                gr.update(interactive=is_processed),  # query_btn
                gr.update(interactive=is_processed),  # clear_chat_btn  
                gr.update(interactive=is_processed)   # export_chat_btn
            )
        
        def update_llm_status():
            return get_llm_status()
        
        def update_stats(status_output):
            return get_repo_stats(), update_buttons_after_load(status_output), update_query_button_after_process(status_output)
        
        # Wire up the interface
        input_type.change(
            fn=toggle_inputs,
            inputs=[input_type],
            outputs=[github_url, zip_file]
        )
        
        load_btn.click(
            fn=process_repository,
            inputs=[input_type, github_url, zip_file],
            outputs=[status_output, structure_output, process_btn, query_btn]
        ).then(
            fn=update_stats,
            inputs=[status_output],
            outputs=[stats_output, process_btn, query_btn]
        )
        
        process_btn.click(
            fn=process_chunks,
            outputs=[status_output, query_btn]
        ).then(
            fn=update_stats,
            inputs=[status_output],
            outputs=[stats_output, process_btn, query_btn]
        )
        
        # Query handling
        query_btn.click(
            fn=handle_query_with_llm,
            inputs=[query_input, use_llm_toggle],
            outputs=[query_output]
        ).then(
            fn=update_llm_status,
            outputs=[llm_status]
        )
        
        # Chat management
        clear_chat_btn.click(
            fn=clear_conversation,
            outputs=[query_output]
        ).then(
            fn=update_llm_status,
            outputs=[llm_status]
        )
        
        # Allow Enter key to submit query
        query_input.submit(
            fn=handle_query_with_llm,
            inputs=[query_input, use_llm_toggle],
            outputs=[query_output]
        )
         # LLM initialization
        init_llm_btn.click(
            fn=initialize_llm,
            outputs=[llm_status]
        ).then(
            fn=update_llm_status,
            outputs=[llm_status]
        )
        # Add some helpful examples
        gr.Markdown("""
        ### 📝 Example Repositories to Try:
        - `https://github.com/pallets/flask` - Popular Python web framework
        - `https://github.com/requests/requests` - HTTP library for Python
        - `https://github.com/fastapi/fastapi` - Modern Python web framework
        - `https://github.com/psf/black` - Python code formatter
        
        ### 💡 Example Queries:
        - "What is the main purpose of this repository?"
        - "Show me all the authentication functions"
        - "How is error handling implemented?"
        - "What are the main classes and their responsibilities?"  
        - "Find functions that handle file operations"
        - "Show me the configuration management code"
        
        ### ⚙️ Setup Requirements:
        1. **Pinecone API Key**: Get a free API key from [Pinecone.io](https://www.pinecone.io/)
        2. **Environment Variables**: Set `PINECONE_API_KEY` in your environment or enter it in Advanced Options
        3. **Internet Connection**: Required for downloading repositories and accessing Pinecone
        
        ### 🚀 How It Works:
        1. **Load**: Repository is downloaded/extracted and validated
        2. **Process**: Code is analyzed and split into hierarchical chunks (file → class → function → block)
        3. **Store**: Chunks are embedded using AI and stored in Pinecone vector database  
        4. **Query**: Your questions are semantically matched against stored code chunks
        """)
    
    return demo

if __name__ == "__main__":
    # Create and launch the interface
    demo = create_interface()
    
    # Launch with some nice settings
    demo.launch(
        server_name="0.0.0.0",  # Allow external access
        server_port=7860,       # Standard port
        share=False,            # Set to True to create public link
        debug=True              # Enable debug mode for development
    )