import tempfile
from pathlib import Path

import joblib
import streamlit as st

from audio_utils import SUPPORTED_AUDIO_EXTENSIONS, convert_to_wav
from features import LABEL_PD, LABEL_TPD, extract_features


BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "models" / "svm_voice_confidence_model.joblib"
CONFIDENCE_THRESHOLD = 0.75
MARGIN_THRESHOLD = 0.25
SUPPORTED_UPLOAD_TYPES = [extension.replace(".", "") for extension in sorted(SUPPORTED_AUDIO_EXTENSIONS)]

LABEL_DESCRIPTION = {
    LABEL_PD: "Percaya Diri",
    LABEL_TPD: "Tidak Percaya Diri",
}


@st.cache_resource
def load_model(model_mtime):
    """
    model_mtime menjadi cache key.
    Jika model dilatih ulang, Streamlit otomatis memuat model terbaru.
    """
    return joblib.load(MODEL_PATH)


def save_uploaded_file(uploaded_file, suffix):
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_audio:
        temp_audio.write(uploaded_file.getbuffer())
        return Path(temp_audio.name)


def predict_uploaded_audio(uploaded_file):
    """
    Alur prediksi:
    upload audio -> convert WAV mono 22050 Hz -> extract_features -> scaler+SVM pipeline.
    """
    model = load_model(MODEL_PATH.stat().st_mtime)
    original_extension = Path(uploaded_file.name).suffix.lower()

    if original_extension not in SUPPORTED_AUDIO_EXTENSIONS:
        allowed = ", ".join(SUPPORTED_UPLOAD_TYPES).upper()
        raise ValueError(f"Format file tidak didukung. Gunakan salah satu: {allowed}")

    temp_input_path = save_uploaded_file(uploaded_file, original_extension)
    temp_wav_path = Path(tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name)

    try:
        convert_to_wav(temp_input_path, temp_wav_path)
        feature_vector = extract_features(temp_wav_path)
        features = feature_vector.reshape(1, -1)

        expected_features = model.named_steps["scaler"].n_features_in_
        if features.shape[1] != expected_features:
            raise ValueError(
                "Jumlah fitur audio tidak sesuai dengan model. "
                f"Audio menghasilkan {features.shape[1]} fitur, "
                f"sedangkan model mengharapkan {expected_features}. "
                "Jalankan ulang `python train_model.py`, lalu restart Streamlit."
            )

        predicted_label = model.predict(features)[0]
        probabilities = model.predict_proba(features)[0]
        class_probabilities = dict(zip(model.classes_, probabilities))
        confidence = class_probabilities[predicted_label]
        probability_pd = class_probabilities.get(LABEL_PD, 0.0)
        probability_tpd = class_probabilities.get(LABEL_TPD, 0.0)
        margin = abs(probability_pd - probability_tpd)

        debug_info = {
            "model_classes": list(model.classes_),
            "raw_probabilities": probabilities.tolist(),
            "feature_shape": features.shape,
            "confidence": float(confidence),
            "margin": float(margin),
        }
    finally:
        temp_input_path.unlink(missing_ok=True)
        temp_wav_path.unlink(missing_ok=True)

    return predicted_label, confidence, class_probabilities, debug_info


st.set_page_config(
    page_title="Klasifikasi Percaya Diri dari Suara",
    layout="centered",
)

st.title("Klasifikasi Percaya Diri dari Suara")
st.write("Upload file audio untuk memprediksi kelas PD atau TPD.")

if not MODEL_PATH.exists():
    st.error("Model belum ditemukan. Jalankan `python train_model.py` terlebih dahulu.")
    st.stop()

uploaded_file = st.file_uploader(
    "Pilih file audio",
    type=SUPPORTED_UPLOAD_TYPES,
)

if uploaded_file is not None:
    uploaded_extension = Path(uploaded_file.name).suffix.lower().replace(".", "")
    st.audio(uploaded_file, format=f"audio/{uploaded_extension}")

    if st.button("Prediksi"):
        with st.spinner("Mengekstraksi fitur dan memprediksi..."):
            try:
                label, confidence, probabilities, debug_info = predict_uploaded_audio(uploaded_file)
            except Exception as error:
                st.error(f"Gagal memproses audio: {error}")
                st.stop()

        probability_pd = probabilities.get(LABEL_PD, 0.0)
        probability_tpd = probabilities.get(LABEL_TPD, 0.0)
        margin = abs(probability_pd - probability_tpd)

        st.subheader("Hasil Prediksi")
        st.write(f"Prediksi: {label}")
        st.write(f"Keterangan: {LABEL_DESCRIPTION[label]}")
        st.write(f"Confidence: {confidence * 100:.2f}%")
        st.write(f"Probabilitas PD: {probability_pd * 100:.2f}%")
        st.write(f"Probabilitas TPD: {probability_tpd * 100:.2f}%")

        if confidence >= CONFIDENCE_THRESHOLD and margin >= MARGIN_THRESHOLD:
            st.success(f"Hasil utama: {label} - {LABEL_DESCRIPTION[label]}")
        else:
            st.warning("Model belum cukup yakin, silakan rekam ulang atau tambah data training.")

        st.progress(float(probability_pd), text=f"PD: {probability_pd * 100:.2f}%")
        st.progress(float(probability_tpd), text=f"TPD: {probability_tpd * 100:.2f}%")

        with st.expander("Debug prediksi"):
            st.write("model.classes_")
            st.json(debug_info["model_classes"])
            st.write("Probabilitas mentah dari predict_proba")
            st.json(debug_info["raw_probabilities"])
            st.write(f"Fitur audio shape: {debug_info['feature_shape']}")
            st.write(f"Confidence: {debug_info['confidence']:.6f}")
            st.write(f"Margin probabilitas: {debug_info['margin']:.6f}")