from datetime import datetime
import tempfile
from pathlib import Path

import joblib
import librosa
import numpy as np
import streamlit as st

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


BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
MODEL_PATH = BASE_DIR / "models" / "svm_voice_confidence_model.joblib"

CONFIDENCE_THRESHOLD = 0.75
MARGIN_THRESHOLD = 0.25
MIN_RECORDING_DURATION_SECONDS = 2.0
LOW_RMS_WARNING_THRESHOLD = 0.003

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_bytes_to_temp_file(audio_bytes, suffix):
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_audio:
        temp_audio.write(audio_bytes)
        return Path(temp_audio.name)


def save_uploaded_file(uploaded_file, suffix):
    return save_bytes_to_temp_file(uploaded_file.getvalue(), suffix)


def validate_audio_quality(audio_path, min_duration=MIN_RECORDING_DURATION_SECONDS):
    """
    Validasi dasar sebelum ekstraksi fitur.
    Error dipakai untuk kasus file kosong/gagal dibaca/terlalu pendek.
    Warning dipakai untuk audio yang masih bisa diproses tetapi kualitasnya lemah.
    """
    try:
        y, sr = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True)
    except Exception as error:
        raise ValueError(f"File audio gagal dibaca: {error}") from error

    if y.size == 0:
        raise ValueError("File audio kosong atau tidak memiliki sinyal suara.")

    duration = librosa.get_duration(y=y, sr=sr)
    if duration < min_duration:
        raise ValueError(
            f"Durasi audio terlalu pendek ({duration:.2f} detik). "
            "Silakan rekam suara 3 sampai 5 detik dengan jelas."
        )

    rms = float(np.sqrt(np.mean(y**2)))
    warning = None
    if rms < LOW_RMS_WARNING_THRESHOLD:
        warning = (
            f"Suara terdeteksi cukup pelan (RMS={rms:.5f}). "
            "Jika hasil kurang tepat, rekam ulang dengan suara lebih jelas."
        )

    return {
        "duration": duration,
        "rms": rms,
        "warning": warning,
    }


def predict_audio_path(audio_path, validate_quality=True):
    """
    Fungsi prediksi umum untuk upload dan rekaman.

    Urutan:
    audio_path -> convert_to_wav -> validate -> extract_features -> predict_proba.
    Label utama diambil dari probabilitas terbesar, bukan model.predict().
    """
    model = load_model(MODEL_PATH.stat().st_mtime)
    audio_path = Path(audio_path)
    extension = audio_path.suffix.lower()

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

    temp_wav_path = Path(tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name)

    try:
        convert_to_wav(audio_path, temp_wav_path)
        quality_info = validate_audio_quality(temp_wav_path) if validate_quality else None
        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."
            )

        probabilities = model.predict_proba(features)[0]
        class_probabilities = dict(zip(model.classes_, probabilities))
        predicted_label = max(class_probabilities, key=class_probabilities.get)
        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),
            "quality_info": quality_info,
        }
    finally:
        temp_wav_path.unlink(missing_ok=True)

    return predicted_label, confidence, class_probabilities, debug_info


def render_prediction_result(label, confidence, probabilities, debug_info):
    probability_pd = probabilities.get(LABEL_PD, 0.0)
    probability_tpd = probabilities.get(LABEL_TPD, 0.0)
    margin = abs(probability_pd - probability_tpd)

    quality_info = debug_info.get("quality_info")
    if quality_info and quality_info.get("warning"):
        st.warning(quality_info["warning"])

    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}%")
    st.write(f"Margin: {margin * 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}")
        if quality_info:
            st.write(f"Durasi audio: {quality_info['duration']:.2f} detik")
            st.write(f"RMS audio: {quality_info['rms']:.6f}")


def get_audio_recorder_input():
    """
    Menggunakan st.audio_input jika tersedia.
    Jika belum tersedia, coba fallback ke streamlit-mic-recorder.
    """
    if hasattr(st, "audio_input"):
        return st.audio_input("Rekam suara")

    try:
        from streamlit_mic_recorder import mic_recorder
    except ImportError:
        st.error(
            "Versi Streamlit ini belum mendukung st.audio_input. "
            "Install fallback recorder dengan perintah: pip install streamlit-mic-recorder"
        )
        return None

    audio = mic_recorder(
        start_prompt="Mulai Rekam",
        stop_prompt="Berhenti Rekam",
        just_once=False,
        use_container_width=True,
        key="mic_recorder",
    )

    if audio and audio.get("bytes"):
        suffix = ".wav"
        return {
            "bytes": audio["bytes"],
            "suffix": suffix,
            "mime_type": "audio/wav",
        }

    return None


def get_recording_bytes(recording):
    if recording is None:
        return None, ".wav", "audio/wav"

    if isinstance(recording, dict):
        return recording["bytes"], recording.get("suffix", ".wav"), recording.get("mime_type", "audio/wav")

    suffix = Path(recording.name).suffix.lower() or ".wav"
    return recording.getvalue(), suffix, recording.type or "audio/wav"


def save_recording_to_dataset(source_audio_path, label):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    target_dir = DATA_DIR / label
    target_dir.mkdir(parents=True, exist_ok=True)
    target_path = target_dir / f"recorded_{label}_{timestamp}.wav"

    convert_to_wav(source_audio_path, target_path)
    return target_path


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

st.title("Klasifikasi Percaya Diri dari Suara")
st.write("Input audio, ekstraksi fitur, prediksi SVM, lalu tampilkan PD atau TPD.")

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

upload_tab, record_tab = st.tabs(["Upload Audio", "Rekam Audio"])

with upload_tab:
    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 Upload"):
            temp_input_path = save_uploaded_file(uploaded_file, Path(uploaded_file.name).suffix.lower())
            try:
                with st.spinner("Mengekstraksi fitur dan memprediksi..."):
                    label, confidence, probabilities, debug_info = predict_audio_path(temp_input_path)
                render_prediction_result(label, confidence, probabilities, debug_info)
            except Exception as error:
                st.error(f"Gagal memproses audio: {error}")
            finally:
                temp_input_path.unlink(missing_ok=True)

with record_tab:
    st.write(
        "Silakan rekam suara selama 3-5 detik. Gunakan suara yang jelas, "
        "tidak terlalu pelan, dan hindari noise ruangan."
    )

    recording = get_audio_recorder_input()
    audio_bytes, suffix, mime_type = get_recording_bytes(recording)

    if audio_bytes:
        st.audio(audio_bytes, format=mime_type)

        temp_recording_path = save_bytes_to_temp_file(audio_bytes, suffix)
        st.session_state["latest_recording_path"] = str(temp_recording_path)

        if st.button("Prediksi Rekaman"):
            try:
                with st.spinner("Mengekstraksi fitur dan memprediksi rekaman..."):
                    label, confidence, probabilities, debug_info = predict_audio_path(temp_recording_path)
                render_prediction_result(label, confidence, probabilities, debug_info)
            except Exception as error:
                st.error(f"Gagal memproses rekaman: {error}")

        st.divider()
        st.subheader("Simpan Rekaman ke Dataset")
        selected_label = st.selectbox(
            "Label manual",
            options=[LABEL_PD, LABEL_TPD],
            format_func=lambda label: f"{label} - {LABEL_DESCRIPTION[label]}",
        )

        if st.button("Simpan ke Dataset"):
            try:
                saved_path = save_recording_to_dataset(temp_recording_path, selected_label)
                st.success(
                    "Rekaman berhasil disimpan. Jalankan ulang train_model.py "
                    "untuk melatih ulang model."
                )
                st.write(f"File: {saved_path}")
            except Exception as error:
                st.error(f"Gagal menyimpan rekaman ke dataset: {error}")
