import os
import sys
import logging
import gradio as gr
from huggingface_hub import InferenceClient
import re
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
import asyncio
from crewai import Agent as CrewAgent, Task, Crew
import openai

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Check for OpenAI API key
if 'OPENAI_API_KEY' not in os.environ:
    logger.error("OPENAI_API_KEY environment variable is not set.")
    logger.info("Please set the OPENAI_API_KEY environment variable before running this script.")
    sys.exit(1)

# Set OpenAI API key
openai.api_key = os.environ['OPENAI_API_KEY']

# Initialize the client with the Mistral-7B-Instruct-v0.2 model
try:
    client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
except Exception as e:
    logger.error(f"Failed to initialize InferenceClient: {e}")
    sys.exit(1)

# Shared context for 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."""

# Guardrail functions
def sanitize_input(input_text):
    return re.sub(r'[<>&\']', '', input_text)

approved_topics = ['account opening', 'trading', 'fees', 'platforms', 'funds', 'regulations', 'support']
vectorizer = CountVectorizer()
classifier = MultinomialNB()

X = vectorizer.fit_transform(approved_topics)
y = np.arange(len(approved_topics))
classifier.fit(X, y)

def is_relevant_topic(query):
    query_vector = vectorizer.transform([query])
    prediction = classifier.predict(query_vector)
    return prediction[0] in range(len(approved_topics))

def redact_sensitive_info(text):
    text = re.sub(r'\b\d{10,12}\b', '[REDACTED]', text)
    text = re.sub(r'[A-Z]{5}[0-9]{4}[A-Z]', '[REDACTED]', text)
    return text

def check_response_content(response):
    unauthorized_patterns = [
        r'\b(guarantee|assured|certain)\b.*\b(returns|profit)\b',
        r'\b(buy|sell)\b.*\b(specific stocks?|shares?)\b'
    ]
    return not any(re.search(pattern, response, re.IGNORECASE) for pattern in unauthorized_patterns)

def check_confidence(response):
    uncertain_phrases = ["I'm not sure", "It's possible", "I don't have enough information"]
    return not any(phrase.lower() in response.lower() for phrase in uncertain_phrases)

# Agent setup using CrewAI only
communication_expert_crew = CrewAgent(
    role='Communication Expert',
    goal='Interpret and rephrase user queries with empathy and respect',
    backstory=SHARED_CONTEXT,
    verbose=True,
    allow_delegation=False
)

response_expert_crew = CrewAgent(
    role='Response Expert',
    goal='Provide accurate, helpful, and emotionally intelligent responses to user queries',
    backstory=SHARED_CONTEXT,
    verbose=True,
    allow_delegation=False
)

# Main function for handling user queries
async def zerodha_support(message, history):
    try:
        sanitized_message = sanitize_input(message)
        
        if not is_relevant_topic(sanitized_message):
            return "I'm sorry, but I can only assist with queries related to Zerodha's services and trading. Could you please ask a question about your Zerodha account, trading, or our platforms?"
        
        sanitized_message = redact_sensitive_info(sanitized_message)

        # CrewAI for rephrasing query
        rephrase_task = Task(
            description=f"Rephrase the following user query with empathy and respect: '{sanitized_message}'",
            agent=communication_expert_crew
        )
        
        crew = Crew(
            agents=[communication_expert_crew],
            tasks=[rephrase_task],
            verbose=2
        )

        rephrased_query = await crew.kickoff()

        # CrewAI for generating the response
        response_task = Task(
            description=f"Generate a comprehensive and empathetic response to: '{rephrased_query}'",
            agent=response_expert_crew
        )
        
        response_crew = Crew(
            agents=[response_expert_crew],
            tasks=[response_task],
            verbose=2
        )

        response = await response_crew.kickoff()

        if not check_response_content(response):
            response += "\n\nPlease note that I cannot provide specific investment advice or guarantee returns. For personalized guidance, please consult with a qualified financial advisor."

        if not check_confidence(response):
            return "I apologize, but I'm not confident in providing an accurate answer to this query. For the most up-to-date and accurate information, please contact Zerodha's customer support directly."

        final_response = post_process_response(response)

        return final_response
    except Exception as e:
        logger.error(f"Error in zerodha_support: {e}")
        return "I apologize, but an error occurred while processing your request. Please try again later."

# Gradio interface setup
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__":
    demo.launch()