# Core Imports
import time
import os
import pandas as pd
import re
import json
import torch
import gc
from datetime import datetime
from typing import TypedDict, Annotated, Literal

# LangChain Imports
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field, ValidationError

# Local Imports
from solution_rag import CodebaseRAGAssistant
from tools import (
    scrape_english_reviews_with_aspects, 
    filter_reviews_by_aspect,
    check_review_relevance_llm_json,
    infer_review_with_local_model,
    get_app_changelog_versions,
    save_app_info_raw_text,
    get_recent_changes,
    filter_negative_aspects,
    initialize_absa_model,
    parse_aspects_from_query
)

# =============================================================================
# CONFIGURATION & INITIALIZATION
# =============================================================================

CHAT_MODEL = 'llama3.1:latest'
llm = init_chat_model(CHAT_MODEL, model_provider='ollama')

# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def log_to_file(message, log_file="system_logs.txt"):
    """Log messages to a text file with timestamp"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"[{timestamp}] {message}\n"
    
    os.makedirs(os.path.dirname(log_file) if os.path.dirname(log_file) else ".", exist_ok=True)
    
    with open(log_file, "a", encoding="utf-8") as f:
        f.write(log_entry)

def clear_gpu_memory():
    """Clear GPU memory to prevent CUDA OOM errors"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    gc.collect()

def extract_review_from_query(query):
    """Extract review text from query"""
    patterns = [
        r"analyze\s+this\s+review:?\s*(.*)",
        r"what\s+are\s+(?:the\s+)?aspects?\s+(?:and\s+opinions?\s+)?(?:of\s+)?(?:this\s+)?review:?\s*(.*)",
        r"review:?\s*(.*)"
    ]
    
    for pattern in patterns:
        match = re.search(pattern, query, re.IGNORECASE)
        if match:
            review_text = match.group(1).strip()
            if len(review_text) > 10:
                return review_text
    return None

# =============================================================================
# STATE MANAGEMENT & MODELS
# =============================================================================

class DeveloperQueryState(TypedDict):
    messages: Annotated[list, add_messages]
    query_type: str | None
    raw_query: str | None
    aspect: str | None
    version: str | None
    app_ids: list | None
    reviews: list | None
    app_info: str | None
    recent_changes: str | None
    filtered_reviews: list | None
    absa_results: list | None
    negative_aspects: list | None
    problems: list | None
    solutions: list | None
    final_response: str | None
    review_text: str | None
    conversation_context: dict | None

class QueryClassifier(BaseModel):
    query_type: Literal[
        "version_issues", "aspect_issues", "version_aspect_issues", 
        "solutions_request", "reviews_request", "single_review_analysis",
        "general_analysis", "casual_conversation", "direct_absa",
        "solution_only", "aspect_opinion_query", "problems_only"
    ] = Field(..., description="Classify the developer query type")

# =============================================================================
# CORE AGENTS
# =============================================================================

def classify_query(state: DeveloperQueryState):
    """Enhanced classifier with better solution detection"""
    last_message = state["messages"][-1]
    query = last_message.content
    
    classifier_llm = llm.with_structured_output(QueryClassifier)
    result = classifier_llm.invoke([
        {"role": "system", 
         "content": """Classify developer queries as:
         - 'direct_absa': asks to analyze specific review text (e.g., "analyze this review: [text]")
         - 'aspect_opinion_query': asks about aspects/opinions WITHOUT providing review text
         - 'solution_only': asks ONLY for solutions without analysis (e.g., "what are solutions for login issues", "how to fix payment problems")
         - 'problems_only': asks ONLY about problems/issues (e.g., "what are the problems", "show me issues")
         - 'version_issues': asks about issues in specific version with solutions implied
         - 'aspect_issues': asks about specific aspect issues with solutions implied
         - 'version_aspect_issues': asks about aspect in specific version with solutions implied
         - 'solutions_request': asks for solutions after showing problems (conversational follow-up)
         - 'reviews_request': asks to see reviews
         - 'general_analysis': asks for overall app analysis including both problems and solutions
         - 'casual_conversation': general chat not related to app analysis
         """},
        {"role": "user", "content": query}
    ])
    
    return {
        "query_type": result.query_type,
        "raw_query": query
    }

def parse_query_details(state: DeveloperQueryState):
    """Extract aspect, version, and review text from query"""
    query = state["raw_query"]
    query_type = state["query_type"]
    
    # Extract review text for direct ABSA queries
    review_text = None
    if query_type == "direct_absa":
        review_text = extract_review_from_query(query)
    
    # Extract version using regex
    version_pattern = r'(?:version\s+)?(\d+\.\d+(?:\.\d+)?(?:\.\d+)?)'
    version_match = re.search(version_pattern, query.lower())
    version = version_match.group(1) if version_match else None
    
    # Extract aspect
    aspects = parse_aspects_from_query(query)
    aspect = aspects[0] if aspects else None
    
    # Use app_ids from state if available, otherwise use default
    existing_app_ids = state.get("app_ids")
    if existing_app_ids and len(existing_app_ids) > 0:
        app_ids = existing_app_ids
    else:
        app_ids = ["com.a10minuteschool.tenminuteschool"]
    
    return {
        "aspect": aspect,
        "version": version,
        "app_ids": app_ids,
        "review_text": review_text
    }

def absa_model_agent(state: DeveloperQueryState):
    """Unified ABSA Agent with Ollama flag - handles both direct review analysis and scraped reviews"""
    
    # Configuration flag - set to True to use Ollama, False to use existing Unsloth
    OLLAMA_FLAG = False
    
    query_type = state["query_type"]
    review_text = state.get("review_text")
    filtered_reviews = state.get("filtered_reviews", [])
    
    # Import the appropriate ABSA function based on flag
    if OLLAMA_FLAG:
        try:
            from ollamaabsa import infer_review_with_local_model as ollama_inference
            inference_function = ollama_inference
            print("🦙 Using Ollama ABSA")
        except ImportError as e:
            print(f"❌ Ollama ABSA import failed: {e}")
            print("🔄 Falling back to existing Unsloth model")
            inference_function = infer_review_with_local_model
    else:
        inference_function = infer_review_with_local_model
        print("🔥 Using Unsloth ABSA model")
    
    # Case 1: Direct review analysis
    if query_type == "direct_absa" and review_text:
        print(f"🧠 ABSA Agent: Direct analysis of: '{review_text[:50]}...'")
        
        try:
            if not OLLAMA_FLAG:
                clear_gpu_memory()
            
            absa_result = inference_function(review_text)
            
            if not OLLAMA_FLAG:
                clear_gpu_memory()
            
            return {
                "absa_results": [{
                    "review": review_text,
                    "analysis": absa_result
                }],
                "negative_aspects": []
            }
            
        except Exception as e:
            print(f"❌ Direct ABSA error: {e}")
            return {"absa_results": [], "negative_aspects": []}
    
    # Case 2: Scraped reviews analysis
    elif filtered_reviews:
        print(f"🧠 ABSA Agent: Analyzing {len(filtered_reviews)} scraped reviews")
        
        try:
            all_absa_results = []
            
            for i, review in enumerate(filtered_reviews):
                try:
                    print(f"Processing review {i+1}/{len(filtered_reviews)}")
                    
                    absa_result = inference_function(review)
                    if absa_result:
                        all_absa_results.append({
                            "review": review,
                            "analysis": absa_result
                        })
                        
                except Exception as e:
                    print(f"ABSA error for review {i+1}: {e}")
                    continue
            
            # Extract negative aspects for problem analysis
            negative_aspects = []
            if query_type in ["version_issues", "aspect_issues", "version_aspect_issues", "general_analysis"]:
                for result in all_absa_results:
                    analysis = result["analysis"]
                    
                    # Handle both dict and list analysis formats
                    if isinstance(analysis, dict):
                        for aspect_name, aspect_data in analysis.items():
                            if isinstance(aspect_data, dict) and 'sentiment' in aspect_data:
                                sentiment = aspect_data['sentiment'].lower()
                                if sentiment in ['negative', 'neutral']:
                                    negative_aspects.append({
                                        "review": result["review"],
                                        "aspect_data": {
                                            "aspect": aspect_name,
                                            "sentiment": aspect_data['sentiment'],
                                            "opinion": aspect_data.get('opinion', '')
                                        }
                                    })
                    elif isinstance(analysis, list):
                        for aspect_item in analysis:
                            if isinstance(aspect_item, dict) and 'sentiment' in aspect_item:
                                sentiment = aspect_item['sentiment'].lower()
                                if sentiment in ['negative', 'neutral']:
                                    negative_aspects.append({
                                        "review": result["review"],
                                        "aspect_data": {
                                            "aspect": aspect_item.get('aspect', 'unknown'),
                                            "sentiment": aspect_item['sentiment'],
                                            "opinion": aspect_item.get('opinion', '')
                                        }
                                    })
            
            return {
                "absa_results": all_absa_results,
                "negative_aspects": negative_aspects
            }
            
        except Exception as e:
            print(f"❌ ABSA Agent error: {e}")
            return {"absa_results": [], "negative_aspects": []}
    
    # Case 3: No reviews to analyze
    else:
        return {"absa_results": [], "negative_aspects": []}

def review_scraper_agent(state: DeveloperQueryState):
    """Agent 1: Get reviews based on query and version"""
    
    app_ids = state["app_ids"]
    version = state["version"]
    aspect = state["aspect"]
    query_type = state["query_type"]
    
    print(f"🔍 Scraper Agent: Getting reviews for {app_ids}, version: {version}, aspect: {aspect}")
    
    try:
        # Get version-specific reviews
        versions_list = [version] if version else None
        scrape_english_reviews_with_aspects(
            app_ids, 
            num_reviews=100, 
            versions=versions_list
        )
        
        # Load reviews
        reviews_df = pd.read_csv("english_app_reviews_with_aspects.csv")
        print(f"📊 Loaded {len(reviews_df)} reviews")
        
        # Filter by aspect ONLY if aspect is provided and query needs filtering
        if aspect and query_type in ["aspect_issues", "version_aspect_issues", "solutions_request"]:
            aspect_reviews = filter_reviews_by_aspect(reviews_df, aspect)
            print(f"🎯 Found {len(aspect_reviews)} aspect-related reviews")
        else:
            aspect_reviews = reviews_df['review_description'].tolist()
            print(f"📄 Returning all {len(aspect_reviews)} reviews (no aspect filtering)")
        
        return {"reviews": aspect_reviews}
        
    except Exception as e:
        print(f"❌ Scraper Agent error: {e}")
        return {"reviews": []}

def context_agent(state: DeveloperQueryState):
    """Agent 2: Get app context and filter for relevance"""
    
    reviews = state["reviews"]
    app_ids = state["app_ids"]
    version = state["version"]
    
    if not reviews or not app_ids:
        return {"filtered_reviews": []}
    
    print(f"🏢 Context Agent: Getting app info and filtering {len(reviews)} reviews")
    
    try:
        app_id = app_ids[0]
        
        # Get app context
        app_info = save_app_info_raw_text(app_id)
        recent_changes = get_recent_changes(app_id)
        
        # Filter for relevance
        filtered_reviews = check_review_relevance_llm_json(
            reviews, app_info, recent_changes, version
        )
        
        print(f"✅ Context Agent: {len(filtered_reviews)} relevant reviews")
        
        return {
            "app_info": app_info,
            "recent_changes": recent_changes,
            "filtered_reviews": filtered_reviews
        }
        
    except Exception as e:
        print(f"❌ Context Agent error: {e}")
        return {"filtered_reviews": reviews}

def problem_analysis_agent(state: DeveloperQueryState):
    """Agent 4: Generate structured list of problems"""
    
    negative_aspects = state["negative_aspects"]
    raw_query = state["raw_query"]
    
    if not negative_aspects:
        return {"problems": ["No significant issues found in the analyzed reviews."]}
    
    print(f"🔍 Problem Agent: Analyzing {len(negative_aspects)} negative aspects")
    
    try:
        # Prepare data for LLM
        aspects_summary = []
        for item in negative_aspects:
            aspects_summary.append({
            "review": item["review"][:100] + "...",
            "aspect": item["aspect_data"]
            })
        
        prompt = f"""Based on the negative aspects from user reviews, identify the main problems:

DEVELOPER QUERY: {raw_query}

NEGATIVE ASPECTS: {json.dumps(aspects_summary, indent=2)}

Generate a clear list of the MAIN PROBLEMS. Be specific and actionable.

Format as:
1. [Specific problem description]
2. [Specific problem description]
3. [Specific problem description]
...

Focus on the most frequently mentioned and critical issues. and In specific, not general analysis."""

        response = llm.invoke([
            {"role": "system", "content": "You are an expert at analyzing app review data to identify key problems and In specific, not general analysis."},
            {"role": "user", "content": prompt}
        ])
        
        # Parse problems from response
        problems = []
        lines = response.content.split('\n')
        for line in lines:
            line = line.strip()
            if re.match(r'^\d+\.', line):
                problem = re.sub(r'^\d+\.\s*', '', line)
                problems.append(problem)
        
        print(f"🎯 Problem Agent: Identified {len(problems)} problems")
        return {"problems": problems}
        
    except Exception as e:
        print(f"❌ Problem Agent error: {e}")
        return {"problems": ["Error analyzing problems"]}

# =============================================================================
# SOLUTIONS AGENT WITH RAG
# =============================================================================

def solutions_agent(state: DeveloperQueryState):
    """Enhanced Solutions Agent using RAG system - handles both full workflow and direct solution requests"""
    
    problems = state.get("problems", [])
    raw_query = state["raw_query"]
    query_type = state["query_type"]
    aspect = state.get("aspect")
    
    print(f"💡 RAG Solutions Agent: Generating solutions for '{raw_query}' (Type: {query_type})")
    
    try:
        # Initialize RAG Assistant
        rag_assistant = CodebaseRAGAssistant(
            model_name="codellama:7b",
            embedding_model="nomic-embed-text", 
            temperature=0.1
        )
        
        # Try to load existing vector store, build if needed
        if not rag_assistant.load_vector_store():
            print("🔄 Building vector store for codebase analysis...")
            root_dir = "Movie-Maina-main"
            if not rag_assistant.build_vector_store(root_dir):
                print("⚠️ RAG system unavailable, falling back to LLM-only solutions")
                return fallback_llm_solutions(problems, raw_query, query_type)
        
        solutions = []
        
        # Branch 1: Solution-only request (no problems provided)
        if query_type == "solution_only" or not problems or problems == ["No significant issues found in the analyzed reviews."]:
            print("💡 Generating direct RAG solutions without problem analysis...")
            
            rag_solution = rag_assistant.solve_problem(raw_query)
            formatted_solution = format_rag_solution(rag_solution, raw_query)
            solutions = [formatted_solution]
            
        # Branch 2: Full workflow (problems already analyzed)
        else:
            print(f"💡 Generating RAG solutions for {len(problems)} analyzed problems...")
            
            all_solutions = []
            
            # Process each problem with RAG
            for i, problem in enumerate(problems, 1):
                print(f"🔍 Processing problem {i}/{len(problems)}: {problem[:50]}...")
                
                try:
                    rag_solution = rag_assistant.solve_problem(problem)
                    
                    formatted_problem_solution = f"""
## Problem {i}: {problem}

### 🔍 Analysis
{rag_solution.get('analysis', 'No analysis available')}

### 🧪 {rag_solution.get('test_type', 'Test Code')}
```java
{rag_solution.get('test_code', 'No test code available')}
```

### 💻 Complete Solution
{rag_solution.get('complete_solution', 'No solution code available')}

**Relevant Files**: {', '.join(rag_solution.get('relevant_files', []))}
**Processing Time**: {rag_solution.get('processing_time_seconds', 0):.2f}s

---"""
                    
                    all_solutions.append(formatted_problem_solution)
                    
                except Exception as problem_error:
                    print(f"❌ Error processing problem {i}: {problem_error}")
                    fallback_solution = f"""
## Problem {i}: {problem}

❌ **RAG Analysis Failed**: Using fallback solution

### Basic Solution Approach
{generate_fallback_solution_for_problem(problem)}

---"""
                    all_solutions.append(fallback_solution)
            
            # Combine all solutions
            combined_solution = f"""# 🚀 RAG-Enhanced Solutions for: "{raw_query}"

**Total Problems Analyzed**: {len(problems)}
**RAG System**: Codebase-aware solutions with tests

{''.join(all_solutions)}

## 📋 Summary
- **Problem Type**: Mobile app development issues
- **Solution Approach**: RAG-enhanced with existing codebase analysis
- **Tests Included**: UI/Unit tests based on problem type
- **Code Examples**: Production-ready implementations

*Solutions generated using codebase RAG analysis for context-aware recommendations*"""
            
            solutions = [combined_solution]
        
        print(f"✅ RAG Solutions Agent: Generated comprehensive solutions")
        return {"solutions": solutions}
        
    except Exception as e:
        print(f"❌ RAG Solutions Agent error: {e}")
        print("🔄 Falling back to LLM-only solutions...")
        return fallback_llm_solutions(problems, raw_query, query_type)

def format_rag_solution(rag_solution, query):
    """Format RAG solution for single query response"""
    
    return f"""# 🚀 RAG-Enhanced Solution for: "{query}"

## 🔍 Problem Analysis
{rag_solution.get('analysis', 'No analysis available')}

## 🧪 {rag_solution.get('test_type', 'Test Code')}
```java
{rag_solution.get('test_code', 'No test code available')}
```

## 💻 Complete Implementation Solution
{rag_solution.get('complete_solution', 'No solution available')}

## 📁 Codebase Context
**Relevant Files Analyzed**: {', '.join(rag_solution.get('relevant_files', ['None']))}
**Code Chunks Found**: {rag_solution.get('code_chunks_found', 0)}
**Problem Type**: {rag_solution.get('problem_type', 'Unknown')}
**Processing Time**: {rag_solution.get('processing_time_seconds', 0):.2f} seconds

---
*Solution generated using RAG analysis of existing codebase for context-aware recommendations*"""

def generate_fallback_solution_for_problem(problem):
    """Generate basic fallback solution when RAG fails"""
    try:
        prompt = f"""Provide a basic technical solution for this mobile app problem:

PROBLEM: {problem}

Provide:
1. Root cause analysis
2. Implementation steps  
3. Code example (if applicable)
4. Testing approach

Be specific and practical."""

        response = llm.invoke([
            {"role": "system", "content": "You are a mobile app developer providing technical solutions."},
            {"role": "user", "content": prompt}
        ])
        
        return response.content
        
    except Exception:
        return f"Basic approach: Analyze the '{problem}' issue, implement standard mobile development best practices, and test thoroughly."

def fallback_llm_solutions(problems, raw_query, query_type):
    """Fallback to original LLM-based solutions when RAG is unavailable"""
    print("🔄 Using LLM-only fallback solutions...")
    
    try:
        # Branch 1: Solution-only request  
        if query_type == "solution_only" or not problems or problems == ["No significant issues found in the analyzed reviews."]:
            prompt = f"""Generate specific technical solutions for the following request:

DEVELOPER QUERY: {raw_query}

Provide practical, implementable solutions for mobile app development. Focus on common issues and best practices.

Format as numbered list with detailed specific explanations:
1. **Solution Name**: Specific implementation steps and technical details
2. **Solution Name**: Specific implementation steps and technical details

Focus on actionable technical recommendations that developers can implement."""
 
            response = llm.invoke([
                {"role": "system", "content": "You are an expert mobile app developer providing technical solutions based on common development challenges."},
                {"role": "user", "content": prompt}
            ])
            
            solutions = [f"# 💡 Solutions for: \"{raw_query}\"\n\n{response.content}\n\n---\n*Fallback solutions - RAG system unavailable*"]
            
        # Branch 2: Full workflow (problems already analyzed)
        else:
            problems_text = "\n".join([f"{i+1}. {p}" for i, p in enumerate(problems)])
            
            prompt = f"""Based on these identified problems from app reviews, suggest specific technical solutions:

DEVELOPER QUERY: {raw_query}

IDENTIFIED PROBLEMS:
{problems_text}

For each problem, provide a specific, actionable technical solution. Format as:

**Problem 1**: [Problem description]
**Solution**: [Specific technical solution with implementation details]

**Problem 2**: [Problem description]  
**Solution**: [Specific technical solution with implementation details]

Continue for all problems. Focus on practical, implementable solutions."""

            response = llm.invoke([
                {"role": "system", "content": "You are an expert mobile app developer providing technical solutions based on user feedback analysis."},
                {"role": "user", "content": prompt}
            ])
            
            solutions = [f"# 💡 Solutions Analysis\n\n{response.content}\n\n---\n*Fallback solutions - RAG system unavailable*"]
        
        return {"solutions": solutions}
        
    except Exception as e:
        return {"solutions": [f"❌ Error generating fallback solutions: {e}"]}

# =============================================================================
# COORDINATOR AGENT
# =============================================================================

def coordinator_agent(state: DeveloperQueryState):
    """Enhanced coordinator that handles all query types and formats final response"""
    
    query_type = state["query_type"]
    raw_query = state["raw_query"]
    reviews = state.get("reviews", [])
    filtered_reviews = state.get("filtered_reviews", [])
    problems = state.get("problems", [])
    solutions = state.get("solutions", [])
    absa_results = state.get("absa_results", [])
    
    print(f"📋 Coordinator: Formatting response for {query_type}")
    
    # Handle direct ABSA analysis
    if query_type == "direct_absa":
        if absa_results:
            result = absa_results[0]
            review_text = result["review"]
            analysis = result["analysis"]
            
            final_response = f"""# 🔍 Review Analysis

**Review**: "{review_text}"

## Analysis Results:
```json
{json.dumps(analysis, indent=2)}
```

## Summary:
"""
            if isinstance(analysis, dict):
                for aspect, data in analysis.items():
                    if isinstance(data, dict):
                        sentiment = data.get('sentiment', 'Unknown')
                        opinion = data.get('opinion', 'No opinion')
                        final_response += f"- **{aspect}**: {sentiment} - {opinion}\n"
            elif isinstance(analysis, list):
                for i, aspect in enumerate(analysis, 1):
                    final_response += f"{i}. **Aspect**: {aspect}\n"
            else:
                final_response += f"**Result**: {analysis}\n"
            
            return {"final_response": final_response}
        else:
            return {"final_response": "❌ Could not analyze the provided review."}
    
    # Handle casual conversation
    elif query_type == "casual_conversation":
        try:
            system_prompt = """You are a helpful and friendly App Review Analysis Assistant. 
            
            Your capabilities include:
            - ABSA (Aspect-Based Sentiment Analysis) of app reviews
            - Identifying problems and issues from user feedback  
            - Generating technical solutions for identified problems
            - Extracting aspects and opinions from reviews
            - Version-specific analysis
            - App review scraping and filtering
            
            Instructions:
            - Be warm, friendly, and professional
            - For greetings, respond naturally and mention your app review analysis capabilities
            - For questions about capabilities, provide a comprehensive overview
            - For thank you messages, acknowledge gracefully
            - For goodbyes, bid farewell warmly
            - Always redirect conversation toward app review analysis when appropriate
            - Include helpful examples and quick tips when relevant
            - Use emojis to make responses engaging
            
            Example queries you can help with:
            • "What are the main problems in the app?"
            • "Analyze this review: [paste review text]"
            • "What are solutions for login issues?"
            • "Show me problems in version 2.1.0"
            • "What aspects are users complaining about?"
            """
            
            response = llm.invoke([
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": raw_query}
            ])
            
            final_response = response.content
            
            # Add contextual tips for simple greetings
            greeting_words = ['hi', 'hello', 'hey', 'good morning', 'good afternoon', 'good evening']
            if any(word in raw_query.lower() for word in greeting_words):
                final_response += "\n\n💡 **Quick Tips:**\n• Load an app first using the App ID\n• Try asking: 'What are the main problems?' or 'Show me solutions for login issues'"
            
            return {"final_response": final_response}
            
        except Exception as e:
            print(f"❌ Casual conversation error: {e}")
            return {"final_response": "I'm here to help with app review analysis! What would you like to know about your app? 🤖"}
        
    # Handle reviews request
    elif query_type == "reviews_request":
        if filtered_reviews:
            final_response = f"""# 📱 Reviews for: "{raw_query}"

**Total Reviews Found**: {len(filtered_reviews)}

## Sample Reviews:
"""
            for i, review in enumerate(filtered_reviews[:10], 1):
                final_response += f"{i}. {review}\n\n"
                
            if len(filtered_reviews) > 10:
                final_response += f"\n*Showing first 10 of {len(filtered_reviews)} total reviews*"
        else:
            final_response = f"""# 📱 Reviews for: "{raw_query}"

**Total Reviews Found**: 0

❌ No reviews found matching your criteria. This could be due to:
- App ID not found
- No reviews for the specified version
- Reviews not in English
- Network connectivity issues

💡 **Tips:**
- Check the app ID format (e.g., com.company.appname)
- Try without version filtering
- Ensure the app exists on Google Play Store"""
        
        return {"final_response": final_response}
        
    # Handle aspect opinion queries
    elif query_type == "aspect_opinion_query":
        final_response = f"""# 🔍 Aspects and Opinions Analysis: "{raw_query}"

**Reviews Analyzed**: {len(filtered_reviews)}

## Extracted Aspects and Opinions:
"""
        if absa_results:
            for i, result in enumerate(absa_results, 1):
                final_response += f"\n**Review {i}**: {result['review'][:60]}...\n"
                final_response += f"**Analysis**: {json.dumps(result['analysis'], indent=2)}\n"
                final_response += "---\n"
        else:
            final_response += "❌ No aspect analysis available. This could be due to:\n"
            final_response += "- No reviews found for analysis\n"
            final_response += "- ABSA model processing issues\n"
            final_response += "- Reviews not suitable for aspect extraction\n"
            
        return {"final_response": final_response}
    
    # Handle solution-only requests
    elif query_type == "solution_only":
        if solutions and solutions[0]:
            final_response = f"""# 💡 Solutions for: "{raw_query}"

{solutions[0]}

---
*Solutions generated based on mobile app development best practices*"""
        else:
            final_response = f"""# 💡 Solutions for: "{raw_query}"

❌ Could not generate solutions. Please try:
- Being more specific about the problem
- Providing more context about the issue
- Checking if the query is related to mobile app development

💡 **Example queries:** 
- "Solutions for login crashes"
- "How to fix payment processing issues"
- "Solutions for slow app performance"""
            
        return {"final_response": final_response}
    
    # Handle problems-only requests
    elif query_type == "problems_only" or (problems and not solutions):
        final_response = f"""# 🚨 Problems Analysis: "{raw_query}"

**Query Type**: {query_type.replace('_', ' ').title()}
**Reviews Analyzed**: {len(filtered_reviews)}
**Issues Found**: {len(problems)}

## 🚨 Identified Problems:
"""
        if problems and problems != ["No significant issues found in the analyzed reviews."]:
            for i, problem in enumerate(problems, 1):
                final_response += f"{i}. {problem}\n"
        else:
            final_response += "✅ No significant issues found in the analyzed reviews.\n"
            final_response += "\n**This could mean:**\n"
            final_response += "- Users are generally satisfied with the app\n"
            final_response += "- Recent updates have resolved major issues\n"
            final_response += "- The analyzed reviews are mostly positive\n"
        
        final_response += f"\n---\n*Analysis based on {len(filtered_reviews)} relevant user reviews*"
        
        if problems and problems != ["No significant issues found in the analyzed reviews."]:
            final_response += f"\n\n💡 *Want solutions? Ask: 'What are the solutions for these problems?'*"
            
        return {"final_response": final_response}
            
    # Handle full analysis: problems + solutions
    else:
        final_response = f"""# 🔍 Complete Analysis: "{raw_query}"

**Query Type**: {query_type.replace('_', ' ').title()}
**Reviews Analyzed**: {len(filtered_reviews)}
**Issues Found**: {len(problems)}

## 🚨 Main Problems:
"""
        if problems and problems != ["No significant issues found in the analyzed reviews."]:
            for i, problem in enumerate(problems, 1):
                final_response += f"{i}. {problem}\n"
        else:
            final_response += "✅ No significant issues found in the analyzed reviews.\n"
            final_response += "\n**This indicates:**\n"
            final_response += "- High user satisfaction\n"
            final_response += "- Effective issue resolution\n"
            final_response += "- Quality app development\n"
        
        # Add solutions section
        final_response += "\n## 💡 Suggested Solutions:\n"
        if solutions and solutions[0] and solutions != ["No solutions needed - no significant issues found."]:
            final_response += solutions[0]
        else:
            if problems and problems != ["No significant issues found in the analyzed reviews."]:
                final_response += "❌ Could not generate specific solutions. Consider:\n"
                final_response += "- Manual review of the identified problems\n"
                final_response += "- Consulting with development team\n"
                final_response += "- Implementing general best practices\n"
            else:
                final_response += "✅ No solutions needed - no significant issues found.\n"
                final_response += "\n**Recommendations:**\n"
                final_response += "- Continue current development practices\n"
                final_response += "- Monitor reviews for emerging issues\n"
                final_response += "- Focus on feature enhancement\n"
        
        final_response += f"\n---\n*Analysis based on {len(filtered_reviews)} relevant user reviews*"
        
        if len(filtered_reviews) == 0:
            final_response += f"\n\n⚠️ **Note**: No reviews were available for analysis. Please check app ID and try again."
            
        return {"final_response": final_response}

# =============================================================================
# ROUTING FUNCTIONS
# =============================================================================

def route_after_classification(state: DeveloperQueryState):
    """Simplified routing after classification"""
    query_type = state["query_type"]
    
    if query_type == "casual_conversation":
        return "coordinator"
    elif query_type == "solution_only":
        return "parse_details"
    else:
        return "parse_details"

def route_after_problems(state: DeveloperQueryState):
    """Route after problem analysis - check if user wants solutions"""
    raw_query = state["raw_query"].lower()
    query_type = state["query_type"]
    
    # Keywords that indicate user wants solutions
    solution_keywords = ["solution", "solve", "fix", "how to", "resolve", "address", "recommend"]
    
    # Check if query asks for solutions
    wants_solutions = any(keyword in raw_query for keyword in solution_keywords)
    
    # Always provide solutions for these query types
    auto_solution_types = ["version_issues", "aspect_issues", "version_aspect_issues", "general_analysis"]
    
    if wants_solutions or query_type in auto_solution_types:
        print("🔀 Routing to solutions agent")
        return "solutions"
    else:
        print("🔀 User only wants problems - routing to coordinator")
        return "coordinator"

def route_after_parsing(state: DeveloperQueryState):
    """Route after parsing - cleaner logic"""
    query_type = state["query_type"]
    
    if query_type == "direct_absa":
        return "absa"
    elif query_type == "solution_only":
        return "solutions"
    else:
        return "scraper"

def route_after_scraping(state: DeveloperQueryState):
    """Route after review scraping"""
    reviews = state.get("reviews", [])
    query_type = state["query_type"]
    
    if not reviews:
        return "coordinator"
    elif query_type == "reviews_request":
        return "context"
    else:
        return "context"

def route_after_context(state: DeveloperQueryState):
    """Route after context filtering"""
    filtered_reviews = state.get("filtered_reviews", [])
    query_type = state["query_type"]
    
    if not filtered_reviews or query_type == "reviews_request":
        return "coordinator"
    else:
        return "absa"

def route_after_absa(state: DeveloperQueryState):
    """Route after ABSA analysis"""
    query_type = state["query_type"]
    
    if query_type in ["direct_absa", "aspect_opinion_query"]:
        return "coordinator"
    else:
        return "problems"

# =============================================================================
# GRAPH BUILDING
# =============================================================================

def build_developer_agent_graph():
    """Build the simplified and more logical graph"""
    
    graph_builder = StateGraph(DeveloperQueryState)
    
    # Add nodes
    graph_builder.add_node("classifier", classify_query)
    graph_builder.add_node("parse_details", parse_query_details)
    graph_builder.add_node("scraper", review_scraper_agent)
    graph_builder.add_node("context", context_agent)
    graph_builder.add_node("absa", absa_model_agent)
    graph_builder.add_node("problems", problem_analysis_agent)
    graph_builder.add_node("solutions", solutions_agent)
    graph_builder.add_node("coordinator", coordinator_agent)
    
    # Build flow
    graph_builder.add_edge(START, "classifier")
    
    graph_builder.add_conditional_edges(
        "classifier",
        route_after_classification,
        {
            "coordinator": "coordinator",
            "parse_details": "parse_details"
        }
    )
    
    graph_builder.add_conditional_edges(
        "parse_details",
        route_after_parsing,
        {
            "absa": "absa",
            "solutions": "solutions",
            "scraper": "scraper"
        }
    )
    
    graph_builder.add_conditional_edges(
        "scraper",
        route_after_scraping,
        {"coordinator": "coordinator", "context": "context"}
    )
    
    graph_builder.add_conditional_edges(
        "context",
        route_after_context,
        {"coordinator": "coordinator", "absa": "absa"}
    )
    
    graph_builder.add_conditional_edges(
        "absa",
        route_after_absa,
        {
            "coordinator": "coordinator",
            "problems": "problems"
        }
    )
    
    graph_builder.add_conditional_edges(
        "problems",
        route_after_problems,
        {"solutions": "solutions", "coordinator": "coordinator"}
    )
    
    graph_builder.add_edge("solutions", "coordinator")
    graph_builder.add_edge("coordinator", END)
    
    return graph_builder.compile()

# =============================================================================
# MAIN EXECUTION FUNCTIONS
# =============================================================================

def run_developer_chatbot():
    """Run the conversational developer assistant with comprehensive logging"""
    
    # Initialize logging
    log_file = "system_logs.txt"
    session_start = time.time()
    
    # Log session start
    log_to_file("="*80, log_file)
    log_to_file("🚀 NEW DEVELOPER ASSISTANT SESSION STARTED", log_file)
    log_to_file("="*80, log_file)
    
    print("🚀 Starting Developer Assistant...")
    log_to_file("System: Starting Developer Assistant initialization", log_file)
    
    # Initialize graph
    graph_init_start = time.time()
    graph = build_developer_agent_graph()
    graph_init_time = time.time() - graph_init_start
    
    log_to_file(f"System: Graph initialization completed in {graph_init_time:.2f} seconds", log_file)
    
    print("🤖 Developer Assistant Ready! Ask me about app issues, versions, or solutions.")
    print("Type 'exit' to quit.\n")
    
    log_to_file("System: Developer Assistant ready for user queries", log_file)
    
    query_count = 0
    
    while True:
        user_input = input("Developer: ")
        
        if user_input.lower() == "exit":
            print("👋 Goodbye!")
            
            # Calculate total session time
            total_session_time = time.time() - session_start
            
            # Log session end
            log_to_file("-"*60, log_file)
            log_to_file(f"📊 SESSION SUMMARY", log_file)
            log_to_file(f"Total Queries Processed: {query_count}", log_file)
            log_to_file(f"Total Session Time: {total_session_time:.2f} seconds", log_file)
            log_to_file(f"Average Query Time: {total_session_time/query_count:.2f} seconds" if query_count > 0 else "No queries processed", log_file)
            log_to_file("👋 Session ended by user", log_file)
            log_to_file("="*80, log_file)
            break
        
        query_count += 1
        query_start_time = time.time()
        
        # Log user query
        log_to_file(f"👤 USER QUERY #{query_count}: {user_input}", log_file)
        
        # Initialize state
        state = {
            "messages": [{"role": "user", "content": user_input}],
            "query_type": None,
            "raw_query": None,
            "aspect": None,
            "version": None,
            "app_ids": None,
            "reviews": None,
            "app_info": None,
            "recent_changes": None,
            "filtered_reviews": None,
            "absa_results": None,
            "negative_aspects": None,
            "problems": None,
            "solutions": None,
            "final_response": None,
            "review_text": None,
            "conversation_context": None
        }
        
        try:
            # Log processing start
            log_to_file(f"🔄 Starting graph processing for query #{query_count}", log_file)
            
            # Run the graph with timing
            graph_start = time.time()
            result = graph.invoke(state)
            graph_processing_time = time.time() - graph_start
            
            # Extract response
            final_response = result.get("final_response", "Sorry, I couldn't process your request.")
            
            # Calculate total query time
            total_query_time = time.time() - query_start_time
            
            # Log processing results
            log_to_file(f"✅ Graph processing completed in {graph_processing_time:.2f} seconds", log_file)
            log_to_file(f"📊 QUERY #{query_count} RESULTS:", log_file)
            log_to_file(f"  - Query Type: {result.get('query_type', 'Unknown')}", log_file)
            log_to_file(f"  - Reviews Found: {len(result.get('reviews', []))}", log_file)
            log_to_file(f"  - Filtered Reviews: {len(result.get('filtered_reviews', []))}", log_file)
            log_to_file(f"  - ABSA Results: {len(result.get('absa_results', []))}", log_file)
            log_to_file(f"  - Problems Found: {len(result.get('problems', []))}", log_file)
            log_to_file(f"  - Solutions Generated: {len(result.get('solutions', []))}", log_file)
            log_to_file(f"  - Graph Processing Time: {graph_processing_time:.2f}s", log_file)
            log_to_file(f"  - Total Query Time: {total_query_time:.2f}s", log_file)
            
            # Log response summary (first 200 chars)
            response_summary = final_response[:200] + "..." if len(final_response) > 200 else final_response
            log_to_file(f"🤖 ASSISTANT RESPONSE SUMMARY: {response_summary}", log_file)
            
            # Print response to user
            print(f"\n🤖 Assistant:\n{final_response}\n")
            print(f"⏱️ Processing Time: {total_query_time:.2f} seconds")
            print("-" * 60)
            
            log_to_file(f"✅ Query #{query_count} completed successfully", log_file)
            log_to_file("-"*40, log_file)
            
        except Exception as e:
            error_time = time.time() - query_start_time
            
            # Log error
            log_to_file(f"❌ ERROR in query #{query_count}: {str(e)}", log_file)
            log_to_file(f"  - Error occurred after {error_time:.2f} seconds", log_file)
            log_to_file(f"  - State at error: {str(state)}", log_file)
            
            print(f"❌ Error: {e}")
            print(f"⏱️ Error occurred after: {error_time:.2f} seconds")
            
            log_to_file("-"*40, log_file)

def run_performance_test():
    """Run performance tests with detailed logging"""
    
    log_file = "performance_test_logs.txt"
    
    # Test queries for different workflows
    test_queries = [
        "analyze this review: this is a good app but not good at learning",  # direct_absa
        "what are the aspects and opinions in the reviews",  # aspect_opinion_query  
        "what are the possible solutions of the login issues",  # solution_only
        "what are the issues in version 3.9.8.4",  # full workflow
        "give me the reviews of version 3.9.8.3",  # reviews_request
        "show me the problems and possible solutions of app latest"  # general_analysis
    ]
    
    log_to_file("="*80, log_file)
    log_to_file("🧪 PERFORMANCE TEST STARTED", log_file)
    log_to_file("="*80, log_file)
    
    print("🧪 Starting Performance Test...")
    
    # Initialize graph
    graph_init_start = time.time()
    graph = build_developer_agent_graph()
    graph_init_time = time.time() - graph_init_start
    
    log_to_file(f"Graph initialization: {graph_init_time:.2f} seconds", log_file)
    
    total_test_start = time.time()
    
    for i, query in enumerate(test_queries, 1):
        print(f"\n🧪 Test {i}/{len(test_queries)}: {query}")
        log_to_file(f"🧪 TEST {i}: {query}", log_file)
        
        query_start = time.time()
        
        # Initialize state
        state = {
            "messages": [{"role": "user", "content": query}],
            "query_type": None,
            "raw_query": None,
            "aspect": None,
            "version": None,
            "app_ids": None,
            "reviews": None,
            "app_info": None,
            "recent_changes": None,
            "filtered_reviews": None,
            "absa_results": None,
            "negative_aspects": None,
            "problems": None,
            "solutions": None,
            "final_response": None,
            "review_text": None,
            "conversation_context": None
        }
        
        try:
            result = graph.invoke(state)
            query_time = time.time() - query_start
            
            log_to_file(f"✅ Test {i} completed in {query_time:.2f} seconds", log_file)
            log_to_file(f"  - Query Type: {result.get('query_type')}", log_file)
            log_to_file(f"  - Reviews: {len(result.get('reviews', []))}", log_file)
            log_to_file(f"  - Problems: {len(result.get('problems', []))}", log_file)
            
            print(f"✅ Completed in {query_time:.2f} seconds")
            
        except Exception as e:
            error_time = time.time() - query_start
            log_to_file(f"❌ Test {i} failed after {error_time:.2f} seconds: {e}", log_file)
            print(f"❌ Failed: {e}")
    
    total_test_time = time.time() - total_test_start
    
    log_to_file("-"*60, log_file)
    log_to_file(f"📊 PERFORMANCE TEST SUMMARY", log_file)
    log_to_file(f"Total Tests: {len(test_queries)}", log_file)
    log_to_file(f"Total Test Time: {total_test_time:.2f} seconds", log_file)
    log_to_file(f"Average Query Time: {total_test_time/len(test_queries):.2f} seconds", log_file)
    log_to_file(f"Graph Init Time: {graph_init_time:.2f} seconds", log_file)
    log_to_file("="*80, log_file)
    
    print(f"\n📊 Performance Test Complete!")
    print(f"Total Time: {total_test_time:.2f} seconds")
    print(f"Average per query: {total_test_time/len(test_queries):.2f} seconds")
    print(f"Check {log_file} for detailed results")

# =============================================================================
# MAIN EXECUTION
# =============================================================================

if __name__ == "__main__":
    # Choose mode
    mode = input("Choose mode - (1) Interactive Chat (2) Performance Test: ").strip()
    
    if mode == "2":
        run_performance_test()
    else:
        # Run interactive chatbot
        run_developer_chatbot()
        
    # Clean up
    try:
        from model_loader import model_manager
        model_manager.unload_model()
    except:
        pass