import pandas as pd
import numpy as np
import torch
from sklearn.preprocessing import LabelEncoder
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr
import gspread
from oauth2client.service_account import ServiceAccountCredentials
import os
import json

# =========================
# PASSWORD (from env)
# =========================
APP_PASSWORD = os.environ.get("APP_PASSWORD")

if APP_PASSWORD is None:
    raise ValueError("APP_PASSWORD environment variable not set")

# =========================
# THEMES
# =========================
THEMES = {
    "Lifestyle": "Lifestyle: everyday life, hobbies, routines, hanging out, personal moments",
    "Travel & Experience": "Travel & Experience: trips, visiting places, tourism, events, sports games, concerts",
    "Reflection": "Reflection: deep thoughts, life lessons, gratitude, personal growth, emotional reflection",
    "Brand & Business": "Brand & Business: marketing, promoting a brand, entrepreneurship, product announcements, advertisements",
    "Education & Value": "Education & Value: advice, tips, teaching something, sharing knowledge",
    "Relationships": "Relationships: family, friends, romantic, social connections",
    "Performance & Sport": "Performance & Sport: competitions, sports achievements, athletic performance"
}

label_encoder = LabelEncoder()
label_encoder.fit(list(THEMES.keys()))

# =========================
# GOOGLE SHEETS AUTH
# =========================
scope = [
    "https://spreadsheets.google.com/feeds",
    "https://www.googleapis.com/auth/drive"
]

sheet = None
try:
    creds_json_str = os.environ.get("GDRIVE_SERVICE_ACCOUNT_JSON")

    if not creds_json_str:
        print("CRITICAL ERROR: Missing GDRIVE_SERVICE_ACCOUNT_JSON")
    else:
        creds_info = json.loads(creds_json_str)
        creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_info, scope)
        client = gspread.authorize(creds)
        sheet = client.open("ClassifierTraining").get_worksheet(0)
        print(f"Connected to sheet: {sheet.title}")
except Exception as e:
    print(f"Auth Failure: {e}")

# =========================
# MODEL LOAD
# =========================
model_name = "nirvairkataria/caption-classifier-model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.to("cuda" if torch.cuda.is_available() else "cpu")
model.eval()

# =========================
# FUNCTIONS
# =========================
def classify_caption_trained_model(caption):
    inputs = tokenizer(caption, return_tensors="pt", truncation=True, padding=True)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
    label_idx = np.argmax(probs)
    pred_label = label_encoder.inverse_transform([label_idx])[0]
    confidence = round(float(probs[label_idx]), 2)
    all_themes = list(THEMES.keys())
    return pred_label, confidence, all_themes

def save_feedback_to_sheet(caption, corrected_theme):
    if sheet is None:
        return "Error: Google Sheet authentication failed."
    try:
        sheet.append_row([caption, corrected_theme])
        return "Feedback submitted successfully! It will be used when the model is retrained in order to improve accuracy for future captions."
    except Exception as e:
        return f"Failed: {e}"

def process_csv(file_obj):
    df = pd.read_csv(file_obj.name)
    predicted_themes, confidences = [], []
    for caption in df["message"]:
        theme, conf, _ = classify_caption_trained_model(caption)
        predicted_themes.append(theme)
        confidences.append(conf)
    df["Predicted Theme"] = predicted_themes
    df["Confidence"] = confidences
    output_file = "classified_results.csv"
    df.to_csv(output_file, index=False)
    return df, output_file

def handle_classification(caption):
    pred_label, confidence, all_themes = classify_caption_trained_model(caption)
    return pred_label, confidence, gr.update(choices=all_themes, value=pred_label)

# =========================
# PASSWORD CHECK
# =========================
def check_password(pw):
    if pw == APP_PASSWORD:
        return gr.update(visible=False), gr.update(visible=True), ""
    else:
        return gr.update(), gr.update(), "Incorrect password"

# =========================
# UI
# =========================
with gr.Blocks() as demo:

    # LOGIN SCREEN
    with gr.Column(visible=True) as login_screen:
        gr.Markdown("## 🔒 Enter Password")
        password_input = gr.Textbox(type="password", label="Password")
        login_btn = gr.Button("Enter")
        login_status = gr.Textbox(label="", interactive=False)

    # MAIN APP (HIDDEN)
    with gr.Column(visible=False) as app_screen:
        gr.Markdown("# Caption Classifier")

        with gr.Tab("Single Caption"):
            caption_input = gr.Textbox(label="Enter Caption", lines=4)
            with gr.Row():
                theme_output = gr.Textbox(label="Predicted Theme", interactive=False)
                confidence_output = gr.Textbox(label="Confidence", interactive=False)

            classify_btn = gr.Button("Classify")

            corrected_theme_dropdown = gr.Dropdown(label="Corrected Theme", choices=list(THEMES.keys()))
            submit_btn = gr.Button("Submit Correction")
            status_output = gr.Textbox(label="Status", interactive=False)

            classify_btn.click(
                handle_classification,
                inputs=caption_input,
                outputs=[theme_output, confidence_output, corrected_theme_dropdown]
            )

            submit_btn.click(
                save_feedback_to_sheet,
                inputs=[caption_input, corrected_theme_dropdown],
                outputs=status_output
            )

        with gr.Tab("CSV Upload"):
            file_input = gr.File(label="Upload CSV")
            output_table = gr.Dataframe()
            download_file = gr.File()
            process_btn = gr.Button("Run")

            process_btn.click(
                process_csv,
                inputs=file_input,
                outputs=[output_table, download_file]
            )

    # LOGIN ACTION
    login_btn.click(
        check_password,
        inputs=password_input,
        outputs=[login_screen, app_screen, login_status]
    )

# =========================
# LAUNCH
# =========================
demo.launch(server_name="0.0.0.0", server_port=7860)