import gradio as gr
import re
from datetime import datetime
from transformers import pipeline

# Load a conversational model
model_name = "microsoft/DialoGPT-small"
chatbot = pipeline("text-generation", model=model_name)

# Simulated databases
PRODUCT_CATALOG = {
    "Cinco Phone": ["phone", "mobile", "cellphone", "smartphone"],
    "Cinco Tablet": ["tablet", "pad"],
    "Cinco Laptop": ["laptop", "notebook", "computer"],
    "Cinco Earbuds": ["earbuds", "earphones", "headphones"]
}

RECEIPT_DATABASE = {
    "RC-1001": {"product": "Cinco Phone", "date": "2023-05-15", "email": "customer1@example.com"},
    "RC-1002": {"product": "Cinco Tablet", "date": "2023-06-20", "email": "customer2@example.com"},
    "RC-1003": {"product": "Cinco Laptop", "date": "2023-07-10", "email": "customer3@example.com"},
}

# Return policy details
RETURN_POLICY = {
    "window": "30 days",
    "process_time": "5-7 business days",
    "conditions": "Product must be in original condition with packaging",
    "contact": "1-800-CINCO-SUPPORT",
    "email": "returns@cinco.com"
}

# State management
class ChatState:
    NORMAL = 0
    RETURN_DETECTED = 1
    PRODUCT_CONFIRMED = 2
    RECEIPT_HANDLING = 3
    ALT_INFO_HANDLING = 4
    RETURN_PROCESSING = 5
    DISENGAGED = 6
    POLICY_QUESTION = 7

def initialize_state():
    return {
        "stage": ChatState.NORMAL,
        "return_product": None,
        "receipt": None,
        "name": None,
        "purchase_date": None,
        "email": None,
        "error_count": 0,
        "conversation_history": []
    }

# Helper functions
def process_return(product):
    return f"✅ Return processed for your {product}!\nConfirmation sent to your email."

def is_return_request(text):
    return bool(re.search(r'\b(return|refund|exchange|send back)\b', text.lower()))

def is_policy_question(text):
    policy_keywords = ["policy", "how long", "can i", "after", "before", "days", 
                       "time", "window", "condition", "return policy", "refund policy"]
    return any(keyword in text.lower() for keyword in policy_keywords)

def is_general_question(text):
    general_keywords = ["what", "who", "where", "when", "why", "how", "you", "your", "cinco"]
    return any(keyword in text.lower() for keyword in general_keywords)

def format_conversation(history, new_input):
    """Format conversation history with context about Cinco and return policy"""
    prompt = (
        "You are CincoBot, a friendly customer service assistant for Cinco Electronics. "
        "You specialize in helping customers with product returns. "
        f"Cinco's return policy: Returns accepted within {RETURN_POLICY['window']} of purchase. "
        f"Conditions: {RETURN_POLICY['conditions']}. "
        f"Processing time: {RETURN_POLICY['process_time']}. "
        "Be helpful, polite, and professional. Answer questions about returns and Cinco products.\n\n"
    )
    
    for message in history:
        if message["role"] == "user":
            prompt += f"Customer: {message['content']}\n"
        else:
            prompt += f"Assistant: {message['content']}\n"
    
    prompt += f"Customer: {new_input}\nAssistant:"
    return prompt

def recognize_product(user_input):
    """Match product using keywords and synonyms"""
    user_input = user_input.lower()
    for product, aliases in PRODUCT_CATALOG.items():
        if any(alias in user_input for alias in aliases):
            return product
        if product.lower() in user_input:
            return product
    return None

def handle_return_flow(user_input, state):
    """Natural conversation handling for return process"""
    response = ""
    
    # Check if this is a policy question
    if is_policy_question(user_input):
        state["stage"] = ChatState.POLICY_QUESTION
        return None, state
    
    # Detect return request naturally
    if state["stage"] == ChatState.NORMAL:
        if is_return_request(user_input):
            state["stage"] = ChatState.RETURN_DETECTED
            response = (
                "I'd be happy to help with your return! Cinco accepts returns within 30 days of purchase. "
                "Could you tell me which product you'd like to return?"
            )
        else:
            # Let AI handle general conversation
            return None, state
    
    # Product identification
    elif state["stage"] == ChatState.RETURN_DETECTED:
        product = recognize_product(user_input)
        if product:
            state["return_product"] = product
            state["stage"] = ChatState.PRODUCT_CONFIRMED
            response = (
                f"Got it, you want to return the {product}. "
                "Could you please share your receipt number? It usually looks like RC-1234. "
                "If you don't have it, just say so."
            )
        else:
            state["error_count"] += 1
            if state["error_count"] >= 2:
                response = "I'm having trouble identifying the product. Let me connect you to a human agent."
                state["stage"] = ChatState.DISENGAGED
            else:
                response = "I understand you want to return something. Could you specify the product name? We have phones, tablets, laptops, and earbuds."
    
    # Receipt handling
    elif state["stage"] == ChatState.PRODUCT_CONFIRMED:
        receipt_match = re.search(r'RC-\d{4}', user_input.upper())
        if receipt_match:
            receipt_id = receipt_match.group(0)
            if receipt_id in RECEIPT_DATABASE:
                state["receipt"] = receipt_id
                state["email"] = RECEIPT_DATABASE[receipt_id]["email"]
                response = "Thanks! I've verified your receipt. One moment while I process your return..."
                state["stage"] = ChatState.RETURN_PROCESSING
            else:
                response = "I couldn't find that receipt in our system. Could you double-check the number?"
                state["error_count"] += 1
        else:
            if "don't" in user_input.lower() or "no" in user_input.lower() or "not" in user_input.lower():
                response = (
                    "No problem! To process your return without a receipt, I'll need:\n"
                    "1. Your full name\n"
                    "2. Purchase date (YYYY-MM-DD)\n"
                    "3. Email address\n"
                    "You can provide them in any order."
                )
                state["stage"] = ChatState.ALT_INFO_HANDLING
            else:
                response = "Could you share your receipt number? If you don't have it, just say so."
                state["error_count"] += 1
        
        if state["error_count"] >= 2:
            response = "I'm having trouble with your receipt. Let me connect you with a human agent."
            state["stage"] = ChatState.DISENGAGED
    
    # Alternative information handling
    elif state["stage"] == ChatState.ALT_INFO_HANDLING:
        # Extract information from user input
        name_match = re.search(r'[A-Za-z]+ [A-Za-z]+', user_input)
        date_match = re.search(r'\d{4}-\d{2}-\d{2}', user_input)
        email_match = re.search(r'[\w\.-]+@[\w\.-]+', user_input)
        
        if name_match:
            state["name"] = name_match.group(0)
        if date_match:
            state["purchase_date"] = date_match.group(0)
        if email_match:
            state["email"] = email_match.group(0)
        
        # Check if we have all information
        if state["name"] and state["purchase_date"] and state["email"]:
            response = "Thanks! I have all I need to process your return."
            state["stage"] = ChatState.RETURN_PROCESSING
        else:
            missing = []
            if not state["name"]:
                missing.append("full name")
            if not state["purchase_date"]:
                missing.append("purchase date (YYYY-MM-DD)")
            if not state["email"]:
                missing.append("email address")
                
            response = f"I still need your {', '.join(missing)}. Could you provide that?"
    
    # Return processing
    elif state["stage"] == ChatState.RETURN_PROCESSING:
        response = process_return(state["return_product"])
        response += "\n\nIs there anything else I can help you with today?"
        state = initialize_state()
    
    # Disengagement
    elif state["stage"] == ChatState.DISENGAGED:
        response = (
            "I'm transferring you to one of our human support agents. "
            f"Please contact us at {RETURN_POLICY['contact']} or {RETURN_POLICY['email']}."
        )
    
    return response, state

def handle_policy_question(question):
    """Handle return policy questions with detailed responses"""
    question = question.lower()
    
    if "policy" in question or "return policy" in question:
        return (
            "📋 Cinco Return Policy:\n"
            f"- Return Window: {RETURN_POLICY['window']} from purchase date\n"
            f"- Conditions: {RETURN_POLICY['conditions']}\n"
            f"- Processing Time: {RETURN_POLICY['process_time']}\n"
            f"- Contact: {RETURN_POLICY['contact']} or {RETURN_POLICY['email']}"
        )
    
    if "after" in question and "days" in question:
        return (
            f"Our return window is {RETURN_POLICY['window']} from purchase date. "
            "After this period, returns are generally not accepted, but you can contact "
            f"our support team at {RETURN_POLICY['contact']} for special cases."
        )
    
    if "how long" in question or "process" in question:
        return f"Returns take {RETURN_POLICY['process_time']} to process once received."
    
    if "condition" in question or "package" in question:
        return f"To be eligible for return: {RETURN_POLICY['conditions']}"
    
    if "contact" in question or "help" in question or "support" in question:
        return (
            "You can contact our support team:\n"
            f"- Phone: {RETURN_POLICY['contact']}\n"
            f"- Email: {RETURN_POLICY['email']}\n"
            "We're available Monday-Friday, 9AM-5PM EST."
        )
    
    return (
        "Our standard return policy:\n"
        f"- {RETURN_POLICY['window']} return window\n"
        f"- {RETURN_POLICY['conditions']}\n"
        f"- Refunds processed in {RETURN_POLICY['process_time']}\n"
        f"Contact us at {RETURN_POLICY['contact']} for more details."
    )

def chat_fn(user_input, chat_history, state):
    """Main chat function with natural conversation handling"""
    # First check if we're in a return flow
    return_response, state = handle_return_flow(user_input, state)
    
    if return_response is not None:
        response = return_response
    elif state["stage"] == ChatState.POLICY_QUESTION:
        response = handle_policy_question(user_input)
        state["stage"] = ChatState.NORMAL
    else:
        # Generate conversational response using AI with context
        prompt = format_conversation(state["conversation_history"], user_input)
        result = chatbot(
            prompt, 
            max_length=100,
            num_return_sequences=1,
            pad_token_id=50256,
            temperature=0.8,
            do_sample=True
        )
        
        # Clean up the response
        generated = result[0]['generated_text']
        if "Assistant:" in generated:
            response = generated.split("Assistant:")[-1].split("\n")[0].strip()
        else:
            response = generated.strip()
    
    # Update conversation history
    state["conversation_history"].append({"role": "user", "content": user_input})
    state["conversation_history"].append({"role": "assistant", "content": response})
    
    # Update chat history for Gradio
    chat_history.append((user_input, response))
    
    return "", chat_history, state

# Gradio interface with clear branding
with gr.Blocks(title="Cinco Returns Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🛍️ Cinco Returns Assistant")
    gr.Markdown(
        "### How can I help with your Cinco return today?\n"
        f"**Return Policy**: {RETURN_POLICY['window']} return window • {RETURN_POLICY['conditions']}"
    )
    
    chatbot_ui = gr.Chatbot(height=400, bubble_full_width=False)
    
    with gr.Row():
        user_input = gr.Textbox(
            placeholder="Ask about returns, policies, or Cinco products...", 
            show_label=False,
            container=False,
            autofocus=True
        )
        submit_btn = gr.Button("Send", variant="primary")
    
    state = gr.State(initialize_state())
    
    with gr.Row():
        clear_btn = gr.Button("Start New Conversation")
        policy_btn = gr.Button("View Return Policy")
    
    def clear_history():
        return [], initialize_state()
    
    def show_policy():
        return [(handle_policy_question("policy"), "")]
    
    def submit_message(user_input, chat_history, state):
        if not user_input.strip():
            return "", chat_history, state
        return chat_fn(user_input, chat_history, state)
    
    submit_btn.click(
        submit_message,
        [user_input, chatbot_ui, state],
        [user_input, chatbot_ui, state]
    )
    
    user_input.submit(
        submit_message,
        [user_input, chatbot_ui, state],
        [user_input, chatbot_ui, state]
    )
    
    clear_btn.click(
        clear_history,
        outputs=[chatbot_ui, state]
    )
    
    policy_btn.click(
        show_policy,
        inputs=[],
        outputs=[chatbot_ui]
    )

if __name__ == "__main__":
    demo.queue().launch()