components/multi_mode.py

"""
Multi-mode interface components for Codebase Agent.

Provides different interaction modes: Chat, Search, Refactor, Generate
"""

import streamlit as st
from typing import Optional, Dict, Any
import os
from pathlib import Path


def get_workspace_root() -> str:
    """
    Get the workspace root directory for the indexed codebase.
    
    Returns:
        Path to the extracted/processed codebase
    """
    # Check if we have a processed data directory
    data_dir = Path("data")
    if data_dir.exists():
        # Find the extracted folder inside data
        for item in data_dir.iterdir():
            if item.is_dir() and not item.name.startswith('.'):
                return str(item)
    
    # Fallback to data directory itself
    return "data"


def render_mode_selector() -> str:
    """
    Render mode selector and return selected mode.
    
    Returns:
        Selected mode: 'chat', 'search', 'refactor', or 'generate'
    """
    # Mode selector with icons
    mode = st.radio(
        "",
        ["💬 Chat", "🔍 Search", "🔧 Refactor", "✨ Generate"],
        horizontal=True,
        key="mode_selector",
        help="Select interaction mode"
    )
    
    # Map display name to mode key
    mode_map = {
        "💬 Chat": "chat",
        "🔍 Search": "search",
        "🔧 Refactor": "refactor",
        "✨ Generate": "generate"
    }
    
    return mode_map[mode]


def render_chat_mode(chat_engine):
    """
    Render standard chat interface.
    
    Args:
        chat_engine: ChatEngine instance
    """
    st.markdown("### 💬 Chat with Your Codebase")
    
    # Show suggested prompts if no history
    if not st.session_state.get("messages", []):
        st.markdown("💡 **Try asking:**")
        
        # Row 1 - 2 suggestions
        col1, col2 = st.columns(2)
        with col1:
            if st.button("🔍 Explain project structure", key="suggest_0", use_container_width=True):
                st.session_state.pending_prompt = "Explain the project structure and main components"
                st.rerun()
        with col2:
            if st.button("📝 List all functions", key="suggest_1", use_container_width=True):
                st.session_state.pending_prompt = "List all the main functions and their purpose"
                st.rerun()
        
        # Row 2 - 2 suggestions
        col3, col4 = st.columns(2)
        with col3:
            if st.button("⚡ Generate code", key="suggest_2", use_container_width=True):
                st.session_state.pending_prompt = "Generate a new utility function for this project"
                st.rerun()
        with col4:
            if st.button("🔧 Suggest improvements", key="suggest_3", use_container_width=True):
                st.session_state.pending_prompt = "What improvements would you suggest for this code?"
                st.rerun()
    
    # Return True to continue with normal chat flow
    return True


def render_search_mode():
    """
    Render MCP code search interface.
    """
    st.markdown("### 🔍 Search Codebase")
    st.caption("Find patterns across your entire codebase using regex")
    
    # Get workspace root
    workspace = get_workspace_root()
    st.info(f"📁 Searching in: `{workspace}`")
    
    # Search input
    col1, col2 = st.columns([3, 1])
    with col1:
        pattern = st.text_input(
            "Search Pattern",
            placeholder="e.g., class (or def.*login)",
            help="Enter a regex pattern to search for"
        )
    with col2:
        is_regex = st.checkbox("Regex", value=True, help="Use regex pattern matching")
    
    # File pattern filter
    file_pattern = st.text_input(
        "File Pattern",
        value="**/*.py",
        help="Glob pattern for files to search (e.g., **/*.py, src/**/*.js)"
    )
    
    # Context lines
    context_lines = st.slider("Context Lines", 0, 5, 2, help="Number of lines to show before/after match")
    
    # Search button
    if st.button("🔍 Search", type="primary", use_container_width=True):
        if not pattern:
            st.warning("Please enter a search pattern")
            return
        
        with st.spinner("Searching codebase..."):
            try:
                from code_chatbot.mcp.mcp_client import MCPClient
                
                client = MCPClient(workspace_root=workspace)
                results = client.search_code(
                    pattern=pattern,
                    file_pattern=file_pattern,
                    context_lines=context_lines,
                    is_regex=is_regex
                )
                
                if results:
                    st.success(f"✅ Found {len(results)} matches")
                    
                    # Display results
                    for i, result in enumerate(results[:20], 1):  # Limit to 20 results
                        with st.expander(f"📄 {result.file_path}:L{result.line_number}"):
                            # Show context before
                            if result.context_before:
                                st.code("\n".join(result.context_before), language="python")
                            
                            # Highlight matching line
                            st.markdown(f"**→ Line {result.line_number}:**")
                            st.code(result.line_content, language="python")
                            
                            # Show context after
                            if result.context_after:
                                st.code("\n".join(result.context_after), language="python")
                    
                    if len(results) > 20:
                        st.info(f"Showing first 20 of {len(results)} results")
                else:
                    st.info("No matches found. Try a different pattern.")
                    
            except Exception as e:
                st.error(f"Search failed: {e}")
                st.exception(e)


def render_refactor_mode():
    """
    Render MCP refactoring interface.
    """
    st.markdown("### 🔧 Refactor Code")
    st.caption("Perform automated refactorings across your codebase")
    
    # Get workspace root
    workspace = get_workspace_root()
    st.info(f"📁 Refactoring in: `{workspace}`")
    
    # Refactoring type selector
    refactor_type = st.selectbox(
        "Refactoring Type",
        ["Custom Regex", "Common Patterns"],
        help="Choose refactoring approach"
    )
    
    if refactor_type == "Custom Regex":
        # Custom regex refactoring
        col1, col2 = st.columns(2)
        with col1:
            search_pattern = st.text_input(
                "Search Pattern",
                placeholder="e.g., print\\((.*)\\)",
                help="Regex pattern to find"
            )
        with col2:
            replace_pattern = st.text_input(
                "Replace Pattern",
                placeholder="e.g., logger.info(\\1)",
                help="Replacement (supports capture groups like \\1)"
            )
        
        file_pattern = st.text_input(
            "File Pattern",
            value="**/*.py",
            help="Files to process"
        )
        
        dry_run = st.checkbox("Dry Run (Preview Only)", value=True, help="Preview changes without applying")
        
        if st.button("🔧 Refactor", type="primary", use_container_width=True):
            if not search_pattern or not replace_pattern:
                st.warning("Please enter both search and replace patterns")
                return
            
            with st.spinner("Processing refactoring..."):
                try:
                    from code_chatbot.mcp.mcp_client import MCPClient
                    
                    client = MCPClient(workspace_root=workspace)
                    result = client.refactor_code(
                        search_pattern=search_pattern,
                        replace_pattern=replace_pattern,
                        file_pattern=file_pattern,
                        dry_run=dry_run
                    )
                    
                    if result.success:
                        mode_text = "Preview" if dry_run else "Applied"
                        st.success(f"✅ Refactoring {mode_text}: {result.files_changed} files, {result.total_replacements} replacements")
                        
                        # Show changes
                        if result.changes:
                            for change in result.changes[:10]:  # Limit to 10 files
                                with st.expander(f"📄 {change['file_path']} ({change['replacements']} replacements)"):
                                    if change.get('preview'):
                                        st.code(change['preview'], language="diff")
                            
                            if len(result.changes) > 10:
                                st.info(f"Showing first 10 of {len(result.changes)} changed files")
                        else:
                            st.info("No matches found for the given pattern")
                        
                        if dry_run and result.files_changed > 0:
                            st.info("💡 Uncheck 'Dry Run' to apply these changes")
                    else:
                        st.error(f"Refactoring failed: {result.error}")
                        
                except Exception as e:
                    st.error(f"Refactoring failed: {e}")
                    st.exception(e)
    
    else:
        # Common patterns
        st.markdown("#### Common Refactoring Patterns")
        
        common_patterns = {
            "print() → logging": {
                "search": r"print\((.*)\)",
                "replace": r"logger.info(\1)",
                "description": "Replace print statements with logging"
            },
            "assertEqual → assert ==": {
                "search": r"assertEqual\(([^,]+),\s*([^)]+)\)",
                "replace": r"assert \1 == \2",
                "description": "Convert unittest to pytest assertions"
            },
            "Remove trailing whitespace": {
                "search": r"[ \t]+$",
                "replace": "",
                "description": "Clean up trailing whitespace"
            }
        }
        
        pattern_choice = st.selectbox(
            "Select Pattern",
            list(common_patterns.keys())
        )
        
        selected = common_patterns[pattern_choice]
        st.info(selected["description"])
        
        col1, col2 = st.columns(2)
        with col1:
            st.code(f"Search: {selected['search']}", language="regex")
        with col2:
            st.code(f"Replace: {selected['replace']}", language="regex")
        
        dry_run = st.checkbox("Dry Run (Preview Only)", value=True, key="common_dry_run")
        
        if st.button("Apply Refactoring", type="primary", use_container_width=True):
            with st.spinner("Processing..."):
                try:
                    from code_chatbot.mcp.mcp_client import MCPClient
                    
                    client = MCPClient(workspace_root=workspace)
                    result = client.refactor_code(
                        search_pattern=selected["search"],
                        replace_pattern=selected["replace"],
                        file_pattern="**/*.py",
                        dry_run=dry_run
                    )
                    
                    if result.success:
                        st.success(f"✅ {result.files_changed} files, {result.total_replacements} replacements")
                        if result.changes:
                            for change in result.changes[:5]:
                                with st.expander(f"📄 {change['file_path']}"):
                                    st.code(change.get('preview', 'No preview'), language="diff")
                    else:
                        st.error(f"Failed: {result.error}")
                except Exception as e:
                    st.error(f"Failed: {e}")


def render_generate_mode(chat_engine):
    """
    Render code generation interface using ChatEngine.
    
    Args:
        chat_engine: ChatEngine instance
    """
    st.markdown("### ✨ Generate New Features")
    st.caption("Use AI to scaffold complete features from descriptions")
    
    # Feature description
    feature_desc = st.text_area(
        "Describe the feature you want to build",
        placeholder="Example: Create a user authentication system with JWT tokens, login/logout endpoints, password hashing with bcrypt, and session management",
        height=120,
        help="Be as detailed as possible"
    )
    
    # Options
    col1, col2, col3 = st.columns(3)
    with col1:
        include_tests = st.checkbox("Generate Tests", value=True)
    with col2:
        include_docs = st.checkbox("Generate Docs", value=True)
    with col3:
        include_examples = st.checkbox("Include Examples", value=True)
    
    # Framework selection
    framework = st.selectbox(
        "Framework/Stack",
        ["Auto-detect from codebase", "FastAPI", "Flask", "Django", "Express.js", "React", "Vue.js"],
        help="Technology stack for the feature"
    )
    
    if st.button("🚀 Generate Feature", type="primary", use_container_width=True):
        if not feature_desc:
            st.warning("Please describe the feature you want to build")
            return
        
        if not chat_engine:
            st.error("⚠️ Chat engine not initialized. Please index your codebase first.")
            return
        
        with st.spinner("🤖 Generating feature... (this may take 30-60 seconds)"):
            try:
                # Build comprehensive AI Engineer prompt
                prompt = f"""You are a **Senior AI/Software Engineer** with 15+ years of experience building production systems at top tech companies. Your expertise spans system design, security, scalability, and clean code architecture.

## 🎯 MISSION
Analyze the existing codebase and generate a **production-ready, enterprise-grade** implementation for the requested feature.

---

## 📋 FEATURE REQUEST
{feature_desc}

---

## ⚙️ CONFIGURATION
| Setting | Value |
|---------|-------|
| **Framework** | {framework} |
| **Include Tests** | {include_tests} |
| **Include Documentation** | {include_docs} |
| **Include Examples** | {include_examples} |

---

## 🧠 YOUR APPROACH (Follow This Process)

### Phase 1: Architecture Analysis
Before writing code, analyze the existing codebase to understand:
- **Project structure** and conventions
- **Naming patterns** (snake_case, camelCase, etc.)
- **Import style** and module organization
- **Error handling** patterns
- **Logging** approach
- **Configuration** management style

### Phase 2: Design the Solution
- Choose appropriate **design patterns** (Factory, Repository, Service Layer, etc.)
- Plan **database schema** changes if needed
- Define **API contracts** (request/response schemas)
- Consider **edge cases** and error scenarios
- Plan for **scalability** and performance

### Phase 3: Implementation
Generate code that includes:

1. **🏗️ Architecture Overview**
   - High-level system diagram (ASCII or Mermaid)
   - Component relationships and data flow
   
2. **📁 File Structure**
   ```
   feature_name/
   ├── __init__.py
   ├── models.py       # Data models/schemas
   ├── service.py      # Business logic
   ├── routes.py       # API endpoints (if applicable)
   ├── utils.py        # Helper functions
   └── tests/
       ├── test_service.py
       └── test_routes.py
   ```

3. **💻 Complete Code** for each file with:
   - **Type hints** on all functions
   - **Docstrings** with Args, Returns, Raises
   - **Input validation** and sanitization
   - **Error handling** with custom exceptions
   - **Logging** at appropriate levels
   - **Security** considerations (auth, injection prevention, etc.)

4. **🧪 Test Suite** (if enabled):
   - Unit tests with pytest
   - Edge case coverage
   - Mock external dependencies
   - Minimum 80% code coverage target

5. **📖 Documentation** (if enabled):
   - API documentation with examples
   - Usage guide
   - Configuration options

6. **🚀 Integration Guide**:
   - Step-by-step setup instructions
   - Environment variables needed
   - Dependencies to install
   - How to integrate with existing code

---

## 📝 CODE FILE FORMAT
For each file, use this exact format:

### `path/to/filename.py`
```python
\"\"\"
Module docstring explaining purpose.
\"\"\"
# imports here

# code here with full implementation
```

---

## ✅ QUALITY CHECKLIST
Ensure your code:
- [ ] Follows existing codebase conventions
- [ ] Has no hardcoded values (use config/env vars)
- [ ] Handles all error cases gracefully
- [ ] Is thread-safe if applicable
- [ ] Has no security vulnerabilities
- [ ] Is optimized for performance
- [ ] Is maintainable and readable

---

## 🎨 STYLE REQUIREMENTS
- Clean, readable code over clever code
- Self-documenting function/variable names
- Comments for complex logic only
- Consistent formatting with project style
- DRY (Don't Repeat Yourself) principle

Now generate the complete, production-ready implementation:"""
                
                # Use chat engine
                answer, sources = chat_engine.chat(prompt)
                
                st.success("✅ Feature generated!")
                
                # Display generated content
                st.markdown("---")
                st.markdown("#### 📝 Generated Feature")
                st.markdown(answer)
                
                # Show sources if available
                if sources:
                    st.markdown("---")
                    with st.expander("📚 Reference Files Used"):
                        for source in sources:
                            if isinstance(source, dict):
                                st.write(f"- `{source.get('file_path', 'Unknown')}`")
                            else:
                                st.write(f"- `{source}`")
                        
            except Exception as e:
                st.error(f"Generation failed: {e}")
                st.exception(e)


# Export functions
__all__ = [
    'render_mode_selector',
    'render_chat_mode',
    'render_search_mode',
    'render_refactor_mode',
    'render_generate_mode'
]