import os
import re
import json
import shutil
from typing import List, Dict, Any
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document
import time
RELEVANT_EXTENSIONS = {'.java', '.kt', '.xml', '.md', '.gradle', '.properties'}

class CodebaseRAGAssistant:
    def __init__(self, model_name="codellama:7b", embedding_model="nomic-embed-text", temperature=0.1):
        self.model_name = model_name
        self.embedding_model = embedding_model
        self.temperature = temperature
        
        # Initialize models
        self.embeddings = OllamaEmbeddings(model=embedding_model)
        self.llm = OllamaLLM(model=model_name, temperature=temperature)
        self.vectordb = None
        
        # Problem category mapping
        self.problem_categories = {
            'authentication': ['auth', 'login', 'sign', 'token', 'password', 'verify', 'otp', 'phone'],
            'ui': ['layout', 'view', 'button', 'text', 'image', 'color', 'theme', 'style', 'responsive', 'display', 'screen', 'ui'],
            'network': ['api', 'http', 'request', 'response', 'retrofit', 'okhttp', 'internet', 'parsing', 'json'],
            'database': ['room', 'sqlite', 'database', 'dao', 'entity', 'query'],
            'permissions': ['permission', 'camera', 'location', 'storage', 'manifest'],
            'navigation': ['fragment', 'activity', 'intent', 'navigation', 'back', 'transition'],
            'performance': ['memory', 'leak', 'slow', 'crash', 'anr', 'optimize'],
            'recyclerview': ['recyclerview', 'adapter', 'viewholder', 'list', 'grid']
        }

    def ingest_codebase(self, root_dir: str):
        """Recursively collect all relevant code files from the codebase."""
        code_files = []
        for dirpath, _, filenames in os.walk(root_dir):
            for fname in filenames:
                ext = os.path.splitext(fname)[1].lower()
                if ext in RELEVANT_EXTENSIONS:
                    code_files.append(os.path.join(dirpath, fname))
        return code_files

    def clean_codebase(self, file_list: List[str]) -> List[str]:
        """Filter out unwanted files but keep test files for this RAG system."""
        cleaned_files = []
        for f in file_list:
            # Skip build directories but keep test files
            if 'build' in f and 'gradle' not in f:
                continue
            # Skip .idea files except important ones
            if '.idea' in f and not any(x in f for x in ['compiler.xml', 'misc.xml']):
                continue
            cleaned_files.append(f)
        return cleaned_files

    def extract_dependencies(self, content: str, file_path: str, file_type: str) -> Dict[str, Any]:
        """Extract dependencies based on file type."""
        dependencies = {
            'imports': [],
            'layout_references': [],
            'findViewById_calls': [],
            'class_references': [],
            'view_ids': [],
            'test_methods': [],
            'api_calls': [],
            'permissions': []
        }
        
        if file_type in ['java', 'kt']:
            # Extract imports
            dependencies['imports'] = re.findall(r'import\s+([^;]+);?', content)
            
            # Extract layout references
            dependencies['layout_references'] = list(set(re.findall(r'R\.layout\.(\w+)', content)))
            
            # Extract findViewById calls
            dependencies['findViewById_calls'] = list(set(re.findall(r'findViewById\(R\.id\.(\w+)\)', content)))
            
            # Extract class references
            class_refs = re.findall(r'new\s+(\w+)\(|(\w+)\.class', content)
            dependencies['class_references'] = [ref for group in class_refs for ref in group if ref and ref[0].isupper()]
            
            # Extract test methods (for test files)
            dependencies['test_methods'] = re.findall(r'@Test[^{]*?(?:public|private)\s+\w+\s+(\w+)\s*\(', content)
            
            # Extract API calls
            dependencies['api_calls'] = re.findall(r'\.(\w*api\w*|\w*request\w*|\w*call\w*)\(', content, re.IGNORECASE)
            
        elif file_type == 'xml':
            # Extract view IDs
            dependencies['view_ids'] = list(set(re.findall(r'android:id="@\+?id/([^"]+)"', content)))
            
            # Extract permissions (for manifest files)
            if 'manifest' in file_path.lower():
                dependencies['permissions'] = re.findall(r'<uses-permission[^>]*name="([^"]+)"', content)
        
        return dependencies

    def chunk_file_enhanced(self, file_path: str) -> List[tuple]:
        """Enhanced chunking that handles test files and creates better context."""
        ext = os.path.splitext(file_path)[1].lower()
        file_type = ext[1:] if ext.startswith('.') else ext
        
        try:
            with open(file_path, encoding="utf-8", errors="ignore") as f:
                content = f.read()
        except Exception as e:
            print(f"Error reading {file_path}: {e}")
            return []

        chunks = []
        dependencies = self.extract_dependencies(content, file_path, file_type)
        
        base_metadata = {
            "file": file_path,
            "file_type": file_type,
            "is_test": 'test' in file_path.lower(),
            "dependencies": dependencies
        }

        if file_type in ['java', 'kt']:
            chunks.extend(self._chunk_java_kotlin(content, file_path, base_metadata))
        elif file_type == 'xml':
            chunks.extend(self._chunk_xml(content, file_path, base_metadata))
        elif file_type == 'gradle':
            chunks.extend(self._chunk_gradle(content, file_path, base_metadata))
        else:
            # Default chunking for other files
            lines = content.splitlines()
            for i in range(0, len(lines), 50):
                chunk = "\n".join(lines[i:i+50])
                if chunk.strip():
                    chunks.append((chunk.strip(), {**base_metadata, "chunk_type": "file_section", "lines": f"{i}-{i+50}"}))
        
        return chunks

    def _chunk_java_kotlin(self, content: str, file_path: str, base_metadata: Dict) -> List[tuple]:
        """Chunk Java/Kotlin files by classes, methods, and test methods."""
        chunks = []
        
        # Find classes
        class_pattern = r'(?:public\s+)?(?:abstract\s+)?(?:class|interface)\s+(\w+)[^{]*\{'
        class_matches = list(re.finditer(class_pattern, content))
        
        if class_matches:
            for match in class_matches:
                class_start = match.start()
                class_name = match.group(1)
                
                # Find class end by counting braces
                brace_count = 0
                class_end = class_start
                for i, char in enumerate(content[class_start:]):
                    if char == '{':
                        brace_count += 1
                    elif char == '}':
                        brace_count -= 1
                        if brace_count == 0:
                            class_end = class_start + i + 1
                            break
                
                class_content = content[class_start:class_end]
                
                # Create enhanced metadata
                metadata = {
                    **base_metadata,
                    "chunk_type": "class",
                    "class_name": class_name,
                    "start": class_start,
                    "end": class_end
                }
                
                chunks.append((class_content, metadata))
                
                # Also chunk individual methods within the class
                method_chunks = self._extract_methods(class_content, class_name, base_metadata)
                chunks.extend(method_chunks)
        
        # If no classes found, treat as utility file
        if not chunks:
            chunks.append((content, {**base_metadata, "chunk_type": "utility_file"}))
        
        return chunks

    def _extract_methods(self, class_content: str, class_name: str, base_metadata: Dict) -> List[tuple]:
        """Extract individual methods from a class."""
        chunks = []
        
        # Enhanced method pattern including test methods
        method_pattern = r'(@Test\s+)?(@\w+\s+)*(?:public|private|protected)\s+(?:static\s+)?(?:\w+\s+)*(\w+)\s*\([^{]*\)\s*\{'
        method_matches = list(re.finditer(method_pattern, class_content, re.MULTILINE))
        
        for match in method_matches:
            method_start = match.start()
            method_name = match.group(3)
            is_test = match.group(1) is not None
            
            # Find method end by counting braces
            brace_count = 0
            method_end = method_start
            for i, char in enumerate(class_content[method_start:]):
                if char == '{':
                    brace_count += 1
                elif char == '}':
                    brace_count -= 1
                    if brace_count == 0:
                        method_end = method_start + i + 1
                        break
            
            method_content = class_content[method_start:method_end]
            
            metadata = {
                **base_metadata,
                "chunk_type": "test_method" if is_test else "method",
                "class_name": class_name,
                "method_name": method_name,
                "is_test": is_test
            }
            
            chunks.append((method_content, metadata))
        
        return chunks

    def _chunk_xml(self, content: str, file_path: str, base_metadata: Dict) -> List[tuple]:
        """Chunk XML files by layouts and components."""
        chunks = []
        
        if 'layout' in file_path:
            # Layout files - chunk by major containers
            layout_pattern = r'<(LinearLayout|RelativeLayout|ConstraintLayout|FrameLayout|ScrollView|RecyclerView|ListView)[^>]*>.*?</\1>'
            matches = list(re.finditer(layout_pattern, content, re.DOTALL))
            
            for match in matches:
                layout_type = match.group(1)
                metadata = {**base_metadata, "chunk_type": "layout_component", "layout_type": layout_type}
                chunks.append((match.group(), metadata))
            
            # Also add full layout as one chunk
            chunks.append((content, {**base_metadata, "chunk_type": "full_layout"}))
            
        elif 'manifest' in file_path.lower():
            # Manifest file - chunk by components
            component_pattern = r'<(activity|service|receiver|provider)[^>]*>.*?</\1>'
            matches = list(re.finditer(component_pattern, content, re.DOTALL))
            
            for match in matches:
                component_type = match.group(1)
                metadata = {**base_metadata, "chunk_type": "manifest_component", "component_type": component_type}
                chunks.append((match.group(), metadata))
        
        # If no specific chunks found, add whole file
        if not chunks:
            chunks.append((content, {**base_metadata, "chunk_type": "xml_file"}))
        
        return chunks

    def _chunk_gradle(self, content: str, file_path: str, base_metadata: Dict) -> List[tuple]:
        """Chunk Gradle files by sections."""
        chunks = []
        
        # Split by major sections
        sections = ['dependencies', 'android', 'plugins', 'buildTypes', 'productFlavors']
        
        for section in sections:
            pattern = rf'{section}\s*\{{[^}}]*\}}'
            match = re.search(pattern, content, re.DOTALL)
            if match:
                metadata = {**base_metadata, "chunk_type": "gradle_section", "section": section}
                chunks.append((match.group(), metadata))
        
        # Also add full gradle file
        chunks.append((content, {**base_metadata, "chunk_type": "gradle_file"}))
        
        return chunks

    def build_vector_store(self, root_dir: str, persist_directory: str = "rag_chroma_db"):
        """Build and store the vector database."""
        print("🔍 Ingesting codebase...")
        all_files = self.ingest_codebase(root_dir)
        cleaned_files = self.clean_codebase(all_files)
        print(f"📁 Found {len(cleaned_files)} relevant files")
        
        print("⚡ Chunking files...")
        all_chunks = []
        for file_path in cleaned_files:
            chunks = self.chunk_file_enhanced(file_path)
            all_chunks.extend(chunks)
        
        print(f"📚 Created {len(all_chunks)} code chunks")
        
        if len(all_chunks) == 0:
            print("❌ No chunks created! Check your root directory path.")
            return None
        
        # Prepare for vector store
        texts = []
        metadatas = []
        
        for chunk_text, metadata in all_chunks:
            # Enhance text with context
            enhanced_text = self._enhance_chunk_text(chunk_text, metadata)
            texts.append(enhanced_text)
            
            # Flatten metadata for ChromaDB
            flat_metadata = self._flatten_metadata(metadata)
            metadatas.append(flat_metadata)
        
        print("🚀 Building vector store...")
        self.vectordb = Chroma.from_texts(
            texts=texts,
            embedding=self.embeddings,
            metadatas=metadatas,
            persist_directory=persist_directory
        )
        
        print(f"✅ Vector store built with {len(texts)} chunks!")
        return self.vectordb

    def _enhance_chunk_text(self, chunk_text: str, metadata: Dict) -> str:
        """Enhanced text for better Movie Mania search."""
        enhanced_text = chunk_text
        
        # Add file context
        file_name = os.path.basename(metadata.get('file', ''))
        enhanced_text += f"\n\n// File: {file_name}"
        
        # Add type context
        if metadata.get('chunk_type'):
            enhanced_text += f"\n// Type: {metadata['chunk_type']}"
        
        # Add class/method context
        if metadata.get('class_name'):
            enhanced_text += f"\n// Class: {metadata['class_name']}"
        if metadata.get('method_name'):
            enhanced_text += f"\n// Method: {metadata['method_name']}"
        
        # Add Movie Mania specific context
        if 'api' in file_name.lower():
            enhanced_text += "\n// Context: API network request response handling"
        if 'activity' in file_name.lower():
            enhanced_text += "\n// Context: UI activity layout user interface"  
        if 'adapter' in file_name.lower():
            enhanced_text += "\n// Context: RecyclerView adapter data binding"
        if 'chat' in file_name.lower():
            enhanced_text += "\n// Context: Chat messaging API response"
        if 'movie' in file_name.lower():
            enhanced_text += "\n// Context: Movie data TMDB API response"
        
        # Add dependency context
        deps = metadata.get('dependencies', {})
        context_parts = []
        
        if deps.get('layout_references'):
            context_parts.append(f"Layouts: {', '.join(deps['layout_references'][:3])}")
        if deps.get('findViewById_calls'):
            context_parts.append(f"Views: {', '.join(deps['findViewById_calls'][:3])}")
        if deps.get('api_calls'):
            context_parts.append(f"APIs: {', '.join(deps['api_calls'][:3])}")
        if deps.get('test_methods'):
            context_parts.append(f"Tests: {', '.join(deps['test_methods'][:3])}")
        
        if context_parts:
            enhanced_text += f"\n// Context: {' | '.join(context_parts)}"
        
        return enhanced_text

    def _flatten_metadata(self, metadata: Dict) -> Dict:
        """Flatten metadata for ChromaDB compatibility."""
        flat = {}
        
        for key, value in metadata.items():
            if key == 'dependencies':
                # Flatten dependencies
                if isinstance(value, dict):
                    for dep_key, dep_value in value.items():
                        if isinstance(dep_value, list):
                            flat[f"dep_{dep_key}"] = ", ".join(dep_value) if dep_value else ""
                        else:
                            flat[f"dep_{dep_key}"] = str(dep_value) if dep_value else ""
            elif isinstance(value, (str, int, float, bool)) or value is None:
                flat[key] = value
            elif isinstance(value, list):
                flat[key] = ", ".join(map(str, value)) if value else ""
            else:
                flat[key] = str(value)
        
        return flat

    def load_vector_store(self, persist_directory: str = "rag_chroma_db"):
        """Load existing vector store and validate it has content."""
        try:
            # Check if persist directory exists and has content
            if not os.path.exists(persist_directory):
                print("📁 Vector store directory doesn't exist")
                return False
            
            # List files in the directory
            files = os.listdir(persist_directory)
            if not files:
                print("📁 Vector store directory is empty")
                return False
            
            # Try to load the vector store
            self.vectordb = Chroma(
                persist_directory=persist_directory,
                embedding_function=self.embeddings
            )
            
            # Validate the vector store has content
            try:
                collection = self.vectordb._collection
                count = collection.count()
                
                if count == 0:
                    print("📊 Vector store is empty (0 documents)")
                    return False
                
                print(f"✅ Loaded existing vector store with {count} documents")
                return True
                
            except Exception as e:
                print(f"❌ Vector store validation failed: {e}")
                return False
                
        except Exception as e:
            print(f"❌ Failed to load vector store: {e}")
            return False

    def force_rebuild_vector_store(self, root_dir: str, persist_directory: str = "rag_chroma_db"):
        """Force rebuild the vector store from scratch."""
        print("🗑️ Removing existing vector store...")
        
        # Remove existing directory
        if os.path.exists(persist_directory):
            shutil.rmtree(persist_directory)
            print(f"🗑️ Removed {persist_directory}")
        
        # Build new vector store
        print("🔨 Building new vector store from scratch...")
        return self.build_vector_store(root_dir, persist_directory)

    def analyze_problem(self, problem_description: str) -> Dict[str, Any]:
        """Analyze the problem and categorize it."""
        problem_lower = problem_description.lower()
        
        # Detect problem category
        detected_categories = []
        for category, keywords in self.problem_categories.items():
            if any(keyword in problem_lower for keyword in keywords):
                detected_categories.append(category)
        
        # Extract key technical terms
        technical_terms = []
        for word in problem_description.split():
            if (len(word) > 3 and 
                (word.lower().endswith(('activity', 'fragment', 'view', 'adapter', 'service', 'receiver')) or
                 word.lower() in ['api', 'http', 'json', 'xml', 'database', 'auth', 'login'])):
                technical_terms.append(word.lower())
        
        # Determine if it's a UI problem or functional problem
        ui_keywords = ['layout', 'view', 'button', 'text', 'image', 'color', 'theme', 'style', 'responsive', 'display', 'screen', 'ui', 'xml']
        is_ui_problem = any(keyword in problem_lower for keyword in ui_keywords)
        
        return {
            'categories': detected_categories,
            'technical_terms': technical_terms,
            'is_ui_problem': is_ui_problem,
            'severity': 'high' if any(word in problem_lower for word in ['crash', 'error', 'fail', 'broken']) else 'medium'
        }

    def generate_search_queries(self, problem_description: str, analysis: Dict) -> List[str]:
        """Generate better search queries for Movie Mania codebase."""
        queries = [problem_description]
        
        # Movie Mania specific queries
        if "api" in problem_description.lower():
            queries.extend([
                "RetroInstance retrofit api",
                "ApiInterface movie response", 
                "ChatResponse parsing",
                "MovieResponse body",
                "onResponse onFailure callback"
            ])
        
        if "layout" in problem_description.lower() or "ui" in problem_description.lower():
            queries.extend([
                "activity_main layout RecyclerView",
                "ConstraintLayout responsive design",
                "layout_width match_parent",
                "image_gv_item movie grid"
            ])
        
        # Category-based queries
        for category in analysis['categories']:
            if category == 'authentication':
                queries.extend([
                    "phone number authentication login verification",
                    "OTP verification code implementation",
                    "user authentication flow activity"
                ])
            elif category == 'ui':
                queries.extend([
                    "layout XML view components",
                    "UI rendering display issues",
                    "button click listener implementation"
                ])
            elif category == 'network':
                queries.extend([
                    "API call retrofit network request",
                    "HTTP response handling error",
                    "internet connectivity check"
                ])
        
        # Technical term queries
        for term in analysis['technical_terms']:
            queries.append(f"{term} implementation example")
        
        return list(set(queries))  # Remove duplicates

    def retrieve_relevant_code(self, problem_description: str, k: int = 8) -> List[Document]:
        """Retrieve relevant code chunks based on the problem."""
        if not self.vectordb:
            raise ValueError("Vector store not initialized. Call build_vector_store() first.")
        
        # Analyze problem
        analysis = self.analyze_problem(problem_description)
        
        # Generate multiple queries
        queries = self.generate_search_queries(problem_description, analysis)
        
        # Retrieve for each query
        all_results = []
        for query in queries[:3]:  # Limit to top 3 queries
            results = self.vectordb.similarity_search(query, k=k//3)
            all_results.extend(results)
        
        # Remove duplicates and rank by relevance
        seen_content = set()
        unique_results = []
        for doc in all_results:
            content_hash = hash(doc.page_content[:200])  # Use first 200 chars as hash
            if content_hash not in seen_content:
                seen_content.add(content_hash)
                unique_results.append(doc)
        
        return unique_results[:k]

    def generate_solution(self, problem_description: str, relevant_code: List[Document]) -> Dict[str, str]:
        """Generate comprehensive solution with intelligent test selection."""
        
        # Analyze problem type
        analysis = self.analyze_problem(problem_description)
        
        # Prepare context from retrieved code
        context = self._prepare_context(relevant_code)
        
        # Generate different types of solutions
        solutions = {}
        
        # 1. Code Analysis and Bug Fix
        solutions['analysis'] = self._generate_code_analysis(problem_description, context)
        
        # 2. Generate appropriate test based on problem type
        if analysis['is_ui_problem']:
            solutions['test_code'] = self._generate_ui_tests(problem_description, context)
            solutions['test_type'] = 'UI Test'
        else:
            solutions['test_code'] = self._generate_unit_tests(problem_description, context)
            solutions['test_type'] = 'Unit Test'
        
        # 3. Complete Solution with Code Implementation
        solutions['complete_solution'] = self._generate_complete_solution(problem_description, context)
        
        return solutions

    def _prepare_context(self, relevant_code: List[Document]) -> str:
        """Prepare context from relevant code chunks with size control."""
        context_parts = []
        total_length = 0
        max_context_length = 3000  # Reduced for better performance
        
        for i, doc in enumerate(relevant_code[:4]):  # Limit to 4 most relevant
            metadata = doc.metadata
            file_name = os.path.basename(metadata.get('file', 'unknown'))
            chunk_type = metadata.get('chunk_type', 'unknown')
            
            # Truncate very long chunks
            content = doc.page_content
            if len(content) > 600:  # Limit each chunk
                content = content[:600] + "...[truncated for context]"
            
            context_part = f"// File: {file_name} ({chunk_type})\n{content}\n"
            
            if total_length + len(context_part) > max_context_length:
                break
                
            context_parts.append(context_part)
            total_length += len(context_part)
        
        return "\n".join(context_parts)

    def _generate_code_analysis(self, problem_description: str, context: str) -> str:
        """Generate comprehensive code analysis."""
        prompt = PromptTemplate(
            input_variables=["problem", "context"],
            template="""
You are an expert Android developer analyzing a Movie Mania app issue.

Problem: {problem}

Code Context:
{context}

Provide an analysis:

## Problem Analysis
- Root cause of the issue
- Specific code problems identified
- Impact on app functionality

## Key Issues Found
- List specific bugs or problems
- Code sections that need attention
- Potential side effects

## Recommended Fixes
- Step-by-step solution approach
- Code modifications needed
- Best practices to follow

Be thorough and provide specific, actionable insights based on the actual code provided and be specific.
"""
        )
        
        formatted_prompt = prompt.format(problem=problem_description, context=context)
        response = self.llm.invoke(formatted_prompt)
        return response

    def _generate_ui_tests(self, problem_description: str, context: str) -> str:
        """Generate UI tests using Espresso for UI-related problems."""
        prompt = PromptTemplate(
            input_variables=["problem", "context"],
            template="""
Generate comprehensive UI tests for this Android UI problem:

Problem: {problem}

Code Context:
{context}

Create complete Espresso UI tests including:

## UI Test Class
```java
@RunWith(AndroidJUnit4.class)
public class UITestClass {{
    // Complete test implementation
}}
```

## Key Test Methods
- Test the specific UI functionality mentioned in the problem
- Test user interactions (clicks, text input, navigation)
- Test UI state changes and visual elements
- Test responsive behavior if relevant

## Test Setup and Teardown
- Activity rules and test setup
- Mock data if needed
- Test cleanup

Provide complete, runnable Espresso test code that thoroughly tests the UI issue.
"""
        )
        
        formatted_prompt = prompt.format(problem=problem_description, context=context)
        response = self.llm.invoke(formatted_prompt)
        return response

    def _generate_unit_tests(self, problem_description: str, context: str) -> str:
        """Generate unit tests for functional problems."""
        prompt = PromptTemplate(
            input_variables=["problem", "context"],
            template="""
Generate comprehensive unit tests for this functional Android problem:

Problem: {problem}

Code Context:
{context}

Create complete JUnit unit tests including:

## Unit Test Class
```java
public class FunctionalTest {{
    // Complete test implementation
}}
```

## Test Methods
- Test the specific functionality mentioned in the problem
- Test normal cases and edge cases
- Test error handling and exceptions
- Test data processing and business logic

## Mock Objects and Setup
- Use Mockito for dependencies
- Set up test data and mocks
- Test isolation and independence

## Assertions and Validation
- Verify expected behavior
- Check error conditions
- Validate data integrity

Provide complete, runnable JUnit test code with proper annotations and assertions.
"""
        )
        
        formatted_prompt = prompt.format(problem=problem_description, context=context)
        response = self.llm.invoke(formatted_prompt)
        return response

    def _generate_complete_solution(self, problem_description: str, context: str) -> str:
        """Generate a complete solution with working code."""
        prompt = PromptTemplate(
            input_variables=["problem", "context"],
            template="""
You are a senior Android developer providing a complete solution for the Movie Mania app.

Problem: {problem}

Existing Code Context:
{context}

Provide a comprehensive solution:

## Root Cause Analysis
Explain exactly what's causing this problem in the existing code.

## Complete Solution Implementation

### Step 1: Identify the Issue
- Specific code locations with problems
- Why the current implementation fails

### Step 2: Code Implementation
Provide complete, working code fixes:

```java
// Complete working solution code
public class FixedImplementation {{
    // Full implementation here
}}
```

### Step 3: Integration Steps
- How to integrate the fix into existing code
- Files that need to be modified
- Configuration changes if needed

### Step 4: Verification Steps
- How to test the fix
- Expected behavior after fix
- How to prevent similar issues

## Additional Improvements
- Code optimization opportunities
- Best practices recommendations
- Performance considerations

Provide complete, copy-paste ready code solutions that will definitively fix the problem.
"""
        )
        
        formatted_prompt = prompt.format(problem=problem_description, context=context)
        response = self.llm.invoke(formatted_prompt)
        return response

    def solve_problem(self, problem_description: str) -> Dict[str, Any]:
        """Main method to solve a coding problem using RAG."""
        print(f"🔍 Analyzing problem: {problem_description}")
        
        # Analyze problem type
        analysis = self.analyze_problem(problem_description)
        print(f"📊 Problem type: {'UI Problem' if analysis['is_ui_problem'] else 'Functional Problem'}")
        
        # Retrieve relevant code
        relevant_code = self.retrieve_relevant_code(problem_description)
        print(f"📚 Found {len(relevant_code)} relevant code chunks")
        
        # Generate solutions
        print("🤖 Generating comprehensive solution...")
        start_time = time.time()
        solutions = self.generate_solution(problem_description, relevant_code)
        processing_time = time.time() - start_time
        # Prepare response
        response = {
            'problem': problem_description,
            'problem_type': 'UI Problem' if analysis['is_ui_problem'] else 'Functional Problem',
            'relevant_files': [os.path.basename(doc.metadata.get('file', '')) for doc in relevant_code[:5]],
            'analysis': solutions.get('analysis', ''),
            'test_type': solutions.get('test_type', ''),
            'test_code': solutions.get('test_code', ''),
            'complete_solution': solutions.get('complete_solution', ''),
            'code_chunks_found': len(relevant_code),
            'processing_time_seconds': processing_time
        }
        
        return response

    def debug_search(self, query: str, k: int = 5):
        """Debug search functionality."""
        print(f"🔍 Debugging search for: '{query}'")
        
        if not self.vectordb:
            print("❌ Vector store not initialized!")
            return
        
        try:
            results = self.vectordb.similarity_search(query, k=k)
            print(f"📚 Found {len(results)} results")
            
            for i, result in enumerate(results):
                file_name = os.path.basename(result.metadata.get('file', 'unknown'))
                chunk_type = result.metadata.get('chunk_type', 'unknown')
                print(f"  {i+1}. {file_name} ({chunk_type})")
                print(f"     Preview: {result.page_content[:100]}...")
            
        except Exception as e:
            print(f"❌ Error debugging search: {e}")

def main():
    """Enhanced main function with better problem solving."""
    assistant = CodebaseRAGAssistant(
        model_name="codellama:7b",
        embedding_model="nomic-embed-text",
        temperature=0.1
    )
    
    root_dir = "Movie-Maina-main"
    
    # Try to load existing vector store, rebuild if necessary
    if not assistant.load_vector_store():
        print("Building new vector store...")
        assistant.build_vector_store(root_dir)
    
    # Test problems for Movie Mania
    problems = [
        "API response is not being parsed correctly",
        "RecyclerView adapter is not displaying data correctly", 
        "UI layout is not responsive on different screen sizes",
        "Chat activity API call not working properly"
    ]
    
    # Solve problems one by one
    for problem in problems[:2]:  # Solve first 2 problems
        print("\n" + "="*80)
        print(f"🔧 SOLVING: {problem}")
        print("="*80)
        
        try:
            solution = assistant.solve_problem(problem)
            
            print(f"\n📋 **Problem**: {solution['problem']}")
            print(f"🎯 **Problem Type**: {solution['problem_type']}")
            print(f"📁 **Relevant Files**: {', '.join(solution['relevant_files'])}")
            print(f"📊 **Code Chunks Found**: {solution['code_chunks_found']}")
            
            if solution['code_chunks_found'] > 0:
                print(f"\n" + "="*60)
                print("🔍 **DETAILED ANALYSIS**")
                print("="*60)
                print(solution['analysis'])
                
                print(f"\n" + "="*60)
                print(f"🧪 **{solution['test_type'].upper()}**")
                print("="*60)
                print(solution['test_code'])
                
                print(f"\n" + "="*60)
                print("✅ **COMPLETE SOLUTION WITH CODE**")
                print("="*60)
                print(solution['complete_solution'])
                print(f"\n" + "="*60)
                print(f"\n⏱️ Processing Time: {solution['processing_time_seconds']:.2f} seconds")
            else:
                print("\n❌ No relevant code found!")
                
        except Exception as e:
            print(f"❌ Error solving problem: {e}")
            import traceback

            traceback.print_exc()

if __name__ == "__main__":
    main()