use axum::Router;
use axum::routing::{MethodRouter, any};
use crossbeam::channel;
use serde_json::{Map, Value};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tower::ServiceBuilder;
use tower_http::trace::TraceLayer;

use crate::args::Args;
#[cfg(feature = "static-files")]
use axum::http::header::CONTENT_TYPE;
#[cfg(feature = "static-files")]
use axum::http::{Response, StatusCode};
#[cfg(feature = "static-files")]
use axum::routing::get;
#[cfg(feature = "static-files")]
use std::fs;
use std::net::SocketAddr;
#[cfg(feature = "static-files")]
use std::path::{Path, PathBuf};
#[cfg(feature = "static-files")]
use tower_http::services::ServeDir;

pub fn init_logging() {
    match log4rs::init_file("log4rs.yaml", Default::default()) {
        Ok(_) => {}
        Err(_) => {
            let config =
                serde_yaml_ng::from_str::<log4rs::config::RawConfig>(LOG4RS_DEFAULT_CONFIG)
                    .unwrap();
            log4rs::init_raw_config(config).expect("Unable to initialize with default config.");
            log::info!(
                "No log4rs.yaml file found in working directory. Using default config (stdout)."
            );
        }
    }
}

#[cfg(feature = "static-files")]
fn directory_fallback(path: PathBuf) -> Response<String> {
    let files = fs::read_dir(path).unwrap();
    let mut values = Vec::new();
    for file in files {
        match file {
            Ok(file) => values.push(Value::from(file.file_name().to_str().unwrap())),
            Err(e) => {
                log::error!("{}", e);
            }
        }
    }
    let value = Value::Array(values);
    Response::builder()
        .status(StatusCode::OK)
        .header(CONTENT_TYPE, "application/json")
        .body(value.to_string())
        .unwrap()
}

#[cfg(feature = "static-files")]
fn handle_directory(path: PathBuf) -> Router {
    let mut router = Router::new();
    let dir = fs::read_dir(path.clone()).unwrap();
    let mut has_index = false;
    for file in dir {
        match file {
            Ok(file) => {
                if file.metadata().unwrap().is_dir() {
                    router = router.nest_service(
                        &format!("/{}", file.file_name().to_str().unwrap()),
                        handle_directory(file.path()),
                    );
                } else {
                    if file.file_name() == "index.html" {
                        log::debug!("{:?} has index", path);
                        has_index = true;
                        continue;
                    }
                }
            }
            Err(e) => {
                log::error!("{}", e);
            }
        }
    }
    if !has_index {
        router = router
            .route("/", get(directory_fallback(path.clone())))
            .fallback_service(ServeDir::new(path));
    } else {
        router = router.fallback_service(ServeDir::new(path));
    }
    router
}

#[cfg(feature = "static-files")]
fn handle_directories_with_router(dir: &String) -> Router {
    let mut router = Router::new();

    let path = Path::new(dir);
    if fs::metadata(path).unwrap().is_dir() {
        router = handle_directory(path.to_path_buf());
    }

    router
}

#[derive(Clone)]
pub struct WsProxyState {
    #[allow(unused)]
    pub socket_state: Arc<Mutex<Value>>,
    pub(crate) txs: Arc<Mutex<HashMap<uuid::Uuid, channel::Sender<String>>>>,
    pub(crate) crg_ws_reconnect_rate_s: u64,
    pub(crate) stop: Arc<Mutex<bool>>,
    pub(crate) registration_paths: Option<Vec<String>>,
    pub ext: Arc<Mutex<Map<String, Value>>>,
}

unsafe impl Send for WsProxyState {}
unsafe impl Sync for WsProxyState {}

impl WsProxyState {
    fn new(crg_ws_reconnect_rate_s: u64, registration_paths: Option<Vec<String>>) -> Self {
        WsProxyState {
            socket_state: Default::default(),
            txs: Default::default(),
            stop: Arc::new(Mutex::new(false)),
            registration_paths,
            crg_ws_reconnect_rate_s,
            ext: Default::default(),
        }
    }
}

pub struct WsProxy {
    state: Arc<WsProxyState>,
    ip: Option<String>,
    port: Option<u16>,
    crg: (String, u16),
    #[cfg(feature = "static-files")]
    static_dir: Option<String>,
    services: HashMap<String, Router<Arc<WsProxyState>>>,
    routes: HashMap<String, MethodRouter<Arc<WsProxyState>>>,
}

const LOG4RS_DEFAULT_CONFIG: &str = include_str!("../log4rs.yaml");

async fn shutdown(state: Arc<WsProxyState>) {
    let ctrl_c = tokio::signal::ctrl_c();
    let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
        .expect("failed to install signal handler");
    let mut terminate = Box::pin(terminate.recv());

    tokio::select! {
        _ = ctrl_c => {
            let mut stop = state.stop.lock().await;
            *stop = true;
        },
        _ = &mut terminate => {
            let mut stop = state.stop.lock().await;
            *stop = true;
        },
    }
}

impl WsProxy {
    pub fn builder() -> WsProxyBuilder {
        WsProxyBuilder {
            crg_ws_reconnect_rate_s: 5,
            ..Default::default()
        }
    }

    pub async fn start(&self) -> std::io::Result<()> {
        let ip = self.ip.clone().unwrap_or("127.0.0.1".to_string());
        let port = self.port.clone().unwrap_or(8001);

        let shared_state = self.state.clone();

        crate::crg_ws::init(self.crg.clone(), self.state.clone()).await;

        #[allow(unused_mut)]
        let mut app = Router::new()
            .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()))
            .route("/WS/", any(crate::apex_ws::ws_handler));

        for (path, service) in self.services.clone() {
            app = app.nest(&path, service);
        }

        for (path, route) in &self.routes {
            app = app.route(path, route.clone());
        }

        #[cfg(feature = "static-files")]
        if let Some(ref dir) = self.static_dir {
            let file_service = ServeDir::new(dir.clone());
            let router = handle_directories_with_router(dir).fallback_service(file_service);
            app = app.fallback_service(router);
        }

        let app = app.with_state(shared_state.clone());

        log::info!("Starting server on {}:{}", ip, port);
        let listener = tokio::net::TcpListener::bind(format!("{}:{}", ip, port)).await?;
        axum::serve(
            listener,
            app.into_make_service_with_connect_info::<SocketAddr>(),
        )
        .with_graceful_shutdown(shutdown(shared_state))
        .await
    }
}

#[derive(Default)]
pub struct WsProxyBuilder {
    crg_host: Option<String>,
    crg_port: Option<u16>,
    crg_ws_reconnect_rate_s: u64,
    port: Option<u16>,
    #[cfg(feature = "static-files")]
    static_dir: Option<String>,
    registration_paths: Option<Vec<String>>,
    services: HashMap<String, Router<Arc<WsProxyState>>>,
    routes: HashMap<String, MethodRouter<Arc<WsProxyState>>>,
}

impl WsProxyBuilder {
    pub fn build(self) -> WsProxy {
        let crg_host = self
            .crg_host
            .expect("Unable to start server without CRG host address");
        let crg_port = self
            .crg_port
            .expect("Unable to start server without CRG host port");
        WsProxy {
            state: Arc::new(WsProxyState::new(
                self.crg_ws_reconnect_rate_s,
                self.registration_paths,
            )),
            ip: None,
            port: self.port,
            #[cfg(feature = "static-files")]
            static_dir: self.static_dir,
            crg: (crg_host, crg_port),
            services: self.services,
            routes: self.routes,
        }
    }

    pub fn for_crg(mut self, host: &str, port: u16) -> Self {
        self.crg_host = Some(host.into());
        self.crg_port = Some(port);
        self
    }

    pub fn port(mut self, port: u16) -> Self {
        self.port = Some(port);
        self
    }

    pub fn with_registration_paths(mut self, registration_paths: Vec<String>) -> Self {
        if registration_paths.len() > 0 {
            self.registration_paths = Some(registration_paths);
        }
        self
    }

    #[cfg(feature = "static-files")]
    pub fn with_static(mut self, dir: String) -> Self {
        self.static_dir = Some(dir);
        self
    }

    pub fn with_config(mut self, config: Args) -> Self {
        self.crg_port = Some(config.crg_port);
        self.crg_host = Some(config.crg_host);
        self.port = Some(config.port);
        self.crg_ws_reconnect_rate_s = config.crg_ws_reconnect_s;
        self.registration_paths = if config.registration_paths.len() > 0 {
            Some(config.registration_paths)
        } else {
            None
        };
        self
    }

    pub fn with_route(mut self, path: &str, route: MethodRouter<Arc<WsProxyState>>) -> Self {
        self.routes.insert(path.to_string(), route);
        self
    }

    pub fn with_service(mut self, path: &str, service: Router<Arc<WsProxyState>>) -> Self {
        self.services.insert(path.to_string(), service);
        self
    }
}
