# Groq AI Helper for Natural Language Understanding with Structured Output

from groq import Groq
from pydantic import BaseModel, Field
from typing import Optional, Literal
import re
import json
from voice_config import GROQ_API_KEY, USE_GROQ_AI


# Pydantic Models for Structured Output
class RequestTypeResponse(BaseModel):
    intent: Literal["exchange", "return"] = Field(description="Whether customer wants exchange or return")
    confidence: Literal["high", "medium", "low"] = Field(description="Confidence level")
    reasoning: str = Field(description="Brief explanation")


class OrderIdResponse(BaseModel):
    order_id: Optional[str] = Field(description="Order ID extracted, null if not found")
    found: bool = Field(description="Whether order ID was found")
    interpretation: str = Field(description="How it was interpreted")


class ReasonResponse(BaseModel):
    reason: str = Field(description="Professional reason in 2nd person (You...), max 10 words")
    category: Literal["size", "defect", "wrong_item", "preference", "delivery", "other"] = Field(description="Issue category")


class PreferenceResponse(BaseModel):
    preference: str = Field(description="Product description, max 15 words")
    type: Literal["size", "color", "model", "feature", "other"] = Field(description="Preference type")


class ConfirmationResponse(BaseModel):
    confirmed: bool = Field(description="True if confirming, False if declining")
    confidence: Literal["high", "medium", "low"] = Field(description="Confidence level")


class CorrectionField(BaseModel):
    field: Literal["request_type", "order_id", "reason", "exchange_preference", "everything"] = Field(description="Which field the customer wants to correct")
    reasoning: str = Field(description="Why this field needs correction")


class GroqAI:
    def __init__(self):
        print(f"\n{'='*60}")
        print("🔧 INITIALIZING GROQ AI")
        print(f"USE_GROQ_AI = {USE_GROQ_AI}")
        print(f"GROQ_API_KEY = {'SET ('+str(len(GROQ_API_KEY))+' chars)' if GROQ_API_KEY else 'NOT SET'}")
        print(f"{'='*60}\n")
        
        self.enabled = USE_GROQ_AI and GROQ_API_KEY
        if self.enabled:
            try:
                self.client = Groq(api_key=GROQ_API_KEY)
                self.model_name = "llama-3.3-70b-versatile"  # Fast and powerful model
                print("✅ Groq AI initialized successfully")
            except Exception as e:
                print(f"⚠️ Groq AI initialization failed: {e}")
                self.enabled = False
        else:
            print("⚠️ Groq AI is DISABLED - will use basic pattern matching")
            if not USE_GROQ_AI:
                print("   Reason: USE_GROQ_AI is False")
            if not GROQ_API_KEY:
                print("   Reason: GROQ_API_KEY is not set")
    
    def extract_request_type(self, user_text, question_asked="Would you like exchange or return?"):
        print(f"\n{'='*60}")
        print(f"🎯 EXTRACT_REQUEST_TYPE")
        print(f"📝 User Input (raw): '{user_text}'")
        print(f"📝 User Input (repr): {repr(user_text)}")
        print(f"📝 Input length: {len(user_text)} chars")
        print(f"❓ Question: {question_asked}")
        print(f"🔌 Groq Enabled: {self.enabled}")
        print(f"{'='*60}")
        
        # Clean input
        user_text_clean = user_text.strip().lower()

        # Fast path: deterministic keyword detection to avoid model mistakes
        strong_return_keywords = [
            'return', 'refund', 'money back', 'send back', 'want to return',
            'need to return', 'request a return', 'two kind of return'
        ]
        strong_exchange_keywords = [
            'exchange', 'swap', 'replace', 'replacement', 'change it', 'different product'
        ]

        if any(keyword in user_text_clean for keyword in strong_return_keywords):
            print("   ✅ FAST PATH: RETURN detected via keyword match")
            return 'return'
        if any(keyword in user_text_clean for keyword in strong_exchange_keywords):
            print("   ✅ FAST PATH: EXCHANGE detected via keyword match")
            return 'exchange'
        
        # Try Groq AI first if enabled
        if self.enabled:
            try:
                prompt = (
                    f"You are analyzing customer service dialogue.\n\n"
                    f"QUESTION: {question_asked}\n"
                    f"CUSTOMER SAYS: {user_text}\n\n"
                    f"TASK: Determine what the customer wants.\n\n"
                    f"YOU MUST CHOOSE EXACTLY ONE:\n"
                    f"A) 'exchange' - if they want to swap/replace/exchange the product\n"
                    f"B) 'return' - if they want to send back/refund/return the product\n\n"
                    f"EXAMPLES:\n"
                    f"'exchange' → exchange\n"
                    f"'return' → return\n"
                    f"'I want to exchange this' → exchange\n"
                    f"'can I get a refund' → return\n"
                    f"'I would like you to process an exchange' → exchange\n"
                    f"'swap it for another one' → exchange\n\n"
                    f"Look for these keywords:\n"
                    f"- exchange, swap, replace, change, different = EXCHANGE\n"
                    f"- return, refund, send back, money back = RETURN\n\n"
                    f"RESPOND WITH JSON ONLY:\n"
                    f'{{"intent": "exchange", "confidence": "high", "reasoning": "..."}}\n'
                    f'OR\n'
                    f'{{"intent": "return", "confidence": "high", "reasoning": "..."}}\n\n'
                    f"The 'intent' MUST be EXACTLY either 'exchange' or 'return' (lowercase).\n"
                    f"Choose the BEST match even if you're uncertain. NEVER return anything other than 'exchange' or 'return'."
                )
                
                response = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=[{"role": "user", "content": prompt}],
                    response_format={"type": "json_object"},
                    temperature=0.0,  # More deterministic
                )
                
                json_text = response.choices[0].message.content
                print(f"[LLM_OUTPUT request_type] {json_text}")
                result = RequestTypeResponse.model_validate_json(json_text)
                print(f"✅ Groq Result: {result.intent} ({result.confidence}) - {result.reasoning}")
                
                # Validate result is one of the two options
                if result.intent in ['exchange', 'return']:
                    return result.intent
                else:
                    print(f"⚠️ Invalid intent from Groq: {result.intent}, falling back")
                    
            except Exception as e:
                print(f"❌ Groq Error: {e}")
                print(f"❌ Raw response: {json_text if 'json_text' in locals() else 'No response'}")
                # ensure downstream always gets a value
                return 'exchange'
        
        # Fallback: Basic keyword matching (ALWAYS returns a result)
        print(f"\n🔍 BASIC KEYWORD MATCHING")
        print(f"   Cleaned input: '{user_text_clean}'")
        
        # Check for return keywords first
        return_keywords = ['return', 'refund', 'money back', 'send back', 'don\'t want', 'cancel']
        for keyword in return_keywords:
            if keyword in user_text_clean:
                print(f"   ✅ RETURN detected (keyword: '{keyword}')")
                return 'return'
        
        # Check for exchange keywords
        exchange_keywords = ['exchange', 'swap', 'replace', 'change', 'different']
        for keyword in exchange_keywords:
            if keyword in user_text_clean:
                print(f"   ✅ EXCHANGE detected (keyword: '{keyword}')")
                return 'exchange'
        
        # If no keywords found, make best guess based on partial matches
        print(f"   ⚠️ No exact keywords found, checking partial matches...")
        
        if any(word in user_text_clean for word in ['exch', 'swp', 'replac']):
            print(f"   ✅ EXCHANGE detected (partial match)")
            return 'exchange'
        
        if any(word in user_text_clean for word in ['retur', 'refun', 'back']):
            print(f"   ✅ RETURN detected (partial match)")
            return 'return'
        
        # Absolute last resort: default to exchange
        print(f"   ⚠️ No matches found, defaulting to EXCHANGE")
        return 'exchange'
    
    def extract_order_id(self, user_text, question_asked="Please provide your order ID"):
        print(f"\n🎯 EXTRACT_ORDER_ID: {user_text}")
        
        if not self.enabled:
            return self._basic_extract_order_id(user_text)
        
        try:
            prompt = (
                f"Question: {question_asked}\n"
                f"Response: {user_text}\n\n"
                "Extract any order number/ID from the customer's response.\n"
                "Numbers can be spoken as words (e.g., 'one two three' = 123).\n"
                "Look for patterns like: 'order 123', 'order number ABC', 'ORD-456', etc.\n\n"
                "You MUST respond with valid JSON in this exact format:\n"
                '{"order_id": "12345", "found": true, "interpretation": "Extracted from order 12345"}\n'
                'OR if no order ID found:\n'
                '{"order_id": null, "found": false, "interpretation": "No order ID mentioned"}'
            )
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"},
                temperature=0.1,
            )
            
            json_text = response.choices[0].message.content
            print(f"🔍 Groq JSON Response: {json_text}")
            result = OrderIdResponse.model_validate_json(json_text)
            
            if result.found and result.order_id:
                order_id = re.sub(r'\D', '', result.order_id)
                if order_id:
                    print(f"✅ Order ID: {order_id}")
                    return order_id
            
            return self._basic_extract_order_id(user_text)
                
        except Exception as e:
            print(f"❌ Groq Error: {e}")
            print(f"❌ Raw response: {json_text if 'json_text' in locals() else 'No response'}")
            return self._basic_extract_order_id(user_text)
    
    def extract_reason(self, user_text, request_type, question_asked=None):
        if not self.enabled:
            return user_text
        
        if not question_asked:
            question_asked = f"Why do you want to {request_type}?"
        
        try:
            prompt = (
                f"Question: {question_asked}\n"
                f"Response: {user_text}\n\n"
                "Convert their reason to professional format (max 10 words).\n\n"
                "CRITICAL: Convert 1st person to 2nd person:\n"
                "- 'I don't like' → 'You didn't like'\n"
                "- 'I changed my mind' → 'You changed your mind'\n"
                "- 'It's too small' → 'It's too small' (already neutral)\n"
                "- 'I ordered wrong size' → 'You ordered wrong size'\n\n"
                "Categorize as: size, defect, wrong_item, preference, delivery, other\n\n"
                "You MUST respond with valid JSON in this exact format:\n"
                '{"reason": "You didn\'t like the product", "category": "preference"}\n'
                'OR\n'
                '{"reason": "Wrong size ordered", "category": "size"}\n\n'
                "The 'category' MUST be one of: size, defect, wrong_item, preference, delivery, other"
            )
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"},
                temperature=0.1,
            )
            
            json_text = response.choices[0].message.content
            print(f"🔍 Groq JSON Response: {json_text}")
            result = ReasonResponse.model_validate_json(json_text)
            print(f"✅ Reason: {result.reason} ({result.category})")
            return result.reason
                
        except Exception as e:
            print(f"❌ Groq Error: {e}")
            print(f"❌ Raw response: {json_text if 'json_text' in locals() else 'No response'}")
            return user_text
    
    def extract_exchange_preference(self, user_text, question_asked="What would you prefer instead?"):
        if not self.enabled:
            return user_text
        
        try:
            prompt = (
                f"Question: {question_asked}\n"
                f"Response: {user_text}\n\n"
                "Extract what product they want (max 15 words).\n"
                "Categorize type as: size, color, model, feature, other\n\n"
                "You MUST respond with valid JSON in this exact format:\n"
                '{"preference": "Size large in black", "type": "size"}\n'
                'OR\n'
                '{"preference": "Different color - blue", "type": "color"}\n\n'
                "The 'type' MUST be one of: size, color, model, feature, other"
            )
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"},
                temperature=0.1,
            )
            
            json_text = response.choices[0].message.content
            print(f"🔍 Groq JSON Response: {json_text}")
            result = PreferenceResponse.model_validate_json(json_text)
            print(f"✅ Preference: {result.preference} ({result.type})")
            return result.preference
                
        except Exception as e:
            print(f"❌ Groq Error: {e}")
            print(f"❌ Raw response: {json_text if 'json_text' in locals() else 'No response'}")
            return user_text
    
    def is_confirmation(self, user_text, question_asked="Is this correct?"):
        if not self.enabled:
            return self._basic_is_confirmation(user_text)
        
        try:
            prompt = (
                f"Question: {question_asked}\n"
                f"Response: {user_text}\n\n"
                "Is the customer confirming (yes) or declining (no)?\n\n"
                "AFFIRMATIVE (confirming): yes, yeah, yep, sure, ok, okay, correct, right, go ahead, proceed, confirm\n"
                "NEGATIVE (declining): no, nope, nah, wrong, incorrect, cancel, stop, wait\n\n"
                "You MUST respond with valid JSON in this exact format:\n"
                '{"confirmed": true, "confidence": "high"}\n'
                'OR\n'
                '{"confirmed": false, "confidence": "high"}\n\n'
                "The 'confirmed' field MUST be boolean (true or false).\n"
                "The 'confidence' MUST be 'high', 'medium', or 'low'."
            )
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"},
                temperature=0.1,
            )
            
            json_text = response.choices[0].message.content
            print(f"🔍 Groq JSON Response: {json_text}")
            result = ConfirmationResponse.model_validate_json(json_text)
            print(f"✅ Confirmed: {result.confirmed} ({result.confidence})")
            return result.confirmed
                
        except Exception as e:
            print(f"❌ Groq Error: {e}")
            print(f"❌ Raw response: {json_text if 'json_text' in locals() else 'No response'}")
            return self._basic_is_confirmation(user_text)
    
    def _basic_extract_request_type(self, text):
        """Legacy method - not used anymore, kept for compatibility"""
        text_lower = text.lower().strip()
        
        return_keywords = ['return', 'refund', 'money back', 'send back']
        for keyword in return_keywords:
            if keyword in text_lower:
                return 'return'
        
        exchange_keywords = ['exchange', 'swap', 'replace', 'change', 'different']
        for keyword in exchange_keywords:
            if keyword in text_lower:
                return 'exchange'
        
        # Default to exchange if nothing found
        return 'exchange'
    
    def _basic_extract_order_id(self, text):
        print(f"🔍 Basic order ID extraction: {text}")
        match = re.search(r'order[\s#]*(\d+)|(\d+)', text.lower())
        if match:
            order_id = match.group(1) or match.group(2)
            print(f"✅ Found: {order_id}")
            return order_id
        print("❌ No order ID found")
        return None
    
    def _basic_is_confirmation(self, text):
        confirmations = ['yes', 'yeah', 'yep', 'correct', 'right', 'sure', 'ok', 'okay']
        return any(word in text.lower() for word in confirmations)


    def identify_correction_field(self, user_text, current_data):
        """When user says no to confirmation, identify which field they want to correct"""
        if not self.enabled:
            return "everything"
        
        try:
            confirmation_summary = (
                f"Request type: {current_data.get('request_type', 'N/A')}\n"
                f"Order ID: {current_data.get('order_id', 'N/A')}\n"
                f"Reason: {current_data.get('reason', 'N/A')}\n"
            )
            if current_data.get('request_type') == 'exchange':
                confirmation_summary += f"Exchange preference: {current_data.get('exchange_preference', 'N/A')}\n"
            
            prompt = (
                f"The customer declined the confirmation. Here's what we have:\n"
                f"{confirmation_summary}\n"
                f"Customer said: {user_text}\n\n"
                "Which field do they want to correct?\n\n"
                "Options: request_type, order_id, reason, exchange_preference, everything\n\n"
                "Examples:\n"
                "- 'wrong order number' = order_id\n"
                "- 'I want exchange not return' = request_type\n"
                "- 'reason is wrong' = reason\n"
                "- 'start over' = everything\n\n"
                "You MUST respond with valid JSON in this exact format:\n"
                '{\"field\": \"order_id\", \"reasoning\": \"Customer mentioned wrong order number\"}\n'
                'OR\n'
                '{\"field\": \"everything\", \"reasoning\": \"Customer wants to start over\"}\n\n'
                "The 'field' MUST be one of: request_type, order_id, reason, exchange_preference, everything"
            )
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"},
                temperature=0.1,
            )
            
            json_text = response.choices[0].message.content
            print(f"🔍 Groq JSON Response: {json_text}")
            result = CorrectionField.model_validate_json(json_text)
            print(f"✅ Correction needed for: {result.field} - {result.reasoning}")
            return result.field
                
        except Exception as e:
            print(f"❌ Groq Error: {e}")
            print(f"❌ Raw response: {json_text if 'json_text' in locals() else 'No response'}")
            return "everything"


# Global instance
groq_ai = GroqAI()

# Backwards compatibility alias
gemini_ai = groq_ai