import re
from tools import lookup_order, last_orders

def extract_order_mobile_with_llm(message, text_generator=None):
    """Extract order ID or mobile number using LLM and validate"""
    extracted_info = {"type": None, "value": None, "valid": False, "message": ""}
    
    if text_generator:
        try:
            prompt = f"Extract: {message}\nNumber:"
            response = text_generator(
                prompt, 
                max_new_tokens=15, 
                do_sample=False, 
                pad_token_id=50256,
                eos_token_id=50256,
                return_full_text=False
            )
            generated = response[0]['generated_text'].strip()
            
            numbers = re.findall(r'\d+', generated)
            for num in numbers:
                if len(num) == 6:
                    extracted_info["type"] = "order_id"
                    extracted_info["value"] = num
                    break
                elif len(num) == 10:
                    extracted_info["type"] = "mobile"
                    extracted_info["value"] = num
                    break
        except:
            pass
    
    if not extracted_info["value"]:
        order_match = re.search(r'\b\d{6}\b', message)
        mobile_match = re.search(r'\b\d{10}\b', message)
        
        if order_match:
            extracted_info["type"] = "order_id"
            extracted_info["value"] = order_match.group()
        elif mobile_match:
            extracted_info["type"] = "mobile"
            extracted_info["value"] = mobile_match.group()
    
    if extracted_info["value"]:
        if extracted_info["type"] == "order_id":
            result = lookup_order(extracted_info["value"])
            if result["status"] == "success":
                extracted_info["valid"] = True
                extracted_info["data"] = result
            else:
                extracted_info["message"] = "I couldn't find an order with that order ID. Please enter the correct 6-digit order ID to proceed."
        
        elif extracted_info["type"] == "mobile":
            result = last_orders(extracted_info["value"])
            if result["status"] == "success":
                extracted_info["valid"] = True
                extracted_info["data"] = result
            else:
                extracted_info["message"] = "I couldn't find any orders for that mobile number. Please enter the correct 10-digit mobile number to proceed."
    
    return extracted_info

def generate_refund_question(context, customer_message, text_generator=None):
    """Generate contextual refund questions using LLM"""
    
    if text_generator:
        try:
            if context == "question1":
                prompt = f"Customer wants refund: {customer_message}\nAsk when issue started:"
                response = text_generator(
                    prompt, 
                    max_new_tokens=20, 
                    do_sample=False, 
                    pad_token_id=50256,
                    eos_token_id=50256,
                    return_full_text=False
                )
                generated = response[0]['generated_text'].strip()
                if len(generated) > 10 and '?' in generated:
                    return generated.split('?')[0] + '?'
            
            elif context == "question2":
                prompt = f"Customer response: {customer_message}\nAsk about impact:"
                response = text_generator(
                    prompt, 
                    max_new_tokens=20, 
                    do_sample=False, 
                    pad_token_id=50256,
                    eos_token_id=50256,
                    return_full_text=False
                )
                generated = response[0]['generated_text'].strip()
                if len(generated) > 10 and '?' in generated:
                    return generated.split('?')[0] + '?'
            
            elif context == "question3":
                prompt = f"Customer response: {customer_message}\nAsk about previous contact:"
                response = text_generator(
                    prompt, 
                    max_new_tokens=20, 
                    do_sample=False, 
                    pad_token_id=50256,
                    eos_token_id=50256,
                    return_full_text=False
                )
                generated = response[0]['generated_text'].strip()
                if len(generated) > 10 and '?' in generated:
                    return generated.split('?')[0] + '?'
                    
        except:
            pass
    
    templates = {
        "question1": "When did you first notice this issue with the product?",
        "question2": "How has this issue affected your use of the product?",
        "question3": "Have you contacted us about this issue before?",
        "final_solution": "I've tried my best to find alternative solutions. A refund might be the most appropriate solution. Would you like me to proceed?"
    }
    
    return templates.get(context, "How can I help you further?")

def process_complaint_flow(message, state, text_generator=None):
    """Process complaint handling flow"""
    
    def _check_keywords(msg, keywords):
        return any(word in msg.lower() for word in keywords)
    
    if state['stage'] == 'order_details':
        extracted = extract_order_mobile_with_llm(message, text_generator)
        
        if extracted["valid"]:
            if extracted["type"] == "order_id":
                order = extracted["data"]["order"][0]
                state.update({'order_id': order.get('order_id'), 'customer_name': order.get('customer_name', 'Valued Customer'), 'stage': 'issue_description'})
                return f"Hi {state['customer_name']}! I found your order details. You ordered {order.get('category', 'items')} to be delivered at {order.get('address', 'your address')} with {order.get('discount', '0')}% discount. Could you please describe what happened with your order?"
            elif extracted["type"] == "mobile":
                orders = extracted["data"]["orders"]
                order = orders[-1]
                state.update({'order_id': order.get('order_id'), 'customer_name': order.get('customer_name', 'Valued Customer'), 'stage': 'issue_description'})
                return f"Hi {state['customer_name']}! I found your order details. You ordered {order.get('category', 'items')} to be delivered at {order.get('address', 'your address')} with {order.get('discount', '0')}% discount. Could you please describe what happened with your order?"
        elif extracted["message"]:
            return extracted["message"]
        else:
            return "Please provide your order ID (6 digits) or mobile number (10 digits) so that I can fetch your details."
    
    elif state['stage'] == 'issue_description':
        if _check_keywords(message, ['order', 'query', 'complaint', 'problem', 'issue']) and len(message.split()) <= 5:
            return "Could you please describe what happened with your order? What specific issue are you facing?"
        
        state['issue_type'] = message.lower()
        if _check_keywords(message, ['damaged', 'broken', 'defective', 'wrong size', 'size', 'small', 'large', 'fit', 'wrong color', 'color', 'different color']):
            state['stage'] = 'file_upload_request'
            return "I'm sorry to hear about this issue. To help you better and process your request quickly, could you please upload an image or video showing the problem?"
        elif _check_keywords(message, ['late', 'delay', 'not delivered', 'delivery']):
            state['stage'] = 'resolution_attempt'
            state['attempts_to_resolve'] += 1
            return "I apologize for the delivery delay. When was your expected delivery date? Let me check if we can expedite a replacement or provide you with an updated timeline."
        else:
            state['stage'] = 'file_upload_request'
            return "Could you please provide more specific details about the issue? For better assistance, you can also upload an image or video if you have any."
    
    elif state['stage'] == 'file_upload_request':
        if '[File uploaded:' in message or state['file_uploaded']:
            state['file_uploaded'] = True
            state['stage'] = 'replacement_offer'
            return "Thank you for the file. Based on the problem, I believe a replacement would be the best solution. Would you like me to arrange a replacement?"
        elif _check_keywords(message, ['damaged', 'broken', 'defective', 'wrong', 'issue', 'problem', 'torn', 'bad', 'quality']):
            state['stage'] = 'replacement_offer'
            return "I understand the issue with your product. Based on the problem described, I believe a replacement would be the best solution. Would you like me to arrange a replacement?"
        else:
            return "Could you please upload an image or video of the issue for faster processing? If you're having trouble uploading, you can describe the issue in more detail."
    
    elif state['stage'] == 'replacement_offer':
        if _check_keywords(message, ['yes', 'okay', 'sure', 'please']):
            state['stage'] = 'processing_request'
            return "Perfect! I'm processing your replacement request. Our team will arrange a replacement within 2-3 business days. Is there anything else I can help you with?"
        elif _check_keywords(message, ['no', 'refund', 'money back']):
            state['stage'] = 'resolution_attempt'
            state['attempts_to_resolve'] = 0
            return generate_refund_question("replacement_declined", message, text_generator)
        else:
            state['stage'] = 'refund_consideration'
            return "I understand replacement might not be what you're looking for. Would you prefer to explore refund options instead?"
    
    elif state['stage'] == 'resolution_attempt':
        state['attempts_to_resolve'] += 1
        if state['attempts_to_resolve'] == 1:
            state['stage'] = 'refund_step1'
            return generate_refund_question("question1", message, text_generator)
        elif state['attempts_to_resolve'] == 2:
            state['stage'] = 'refund_step2'
            return generate_refund_question("question2", message, text_generator)
        else:
            state['stage'] = 'refund_step3'
            return generate_refund_question("question3", message, text_generator)
    
    elif state['stage'] == 'refund_step1':
        state['stage'] = 'refund_step2'
        return generate_refund_question("question2", message, text_generator)
    
    elif state['stage'] == 'refund_step2':
        state['stage'] = 'refund_step3'
        return generate_refund_question("question3", message, text_generator)
    
    elif state['stage'] == 'refund_step3':
        state['stage'] = 'refund_process'
        return "I have registered your refund request. Our customer executive will get back to you within 5-7 working days with the complete refund process. Thank you for choosing SparkMart. Would you like to close our conversation or continue with something else?"
    
    elif state['stage'] == 'refund_consideration':
        if _check_keywords(message, ['yes', 'ok', 'okay', 'refund', 'money back', 'proceed']):
            state['stage'] = 'refund_step1'
            return generate_refund_question("question1", message, text_generator)
        else:
            return "I understand. Is there any other way I can assist you with this order or any other concern?"
    
    elif state['stage'] in ['refund_process', 'processing_request']:
        if _check_keywords(message, ['recommend', 'recommendation', 'product', 'suggest', 'buy', 'purchase', 'shopping']):
            state['stage'] = 'recommendation_followup'
            from tools import recommend_products, format_recommendations
            recs = recommend_products("popular")
            return format_recommendations(recs)
        elif _check_keywords(message, ['general', 'query', 'service', 'information', 'company', 'provide', 'terms', 'conditions']):
            state['stage'] = 'general_query'
            from tools import handle_general_query
            return handle_general_query(message)
        elif _check_keywords(message, ['refund', 'want refund', 'i want']) and not _check_keywords(message, ['recommend', 'product']):
            return "I have registered your refund request. Our customer executive will get back to you within 5-7 working days with the complete refund process. Thank you for choosing SparkMart. Would you like to close our conversation or continue with something else?"
        elif _check_keywords(message, ['no', 'continue', 'more']):
            state['stage'] = 'query_type'
            return "Of course! Please let me know what other concern you'd like to discuss."
        elif _check_keywords(message, ['yes']) and not _check_keywords(message, ['recommend', 'product', 'want']):
            return "Thank you for contacting SparkMart! Have a wonderful day!"
        elif _check_keywords(message, ['close', 'end', 'bye', 'goodbye', 'thanks', 'thank you', 'nothing', 'exit']):
            return "Thank you for contacting SparkMart! Have a wonderful day!"
        else:
            state['stage'] = 'general_query'
            from tools import handle_general_query
            return handle_general_query(message)
    
    return "I'm here to help with your complaint. Could you please provide more details?"