use crate::JarvisState;
use crate::{
    model::{
        extension::{Extension, ExtensionInfo},
        manifest::{ExtPackageJsonExtra, MANIFEST_FILE_NAME},
    },
    utils::manifest::load_jarvis_ext_manifest,
};
use std::collections::HashMap;
use std::net::{SocketAddr, TcpListener};
use std::{
    path::PathBuf,
    sync::{Arc, Mutex},
};
use tauri::{AppHandle, Runtime, State, Window};
use tower_http::cors::CorsLayer;
use tower_http::services::ServeDir;

/// manifest_path can be folder of package.json
/// If it's a folder, join it with package.json
#[tauri::command]
pub async fn load_manifest<R: Runtime>(
    _app: tauri::AppHandle<R>,
    _window: tauri::Window<R>,
    manifest_path: PathBuf,
) -> Result<ExtPackageJsonExtra, String> {
    Ok(ExtPackageJsonExtra::from(
        load_jarvis_ext_manifest(manifest_path.clone()).map_err(|e| e.to_string())?,
        manifest_path.parent().unwrap().to_path_buf(),
    ))
}

#[tauri::command]
pub async fn load_all_extensions<R: Runtime>(
    _app: tauri::AppHandle<R>,
    _window: tauri::Window<R>,
    extensions_folder: PathBuf,
) -> Result<Vec<ExtPackageJsonExtra>, String> {
    let mut extensions_with_path: Vec<ExtPackageJsonExtra> = vec![];
    for entry in std::fs::read_dir(extensions_folder).map_err(|e| e.to_string())? {
        let entry = entry.map_err(|e| e.to_string())?;

        if entry.path().join(MANIFEST_FILE_NAME).exists() {
            let ext_manifest = load_jarvis_ext_manifest(entry.path()).map_err(|e| e.to_string());
            if ext_manifest.is_err() {
                continue;
            }
            extensions_with_path.push(ExtPackageJsonExtra::from(
                ext_manifest.unwrap(),
                entry.path(),
            ));
        }
    }
    Ok(extensions_with_path)
}

#[tauri::command]
pub async fn is_window_label_registered<R: Runtime>(
    _app: AppHandle<R>,
    state: State<'_, JarvisState>,
    label: String,
) -> Result<bool, String> {
    Ok(state
        .window_label_ext_map
        .lock()
        .unwrap()
        .contains_key(label.as_str()))
}

/// When ui iframe ext loaded, register the window containing the iframe with the extension info
/// Extension info contains path to extension command and dist path, so Kunkun's custom file server knows where to load
/// static files from when a request is received from a window.
#[tauri::command]
pub async fn register_extension_window<R: Runtime>(
    _app: AppHandle<R>,
    state: State<'_, JarvisState>,
    extension_path: PathBuf,
    dist: Option<String>,
    window_label: Option<String>,
) -> Result<String, String> {
    let window_label_2 = match window_label {
        Some(label) => label,
        None => format!("main:ext:{}", uuid::Uuid::new_v4()),
    };
    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
    let ext = Extension {
        info: ExtensionInfo {
            path: extension_path,
            processes: vec![],
            dist: dist,
        },
        shutdown_handle: Arc::new(Mutex::new(None)),
        server_handle: Arc::new(Mutex::new(None)),
    };
    label_ext_map.insert(window_label_2.clone(), ext);
    Ok(window_label_2)
}

#[tauri::command]
pub async fn register_extension_spawned_process<R: Runtime>(
    _app: AppHandle<R>,
    state: State<'_, JarvisState>,
    window_label: String,
    pid: u32,
) -> Result<(), String> {
    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
    // check if window_label is registered, if not, return error
    if !label_ext_map.contains_key(window_label.as_str()) {
        return Err(format!(
            "Window with label {} is not registered",
            &window_label
        ));
    }
    let ext = label_ext_map.get_mut(window_label.as_str()).unwrap();
    ext.info.processes.push(pid);
    Ok(())
}

#[tauri::command]
pub async fn unregister_extension_spawned_process<R: Runtime>(
    _app: AppHandle<R>,
    state: State<'_, JarvisState>,
    window_label: String,
    pid: u32,
) -> Result<(), String> {
    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
    label_ext_map
        .get_mut(window_label.as_str())
        .unwrap()
        .info
        .processes
        .retain(|p| *p != pid);
    Ok(())
}

#[tauri::command]
pub async fn get_ext_label_map<R: Runtime>(
    _app: AppHandle<R>,
    _window: Window<R>,
    state: State<'_, JarvisState>,
) -> Result<HashMap<String, ExtensionInfo>, String> {
    let label_ext_map = state.window_label_ext_map.lock().unwrap();
    // turn label_ext_map from HashMap<String, Extension> to HashMap<String, ExtensionInfo>
    let label_ext_map_info: HashMap<String, ExtensionInfo> = label_ext_map
        .iter()
        .map(|(label, ext)| (label.clone(), ext.info.clone()))
        .collect();
    Ok(label_ext_map_info)
}

/// unregister extension window
#[tauri::command]
pub async fn unregister_extension_window<R: Runtime>(
    _app: AppHandle<R>,
    state: State<'_, JarvisState>,
    label: String,
) -> Result<(), String> {
    // find extension, if there is shutdown handle, shutdown it
    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
    // find extension info with window_label
    let ext = label_ext_map.get(label.as_str()).cloned();
    // drop(label_ext_map);
    match ext {
        Some(ext) => {
            let shutdown_handle = ext.shutdown_handle.lock().unwrap();
            if let Some(shutdown_handle) = shutdown_handle.as_ref() {
                shutdown_handle.shutdown();
                log::info!("Shutdown extension file server with label {}", label);
            }
            let server_handle = ext.server_handle.lock().unwrap();
            if let Some(server_handle) = server_handle.as_ref() {
                server_handle.abort();
                log::info!("Abort extension file server with label {}", label);
            }
            label_ext_map.remove(label.as_str());
            log::info!("Unregistered extension window with label {}", label);
        }
        None => return Err(format!("Extension with label {} not found", label)),
    }
    // state
    //     .window_label_ext_map
    //     .lock()
    //     .unwrap()
    //     .remove(label.as_str())
    //     .is_some()
    Ok(())
}

/// spawn extension file server, only for Windows
#[tauri::command]
pub async fn spawn_extension_file_server<R: Runtime>(
    _app: AppHandle<R>,
    state: State<'_, JarvisState>,
    window_label: String,
) -> Result<SocketAddr, String> {
    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
    // find extension info with window_label
    let ext = label_ext_map.get(window_label.as_str());
    if ext.is_none() {
        return Err(format!("Extension with label {} not found", window_label));
    }
    let ext = ext.unwrap();
    let mut ext_path = ext.info.path.clone();
    if let Some(dist) = ext.info.dist.clone() {
        ext_path = ext_path.join(dist);
    }
    // TODO: spawn file server
    let shutdown_handle = axum_server::Handle::new();
    let shutdown_handle_clone = shutdown_handle.clone();
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    let rest_router = axum::Router::new()
        .layer(CorsLayer::permissive())
        .nest_service("/", ServeDir::new(ext_path));

    let server_handle = tauri::async_runtime::spawn(async move {
        // axum_server::bind(addr)
        axum_server::from_tcp(listener)
            .handle(shutdown_handle_clone)
            .serve(rest_router.into_make_service())
            // .serve(combined_router.into_make_service())
            .await
            .unwrap();
    });
    // add server handle and shutdown handle to extension
    let ext = label_ext_map.get_mut(window_label.as_str()).unwrap();
    ext.server_handle.lock().unwrap().replace(server_handle);
    ext.shutdown_handle.lock().unwrap().replace(shutdown_handle);

    // TODO: add server handle and shutdown handle
    Ok(addr)
}
