import gradio as gr
from huggingface_hub import InferenceClient
import os
import re

# Initialize the client with the Zephyr 7B model
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Shared context for both agents
SHARED_CONTEXT = """You are part of a multi-agent system designed to provide respectful, empathetic, and accurate support for Zerodha, a leading Indian financial services company. Your role is crucial in ensuring all interactions uphold the highest standards of customer service while maintaining Zerodha's excellent reputation.

Key points about Zerodha:
1. India's largest discount broker, known for innovative technology and low-cost trading.
2. Flat fee structure: ₹20 per executed order for intraday and F&O trades, zero brokerage for delivery equity investments.
3. Main trading platform: Kite (web and mobile).
4. Coin platform for commission-free direct mutual fund investments.
5. Extensive educational resources through Varsity.
6. Additional tools: Sentinel (price alerts) and ChartIQ (advanced charting).
7. Console for account management and administrative tasks.

Always prioritize user safety, ethical investing practices, and transparent communication. Never provide information that could mislead users or bring disrepute to Zerodha."""

# Communication Expert context
COMMUNICATION_EXPERT_CONTEXT = SHARED_CONTEXT + """
As the Communication Expert, your primary role is to interpret user queries with the utmost respect and empathy. You should:
1. Rephrase the user's query to ensure it's understood in the most positive and constructive light.
2. Identify and highlight any emotional subtext or concerns in the query.
3. Frame the query in a way that invites a supportive and informative response.
4. Ensure that any potential complaints or frustrations are acknowledged respectfully.

Your output should be a rephrased version of the user's query that maintains its original intent while setting the stage for an empathetic and respectful response."""

# Response Expert context
RESPONSE_EXPERT_CONTEXT = SHARED_CONTEXT + """
As the Response Expert, your role is to provide accurate, helpful, and emotionally intelligent responses to user queries. You should:
1. Address the user's question or concern directly and comprehensively.
2. Maintain a tone of respect and empathy throughout your response.
3. Provide clear, factual information about Zerodha's services and policies.
4. When discussing financial matters, include appropriate disclaimers and encourage users to seek professional advice for complex decisions.
5. For complaints or concerns, acknowledge them respectfully and provide constructive guidance or escalation paths.
6. Always uphold Zerodha's reputation for transparency and user-centric service.

Your output should be a complete, informative response that addresses the user's query while demonstrating empathy and respect."""

def format_prompt(context, message):
    return f"""<s>[INST] <<SYS>>
{context}
<</SYS>>

{message} [/INST]
"""

def communication_expert(query):
    prompt = format_prompt(COMMUNICATION_EXPERT_CONTEXT, f"Please rephrase the following user query with empathy and respect: '{query}'")
    response = client.text_generation(prompt, max_new_tokens=200, temperature=0.7)
    return response

def response_expert(rephrased_query):
    prompt = format_prompt(RESPONSE_EXPERT_CONTEXT, f"Please provide a respectful and empathetic response to the following query: '{rephrased_query}'")
    response = client.text_generation(prompt, max_new_tokens=500, temperature=0.7)
    return response

def post_process_response(response):
    # Remove any potentially disrespectful language
    response = re.sub(r'\b(stupid|dumb|idiotic|foolish)\b', 'mistaken', response, flags=re.IGNORECASE)
    
    # Ensure the response ends respectfully
    if not re.search(r'(Thank you|Is there anything else|Hope this helps|Let me know if you need more information)\s*$', response, re.IGNORECASE):
        response += "\n\nIs there anything else I can help you with regarding Zerodha's services?"
    
    # Add a disclaimer for financial advice
    if re.search(r'\b(invest|trade|buy|sell|market)\b', response, re.IGNORECASE):
        response += "\n\nPlease note that this information is for educational purposes only and should not be considered as financial advice. Always do your own research and consider consulting with a qualified financial advisor before making investment decisions."
    
    return response

def zerodha_support(message, history):
    # Step 1: Communication Expert rephrases the query
    rephrased_query = communication_expert(message)
    print(f"Rephrased query: {rephrased_query}")  # Debug print
    
    # Step 2: Response Expert generates a response
    initial_response = response_expert(rephrased_query)
    print(f"Initial response: {initial_response}")  # Debug print
    
    # Step 3: Post-process the response
    final_response = post_process_response(initial_response)
    print(f"Final response: {final_response}")  # Debug print
    
    return final_response

demo = gr.ChatInterface(
    zerodha_support,
    chatbot=gr.Chatbot(height=600),
    textbox=gr.Textbox(placeholder="Ask your question about Zerodha here...", container=False, scale=7),
    title="Zerodha Support Assistant",
    description="Ask questions about Zerodha's services, trading, account management, and more. Our multi-agent system ensures respectful and empathetic responses.",
    theme="soft",
    examples=[
        "How do I open a Zerodha account?",
        "I'm frustrated with the recent changes to the Kite platform. Can you help?",
        "What are the risks involved in F&O trading?",
        "I think there's an error in my account statement. What should I do?",
        "Can you explain Zerodha's policy on intraday trading margins?",
        "I'm new to investing. What resources does Zerodha offer for beginners?",
        "How does Zerodha ensure the security of my investments and personal data?"
    ],
)

if __name__ == "__main__":
    public_url = demo.launch(share=True, server_name="0.0.0.0", server_port=7860)
    print(f"\n\nSHAREABLE LINK: {public_url}\n\n")