import io
import os
import uuid
import librosa
import numpy as np
import soundfile as sf
import streamlit as st
from huggingface_hub import HfApi, CommitOperationAdd

DATASET_REPO_ID = "uchi8977/kiss-vs-tsk"
TARGET_SECONDS = 0.5
TARGET_SR = 16000
SILENCE_THRESHOLD = 0.01

st.set_page_config(page_title="投げキッスと舌打ち", page_icon="😘", layout="centered")


def process_audio(audio_file) -> bytes | None:
    try:
        y, sr = librosa.load(io.BytesIO(audio_file.getvalue()), sr=TARGET_SR, mono=True)

        peak = float(np.max(np.abs(y)))
        if peak < SILENCE_THRESHOLD:
            st.warning("⚠️ 音が小さすぎるか、無音のようです。もう一度録音してください。")
            return None
        y = y / peak

        target_samples = int(sr * TARGET_SECONDS)
        if len(y) < target_samples:
            y = np.pad(y, (0, target_samples - len(y)))
        else:
            peak_idx = int(np.argmax(np.abs(y)))
            start = max(0, min(peak_idx - target_samples // 2, len(y) - target_samples))
            y = y[start:start + target_samples]

        buf = io.BytesIO()
        sf.write(buf, y, sr, format='WAV', subtype='PCM_16')
        return buf.getvalue()

    except Exception as e:
        st.error(f"音声処理エラー: {e}")
        return None


def submit_dataset(kiss_bytes: bytes, tsk_bytes: bytes, hf_token: str) -> str | None:
    try:
        submission_id = uuid.uuid4().hex
        operations = [
            CommitOperationAdd(
                path_in_repo=f"data/{label}/{label}_{submission_id}.wav",
                path_or_fileobj=wav_bytes,
            )
            for label, wav_bytes in [("kiss", kiss_bytes), ("tsk", tsk_bytes)]
        ]
        HfApi(token=hf_token).create_commit(
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            operations=operations,
            commit_message=f"Add submission {submission_id}",
        )
        return submission_id
    except Exception as e:
        st.error(f"❌ 送信エラー: {str(e)}")
        return None


def go_to(step: int):
    st.session_state.step = step
    st.rerun()


def back_button():
    if st.button("← 前に戻る", use_container_width=True, key="back"):
        go_to(st.session_state.step - 1)


for key, val in {"step": 0, "kiss_bytes": None, "tsk_bytes": None, "submission_id": None}.items():
    st.session_state.setdefault(key, val)

st.title("投げキッスと舌打ち")
st.write("機械学習のデータセット作成にご協力ください。")
st.divider()

step = st.session_state.step

if step == 0:
    st.warning("⚠️ ご提供いただいた音声データは、機械学習のオープンデータセットとしてHugging Face上で一般公開されます。")

    agree = st.checkbox("音声データが公開されることに同意します")
    st.write("準備ができたら下のボタンを押して録音を開始してください。")

    if st.button("はじめる", type="primary", use_container_width=True, disabled=not agree):
        go_to(1)

elif step in (1, 2):
    label, display_name, next_label = (
        ("kiss", "投げキッス", "次へ（舌打ちの録音）➔")
        if step == 1 else
        ("tsk", "舌打ち", "送信画面へ ➔")
    )

    st.subheader(f"Step {step}/2: {display_name}")
    st.write(f"下のマイクボタンを押して{display_name}を録音してください。")

    audio_file = st.audio_input(
        f"{display_name}を録音",
        label_visibility="collapsed",
        key=f"{label}_input"
    )
    st.write(f"※最大音量の瞬間を中心に【{TARGET_SECONDS}秒】だけを自動で切り抜きます（{TARGET_SECONDS}秒未満は無音で補完）。")

    if audio_file:
        wav_bytes = process_audio(audio_file)
        if wav_bytes:
            st.session_state[f"{label}_bytes"] = wav_bytes
            st.success("✅ 切り抜き完了！再生して確認してください。")
            st.audio(wav_bytes)

            if st.button(next_label, type="primary", use_container_width=True):
                go_to(step + 1)

    back_button()

elif step == 3:
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        st.error("❌ サーバー側の設定エラー：HF_TOKEN が設定されていません。管理者にお問い合わせください。")
        st.stop()

    kiss_bytes = st.session_state.kiss_bytes
    tsk_bytes = st.session_state.tsk_bytes
    if not (kiss_bytes and tsk_bytes):
        st.error("❌ 録音データが見つかりません。最初からやり直してください。")
        if st.button("最初に戻る", use_container_width=True):
            st.session_state.kiss_bytes = None
            st.session_state.tsk_bytes = None
            go_to(0)
        st.stop()

    st.subheader("🚀 最終確認")
    st.write("以下の2つのデータを送信します。よろしいですか？")

    col1, col2 = st.columns(2)
    with col1:
        st.write("投げキッス")
        st.audio(kiss_bytes)
    with col2:
        st.write("舌打ち")
        st.audio(tsk_bytes)

    st.write("※送信ボタンを押すと、Hugging Faceの公開リポジトリにアップロードされます。")

    if st.button("送信する！", type="primary", use_container_width=True):
        with st.spinner("Hugging Faceに送信中..."):
            submission_id = submit_dataset(kiss_bytes, tsk_bytes, hf_token)
            if submission_id:
                st.session_state.submission_id = submission_id
                st.session_state.kiss_bytes = None
                st.session_state.tsk_bytes = None
                go_to(4)

    back_button()

elif step == 4:
    st.balloons()
    st.header("🎉 送信完了！")
    st.write("ご協力ありがとうございました。タブを閉じて終了してください。")

    submission_id = st.session_state.submission_id
    if submission_id:
        st.divider()
        st.write(f"**あなたのデータのID: `{submission_id}`**")
        st.write(f"リポジトリで `kiss_{submission_id}.wav` と `tsk_{submission_id}.wav` として公開されています。")

    st.divider()

    st.write("寄せられた録音データは、Hugging Faceのリポジトリでリアルタイムに公開されています👇")
    st.link_button(
        "他の人の投げキッスと舌打ちを聴きに行く",
        f"https://huggingface.co/datasets/{DATASET_REPO_ID}/tree/main/data",
        use_container_width=True
    )