import os
import gradio as gr
import torch
import sqlite3
from datetime import datetime
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from huggingface_hub import snapshot_download
from sendgrid import SendGridAPIClient
from sendgrid.helpers.mail import Mail

# Student Information
My_info = "Student ID: 6319250G, Name: Aung Hlaing Tun"

# Define Hugging Face Model Repo
MODEL_REPO_ID = "ZAM-ITI-110/Distil_Bert_V3"

# SendGrid API Key (Set in Hugging Face Space Secrets)
SENDGRID_API_KEY = os.getenv("SENDGRID_API_KEY")  # Add this in Space Settings > Secrets

# Load Model & Tokenizer from Hugging Face
def load_model(repo_id):
    """Download and load the model and tokenizer."""
    cache_dir = "/home/user/app/hf_models"
    os.makedirs(cache_dir, exist_ok=True)
    download_dir = snapshot_download(repo_id, cache_dir=cache_dir, local_files_only=False)
    model = AutoModelForSequenceClassification.from_pretrained(download_dir)
    tokenizer = AutoTokenizer.from_pretrained(download_dir)
    return model, tokenizer

# Load Model
model, tokenizer = load_model(MODEL_REPO_ID)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()

# Initialize SQLite Database
def init_db():
    """Create the tickets table if it doesn’t exist."""
    conn = sqlite3.connect("tickets.db")
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS tickets
                 (id INTEGER PRIMARY KEY AUTOINCREMENT, description TEXT, predicted_team TEXT,
                  team_email TEXT, status TEXT, timestamp TEXT)''')
    conn.commit()
    conn.close()

# Prediction Function (Single Ticket)
def predict_team_and_email(text):
    """Predict team and email for a single ticket description."""
    if not text.strip():
        return "", ""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
    with torch.no_grad():
        logits = model(**inputs).logits
    pred = torch.argmax(logits, dim=-1).item()
    
    label_mapping = {
        0: "Code Review Team", 1: "Functional Team", 2: "Infrastructure Team",
        3: "Performance Team", 4: "Security Team"
    }
    email_mapping = {
        0: "codereview.team@company.com", 1: "functional.team@company.com",
        2: "infra.team@company.com", 3: "performance.team@company.com",
        4: "security.team@company.com"
    }
    
    return label_mapping.get(pred, "Unknown"), email_mapping.get(pred, "Unknown")

# Save Ticket to Database
def save_ticket(description, predicted_team, team_email):
    """Save a ticket to the SQLite database."""
    conn = sqlite3.connect("tickets.db")
    c = conn.cursor()
    c.execute("INSERT INTO tickets (description, predicted_team, team_email, status, timestamp) VALUES (?, ?, ?, ?, ?)",
              (description, predicted_team, team_email, "Open", datetime.now().isoformat()))
    conn.commit()
    conn.close()

# Send Email via SendGrid
def send_email(to_email, subject, content):
    """Send an email using SendGrid."""
    if not SENDGRID_API_KEY:
        return "SendGrid API key not set."
    message = Mail(
        from_email="your-verified-email@domain.com",  # Replace with your verified sender email
        to_emails=to_email,
        subject=subject,
        plain_text_content=content)
    try:
        sg = SendGridAPIClient(SENDGRID_API_KEY)
        sg.send(message)
        return "Email sent successfully!"
    except Exception as e:
        return f"Email failed: {str(e)}"

# Send Tickets Function
def send_tickets(*args):
    """Save to DB and send emails for non-empty tickets."""
    tickets = []
    for i, (text, team, email) in enumerate(zip(args[::2], args[1::2], args[2::2]), 1):
        if text.strip() and team and email:
            save_ticket(text, team, email)
            email_status = send_email(email, f"New Ticket Assigned to {team}", text)
            tickets.append(f"Ticket {i}: '{text}' -> {team} ({email}) - {email_status}")
    return "\n".join(tickets) + "\n\nProcessed successfully!" if tickets else "No tickets to send."

# Clear Function
def clear_all():
    """Clear all inputs and outputs."""
    return [""] * 19  # 6 tickets x (input, team, email) + 1 sent_output

# Fetch Ticket History
def get_ticket_history():
    """Retrieve all tickets from the database."""
    conn = sqlite3.connect("tickets.db")
    df = pd.read_sql_query("SELECT * FROM tickets", conn)
    conn.close()
    return df

# Gradio UI Setup
init_db()  # Initialize database on startup
with gr.Blocks(title="AI Ticket Classifier") as interface:
    gr.Markdown("📩 **Development of an AI Ticket Classifier Model Using DistilBERT**")
    gr.Markdown(f"*{My_info}*")
    gr.Markdown(
        """
        **🔍 About this App**
        - Predicts the appropriate **team** and **email** for up to 6 ticket descriptions.
        - Click 'Predict' for each ticket, then 'Send Tickets' to save and notify teams.
        """
    )

    # Ticket Entry Section
    with gr.Column():
        gr.Markdown("### Enter Ticket Descriptions")
        inputs = []
        outputs = []
        buttons = []
        for i in range(6):
            with gr.Row():
                ticket_input = gr.Textbox(lines=2, placeholder=f"Ticket {i+1} description...", label=f"Ticket {i+1}")
                team_output = gr.Textbox(label="Predicted Team", interactive=False)
                email_output = gr.Textbox(label="Team Email", interactive=False)
                predict_btn = gr.Button(f"Predict {i+1}")
                inputs.append(ticket_input)
                outputs.extend([team_output, email_output])
                buttons.append(predict_btn)

    # Action Buttons
    with gr.Row():
        send_btn = gr.Button("Send Tickets")
        clear_btn = gr.Button("Clear")

    # Output for Sent Tickets
    sent_output = gr.Textbox(label="Sent Tickets", interactive=False)

    # Ticket History Section
    with gr.Column():
        gr.Markdown("### Ticket History")
        history_btn = gr.Button("View Tickets")
        history_output = gr.Dataframe()

    # Event Handlers for Predict Buttons
    for i, btn in enumerate(buttons):
        btn.click(
            fn=predict_team_and_email,
            inputs=inputs[i],
            outputs=[outputs[i*2], outputs[i*2 + 1]]
        )

    # Send and Clear Handlers
    send_btn.click(
        fn=send_tickets,
        inputs=inputs + outputs,
        outputs=sent_output
    )
    clear_btn.click(
        fn=clear_all,
        inputs=None,
        outputs=inputs + outputs + [sent_output]
    )

    # History Handler
    history_btn.click(
        fn=get_ticket_history,
        inputs=None,
        outputs=history_output
    )

# Launch the interface
interface.launch()