from langchain_ollama import OllamaLLM
import json
import re

def infer_review_with_local_model(review, model_name="llama3.1"):
    """Production ABSA using Ollama - fully dynamic, no hardcoded aspects"""
    try:
        llm = OllamaLLM(model=model_name)
        
        # Dynamic prompt - no hardcoded aspects
        prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
You are an expert at Aspect-Based Sentiment Analysis (ABSA). Analyze the following review and extract ALL aspects mentioned with their sentiments and opinions. 

Instructions:
1. Identify ANY aspect mentioned in the review (don't limit to predefined categories)
2. For each aspect, determine sentiment: "positive", "negative", or "neutral"  
3. Extract the specific opinion expressed about that aspect
4. Use concise, descriptive aspect names
5. Return as JSON array with "aspect", "sentiment", and "opinion" fields
6. If no specific aspects found, analyze the overall review as "general"

### Input:
{review}

### Response:
"""

        response = llm.invoke(prompt)
        
        # Parse JSON response
        try:
            # Extract JSON array from response
            json_match = re.search(r'\[.*?\]', response, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                result = json.loads(json_str)
                
                # Validate structure only - no content filtering
                if isinstance(result, list):
                    cleaned_result = []
                    for item in result:
                        if isinstance(item, dict) and all(k in item for k in ["aspect", "sentiment", "opinion"]):
                            # Only validate sentiment values
                            sentiment = item["sentiment"].lower().strip()
                            if sentiment not in ["positive", "negative", "neutral"]:
                                sentiment = "neutral"
                            
                            cleaned_result.append({
                                "aspect": item["aspect"].strip(),
                                "sentiment": sentiment,
                                "opinion": item["opinion"].strip()
                            })
                    
                    if cleaned_result:
                        print(f"🔍 ABSA Result: {cleaned_result}")
                        return cleaned_result
            
            # Fallback to LLM-based parsing
            return llm_fallback_parsing(review, llm)
            
        except Exception as parse_error:
            print(f"JSON parsing failed: {parse_error}")
            return llm_fallback_parsing(review, llm)
            
    except Exception as e:
        print(f"❌ Ollama ABSA error: {e}")
        return [{
            "aspect": "general",
            "sentiment": "neutral",
            "opinion": f"Analysis failed: {str(e)[:50]}"
        }]

def llm_fallback_parsing(review, llm):
    """Use LLM for fallback parsing when JSON extraction fails"""
    try:
        fallback_prompt = f"""Analyze this review and identify aspects with sentiments. 
        
Review: "{review}"

Please identify:
1. What specific things are mentioned (aspects)
2. Whether each is viewed positively, negatively, or neutrally
3. The opinion expressed

Format your response as:
Aspect: [aspect_name] | Sentiment: [positive/negative/neutral] | Opinion: [specific_opinion]

For each aspect found, use the format above."""

        response = llm.invoke(fallback_prompt)
        
        # Parse the structured response
        aspects_found = []
        lines = response.split('\n')
        
        for line in lines:
            line = line.strip()
            if 'Aspect:' in line and 'Sentiment:' in line and 'Opinion:' in line:
                try:
                    # Parse each component
                    parts = line.split(' | ')
                    if len(parts) >= 3:
                        aspect = parts[0].replace('Aspect:', '').strip()
                        sentiment_part = parts[1].replace('Sentiment:', '').strip().lower()
                        opinion = parts[2].replace('Opinion:', '').strip()
                        
                        # Validate sentiment
                        if sentiment_part not in ["positive", "negative", "neutral"]:
                            sentiment_part = "neutral"
                        
                        if aspect and opinion:
                            aspects_found.append({
                                "aspect": aspect,
                                "sentiment": sentiment_part,
                                "opinion": opinion
                            })
                except Exception as parse_err:
                    continue
        
        # If still no aspects found, use general analysis
        if not aspects_found:
            return general_sentiment_analysis(review, llm)
        
        return aspects_found
        
    except Exception as e:
        print(f"Fallback parsing error: {e}")
        return general_sentiment_analysis(review, llm)

def general_sentiment_analysis(review, llm):
    """General sentiment analysis when no specific aspects found"""
    try:
        sentiment_prompt = f"""Analyze the overall sentiment of this review:

"{review}"

Respond with only one word: positive, negative, or neutral"""

        sentiment_response = llm.invoke(sentiment_prompt)
        sentiment = sentiment_response.strip().lower()
        
        if sentiment not in ["positive", "negative", "neutral"]:
            sentiment = "neutral"
        
        return [{
            "aspect": "general",
            "sentiment": sentiment,
            "opinion": review[:150] + "..." if len(review) > 150 else review
        }]
        
    except Exception as e:
        return [{
            "aspect": "general", 
            "sentiment": "neutral",
            "opinion": "Could not analyze sentiment"
        }]

def initialize_absa_model():
    """Initialize and test Ollama ABSA"""
    try:
        llm = OllamaLLM(model="llama3.1")
        # Test with a simple prompt
        test_response = llm.invoke("Test: return 'OK'")
        print("✅ Ollama ABSA model ready!")
        return True
    except Exception as e:
        print(f"❌ Ollama initialization failed: {e}")
        return False

# Mock model manager for compatibility
class ABSAModelManager:
    def load_model(self):
        return initialize_absa_model()
    
    def get_model(self):
        return True, True
    
    def unload_model(self):
        print("✅ Ollama ABSA session closed")

# Create global instance
model_manager = ABSAModelManager()

def test_absa():
    """Test the dynamic ABSA system"""
    test_reviews = [
        "The login system crashes frequently but I love the beautiful user interface design",
        "Payment processing is slow and the app freezes during checkout", 
        "Great app! Fast performance and excellent content quality",
        "The camera quality is amazing but the battery drains too quickly",
        "Love the new dark mode feature, makes reading so much easier",
        "The notification sounds are too loud and the settings menu is confusing"
    ]
    
    print("🧪 Testing Dynamic Ollama ABSA System...")
    for i, review in enumerate(test_reviews, 1):
        print(f"\n{i}. Review: {review}")
        result = infer_review_with_local_model(review)
        print(f"   Result: {json.dumps(result, indent=6)}")

if __name__ == "__main__":
    test_absa()