import os
import torch
import argparse
import gradio as gr
from mailersend import emails
from dotenv import load_dotenv
from elevenlabs.client import ElevenLabs
from elevenlabs import play, save
import base64
import psycopg2
from urllib.parse import urlparse, parse_qs

# Load environment variables
load_dotenv()

# Argument parsing
parser = argparse.ArgumentParser()
parser.add_argument("--share", action='store_true', default=False, help="make link public")
args = parser.parse_args()

# Initialize ElevenLabs client
client = ElevenLabs(api_key=os.environ.get("ELEVENLABS_API_KEY"))
device = 'cuda' if torch.cuda.is_available() else 'cpu'
output_dir = 'outputs'
os.makedirs(output_dir, exist_ok=True)

supported_languages = ['zh', 'en']

MAILERSEND_API_KEY = os.getenv("MAILERSEND_API_KEY")
MAILERSEND_DOMAIN = os.getenv("MAILERSEND_DOMAIN")
MAILERSEND_SENDER_EMAIL = f"noreply@{MAILERSEND_DOMAIN}"
MAILERSEND_SENDER_NAME = "Voice Clone App"

# List of blocked words
BLOCKED_WORDS = ['Kill','hurt','shoot','gun','rifle','AR','semi automatic','knife','blade','sword','punch harm','disrupt','blackmail','steal','bitch','cunt','fuck','freaking','nigger','nigga','niggas','cracker','jew','oriental','fag','faggot','account','money','transfer','urgent','help','scared','policy','frightened','accident','fear','scam','address','social security number','assault','injure','maim','destroy','damage','threaten','intimidate','bully','menace','blackmail','extort','exploit','defame','steal','rob','embezzle','defraud Harass','jerk','idiot','stupid','moron','asshole','con','trick','swindle','defraud','payment','credit card','bank account','urgent','immediate','afraid','phone number','email','password']  # Add more words as needed

# Function to check for blocked words
def contains_blocked_words(text):
    return any(word.lower() in text.lower() for word in BLOCKED_WORDS)

# Function to send email with downloadable file using MailerSend
def send_email_with_file(recipient_email, file_path, subject, body):
    try:
        mailer = emails.NewEmail(MAILERSEND_API_KEY)

        mail_body = {}
        mail_from = {
            "name": MAILERSEND_SENDER_NAME,
            "email": MAILERSEND_SENDER_EMAIL,
        }
        recipients = [
            {
                "name": "Recipient",
                "email": recipient_email,
            }
        ]

        mailer.set_mail_from(mail_from, mail_body)
        mailer.set_mail_to(recipients, mail_body)
        mailer.set_subject(subject, mail_body)
        mailer.set_html_content(f"<p>{body}</p>", mail_body)
        mailer.set_plaintext_content(body, mail_body)

        with open(file_path, "rb") as file:
            attachment_content = base64.b64encode(file.read()).decode('utf-8')
        
        attachments = [
            {
                "filename": os.path.basename(file_path),
                "content": attachment_content,
                "disposition": "attachment"
            }
        ]
        mailer.set_attachments(attachments, mail_body)

        response = mailer.send(mail_body)

        if response[0] == 202:
            print("Email sent successfully")
            return True
        else:
            print(f"Failed to send email. Status code: {response[0]}")
            print(f"Response: {response[1]}")
            return False
    except Exception as e:
        print(f"An error occurred while sending email: {e}")
        return False

# Database connection details
connection_string = os.environ.get("DATABASE_URL")
result = urlparse(connection_string)
user = result.username
password = result.password
host = result.hostname
port = result.port
database = result.path[1:]
sslmode = parse_qs(result.query)['sslmode'][0]

# Function to add user information to the database
def add_user_info_to_db(email, order_name, sample_audio, name):
    connection = psycopg2.connect(
        dbname=database,
        user=user,
        password=password,
        host=host,
        port=port,
        sslmode=sslmode
    )
    cursor = connection.cursor()
    
    insert_query = """
    INSERT INTO example_table (email, order_name, sample_audio, name)
    VALUES (%s, %s, %s, %s);
    """
    try:
        cursor.execute(insert_query, (email, order_name, sample_audio, name))
        connection.commit()
        print("User information added to the database successfully")
    except Exception as error:
        print(f"Error occurred during data insertion: {error}")
        connection.rollback()
    finally:
        cursor.close()
        connection.close()

# Predict function
def predict(prompt, style, audio_file_pth, voice_name, customer_email, order_name):
    text_hint = 'Your file will only be saved for 24 hours.\n'
    if len(prompt) < 2:
        text_hint += "[ERROR] Please provide a longer prompt text.\n"
        return text_hint, None, None
    if len(prompt) > 200:
        text_hint += "[ERROR] Text length limited to 200 characters. Please try shorter text.\n"
        return text_hint, None, None
    if contains_blocked_words(prompt):
        text_hint += "[ERROR] Your text contains blocked words. Please remove them and try again.\n"
        return text_hint, None, None
    full_voice_name = f"{voice_name} - {customer_email}"
    print(audio_file_pth)
    voice = client.clone(
        name=full_voice_name,
        description="A trial voice model for testing",
        files=[audio_file_pth],
    )
    audio = client.generate(text=prompt, voice=voice)
    save_path = f'{output_dir}/output.wav'
    save(audio, save_path)
    
    subject = "Your Voice Clone File"
    body = "Thank you for using our Voice Clone service. Your file is attached."
    if send_email_with_file(customer_email, save_path, subject, body):
        text_hint += "Email sent successfully with the voice file.\n"
    else:
        text_hint += "Failed to send email with the voice file. Please try again later.\n"

    # Add user information to the database
    with open(audio_file_pth, "rb") as f:
        sample_audio = f.read()
    add_user_info_to_db(customer_email, order_name, sample_audio, voice_name)

    return text_hint, save_path, audio_file_pth

# Gradio interface setup
with gr.Blocks(gr.themes.Glass()) as demo:
    with gr.Row():
        with gr.Column():
            input_text_gr = gr.Textbox(
                label="Create This",
                info="One or two sentences at a time is better. Up to 200 text characters.",
                value="He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
            )
            style_gr = gr.Dropdown(
                label="Style",
                choices=['default', 'whispering', 'cheerful', 'terrified', 'angry', 'sad', 'friendly'],
                info="Please upload a reference audio file that is at least 1 minute long. For best results, ensure the audio is clear. You can use Adobe Podcast Enhance(https://podcast.adobe.com/enhance) to improve the audio quality before uploading.",
                max_choices=1,
                value="default",
            )
            ref_gr = gr.Audio(
                label="Original Audio",
                type="filepath",
                sources=["upload"],
            )
            voice_name_gr = gr.Textbox(
                label="Your name",
                value="Sam"
            )
            order_gr = gr.Textbox(
                label="Your order",
                value="Sample Order",
            )
            customer_email_gr = gr.Textbox(
                label="Your Email",
                info="We'll send you a downloadable file to this email address."
            )
            tts_button = gr.Button("Start", elem_id="send-btn", visible=True)

        with gr.Column():
            out_text_gr = gr.Text(label="Info")
            audio_gr = gr.Audio(label="Replicated Sound", autoplay=True)
            ref_audio_gr = gr.Audio(label="Original Audio Used ")

            tts_button.click(predict, [input_text_gr, style_gr, ref_gr, voice_name_gr, customer_email_gr, order_gr], outputs=[out_text_gr, audio_gr, ref_audio_gr])

    demo.queue()
    demo.launch(debug=True, show_api=False, share=True)

css = """
footer {visibility: hidden}
audio .btn-container {display: none}
"""

demo.add_css(css)