import os
import sys
import logging
import gradio as gr
import autogen
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

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

# Check Python version
if sys.version_info < (3, 7):
    logging.error("This script requires Python 3.7 or higher")
    sys.exit(1)

# Check and set environment variables
required_env_vars = ['HUGGINGFACE_API_KEY']
for var in required_env_vars:
    if var not in os.environ:
        logging.error(f"Environment variable {var} is not set")
        sys.exit(1)

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

# Rest of your code (SHARED_CONTEXT, guardrail functions, etc.) remains the same

# CrewAI setup
try:
    from crewai import Agent as CrewAgent, Task, Crew
except ImportError:
    logging.error("Failed to import crewai. Make sure it's installed: pip install crewai")
    sys.exit(1)

# CrewAI and AutoGen setup remains the same

# Main function
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)

        # Use crewAI for initial query rephrasing
        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 = crew.kickoff()

        # Use AutoGen for generating the response
        async def get_autogen_response():
            await user_proxy.a_initiate_chat(
                response_expert_autogen,
                message=f"Please provide a respectful and empathetic response to the following query: '{rephrased_query}'"
            )
            return response_expert_autogen.last_message()["content"]

        response = await get_autogen_response()

        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:
        logging.error(f"Error in zerodha_support: {e}")
        return "I apologize, but an error occurred while processing your request. Please try again later."

# Wrap the asynchronous function for Gradio
def zerodha_support_wrapper(message, history):
    return asyncio.run(zerodha_support(message, history))

# Gradio interface setup
demo = gr.ChatInterface(
    zerodha_support_wrapper,
    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__":
    try:
        public_url = demo.launch(share=True, server_name="0.0.0.0", server_port=7860)
        print(f"\n\nSHAREABLE LINK: {public_url}\n\n")
    except Exception as e:
        logging.error(f"Failed to launch Gradio interface: {e}")
        sys.exit(1)