# query_classifier.py
import google.generativeai as genai
import json
import os
from dotenv import load_dotenv

load_dotenv()
genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))

class QueryClassifier:
    def __init__(self):
        self.model = genai.GenerativeModel('gemini-2.5-flash-lite')
        self.process_mapping = {
            'legal_aid': 'Legal Aid Application',
            'consumer_complaint': 'Consumer Complaint Filing', 
            'court_status': 'Court Case Status Check',
            'general': 'General Legal Query'
        }
        
    def classify_query(self, user_query: str) -> dict:
        """Classify user query to determine legal process and current step"""
        
        classification_prompt = f"""
        Analyze this user query and determine:
        1. Which legal process they're asking about
        2. What step they might be at in the process
        3. Whether they need step-by-step guidance or just information

        Query: "{user_query}"

        Available processes:
        - legal_aid: Questions about getting free legal help
        - consumer_complaint: Filing complaints about products/services
        - court_status: Checking case status or court procedures
        - general: General legal questions not fitting above categories

        Possible steps for legal_aid:
        - check_eligibility: Asking about who qualifies
        - gather_documents: Asking what documents needed
        - choose_application_method: How to apply
        - online_application: Help with online form
        - wait_processing: Application submitted, waiting for response
        - contact_lawyer: Have lawyer assigned, need next steps

        Return ONLY a JSON object with this format:
        {{
            "process": "process_name",
            "step": "step_id_or_null",
            "intent": "process_guidance|information_only",
            "confidence": 0.8,
            "reasoning": "Brief explanation"
        }}
        """
        
        try:
            response = self.model.generate_content(classification_prompt)
            # Extract JSON from response
            response_text = response.text.strip()
            if response_text.startswith('```'):
                response_text = response_text[7:-3].strip()
            elif response_text.startswith('```'):
                response_text = response_text[3:-3].strip()
                
            result = json.loads(response_text)
            
            # Validate result
            if result.get('process') not in self.process_mapping:
                result['process'] = 'general'
                
            result['process_name'] = self.process_mapping.get(result['process'])
            
            return result
            
        except Exception as e:
            print(f"Classification error: {e}")
            return {
                "process": "general",
                "step": None,
                "intent": "information_only", 
                "confidence": 0.3,
                "reasoning": "Classification failed, defaulting to general",
                "process_name": "General Legal Query"
            }

# Test the classifier
if __name__ == "__main__":
    classifier = QueryClassifier()
    
    test_queries = [
        "I need free legal help but don't know if I qualify",
        "My landlord is trying to evict me illegally", 
        "How do I check my court case status online?",
        "I bought a defective phone and the company won't refund"
    ]
    
    for query in test_queries:
        result = classifier.classify_query(query)
        print(f"\nQuery: {query}")
        print(f"Classification: {json.dumps(result, indent=2)}")