use axum::{Router, routing::post};
use std::sync::Arc;
use std::{collections::HashMap, path::PathBuf};

pub mod api;
pub mod tts;

use tts::{TtsOpts, TtsPool, load_text_to_speech};

use crate::tts::engine::load_voice_style_map;

#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    dotenvy::dotenv().ok();
    tracing_subscriber::fmt::init();

    let model_dir =
        std::env::var("SUPERTONIC_MODEL_DIR").unwrap_or_else(|_| "./assets".to_string());
    let voice_style_path = std::env::var("SUPERTONIC_VOICE_STYLE")
        .unwrap_or_else(|_| format!("F1={model_dir}/voice_styles/F1.json"));
    let _lang = std::env::var("SUPERTONIC_LANG").unwrap_or_else(|_| "en".to_string());
    let total_step: usize = std::env::var("SUPERTONIC_TOTAL_STEP")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(8);
    let speed: f32 = std::env::var("SUPERTONIC_SPEED")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(1.05);
    let silence_before_duration: f32 = std::env::var("SUPERTONIC_SILENCE_BEFORE_DUR")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0.0);
    let silence_in_content_duration: f32 = std::env::var("SUPERTONIC_SILENCE_IN_CONTENT_DUR")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0.26);
    let silence_after_duration: f32 = std::env::var("SUPERTONIC_SILENCE_AFTER_DUR")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0.4);
    let workers: usize = std::env::var("SUPERTONIC_WORKERS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(2);
    let hf_repo = std::env::var("SUPERTONIC_HF_REPO")
        .unwrap_or_else(|_| "https://huggingface.co/Supertone/supertonic-3".to_string());
    let sound_size_mul: f32 = std::env::var("SUPERTONIC_SIZE_MUL")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(1.5);

    let model_path = PathBuf::from(&model_dir);
    tts::assets::ensure_assets(&model_path, &hf_repo)?;

    let onnx_dir_for_init = model_path.join("onnx").to_string_lossy().into_owned();
    let voice_style_for_init = voice_style_path
        .split(",")
        .filter_map(|i| {
            i.split_once("=").or_else(|| {
                tracing::error!("Voice style '{i}' is not valid.");
                None
            })
        })
        .map(|(k, v)| (k.to_owned(), PathBuf::from(v)))
        .collect::<HashMap<_, _>>();
    let voice_style_map = Arc::new(load_voice_style_map(&voice_style_for_init)?);

    let pool = Arc::new(TtsPool::spawn(
        workers,
        move |id| {
            let span = tracing::info_span!("worker", worker_id = id);
            let _enter = span.enter();
            let tts = load_text_to_speech(&onnx_dir_for_init)?;

            Ok((tts, voice_style_map.clone()))
        },
        TtsOpts {
            total_step,
            speed,
            silence_in_content_duration,
            silence_before_duration,
            silence_after_duration,
            sound_size_mul,
        },
    )?);

    let app = Router::new()
        .route("/", post(api::handler))
        .with_state(pool);

    let addr = std::env::var("ADDR").unwrap_or_else(|_| "0.0.0.0:80".to_string());
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}