use std::{net::SocketAddr, ops::ControlFlow, sync::Arc, time::Duration};

use avam_protocol::{Packet, Packets, SimConnectPacket, SystemPacket};
use axum::{
    extract::{
        ws::{Message, WebSocket},
        ConnectInfo, State, WebSocketUpgrade,
    },
    response::IntoResponse,
};
use axum_extra::{
    headers::{self, authorization::Bearer},
    TypedHeader,
};
use futures::{SinkExt, StreamExt};
use http::StatusCode;
use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
use tokio::{
    sync::mpsc,
    time::{sleep, Instant},
};
use tracing::info;

use crate::{
    domain::api::{
        ports::ApiService,
        prelude::{TokenSubject, User},
    },
    inbound::http::{
        handlers::oauth::{TokenClaims, VerifyClientAuthorizationRequest},
        state::AppState,
    },
};

pub async fn ws_handler<S: ApiService>(
    State(app_state): State<AppState<S>>,
    ws: WebSocketUpgrade,
    auth_token: Option<TypedHeader<headers::Authorization<Bearer>>>,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
) -> Result<impl IntoResponse, StatusCode> {
    let auth_token = match auth_token {
        Some(TypedHeader(token)) => Some(token.token().to_string()),
        None => {
            tracing::error!("No Authorization Header Supplied");
            return Err(StatusCode::UNAUTHORIZED);
        }
    };

    let Some(auth_token) = auth_token else {
        tracing::error!("No Token Supplied");
        return Err(StatusCode::UNAUTHORIZED);
    };

    let jwt_secret = &app_state.config().jwt_secret;

    let claims = decode::<TokenClaims<String>>(
        &auth_token,
        &DecodingKey::from_secret(jwt_secret.as_ref()),
        &Validation::new(Algorithm::HS256),
    )
    .map_err(|e| {
        tracing::error!("Unable to decode token: {}\n{:?}", auth_token, e);
        StatusCode::UNAUTHORIZED
    })?
    .claims;

    let token_subject: TokenSubject = serde_qs::from_str(&claims.sub).map_err(|e| {
        tracing::error!("Unable to parse Token Subject: {}\n{:?}", &claims.sub, e);
        StatusCode::BAD_REQUEST
    })?;

    let user_id = token_subject.user_id();
    let client_id = token_subject.client_id();

    let app = app_state.api_service();
    let Ok(user) = app
        .verify_client_authorization(VerifyClientAuthorizationRequest::new(user_id, client_id))
        .await
    else {
        tracing::error!("Client {} not authorized by {}", client_id, user_id);
        return Err(StatusCode::UNAUTHORIZED);
    };

    if app_state.api_service().has_connection(&user) {
        return Err(StatusCode::CONFLICT);
    }

    Ok(ws.on_upgrade(move |socket| handle_socket(app_state.clone(), socket, user, addr)))
}

/// Actual websocket statemachine (one will be spawned per connection)
async fn handle_socket<S: ApiService>(
    app_state: AppState<S>,
    socket: WebSocket,
    user: User,
    who: SocketAddr,
) {
    // AppState needs to store user to channel handles or something
    // It'd know who's connected at all times and be able to send messages

    let (sender, mut receiver) = mpsc::channel(10);
    app_state.api_service().add_connection(&user, sender);

    // This can probably be hella-abstracted away and be made a lot cleaner
    let (mut writer, mut reader) = socket.split();
    // This probably needs a (mpsc?) channel, then a dedicated thread to send ping on that channel
    let mut ping_timer = Instant::now();
    ping_timer.checked_add(Duration::from_secs(15)); // add 15 seconds so we instantly trigger a ping
    let u = user.clone();
    let writer_handle = tokio::spawn(async move {
        loop {
            let msg = match receiver.try_recv() {
                Ok(msg) => Some(msg),
                Err(mpsc::error::TryRecvError::Disconnected) => break,
                _ => None,
            };

            if let Some(msg) = msg {
                let Ok(encoded_message) = msg.encode() else {
                    tracing::error!("Unable to encode message for sending: {:#?}", msg);
                    sleep(Duration::from_millis(100)).await;
                    continue;
                };
                tracing::info!("> [{}]: {:?}", &u.email(), msg);
                let _ = writer.send(Message::Binary(encoded_message)).await;
            }

            if ping_timer.elapsed() >= Duration::from_secs(15) {
                tracing::trace!(
                    "> [{}]: {:?}",
                    &u.email(),
                    Packets::System(SystemPacket::Ping)
                );
                let _ = writer
                    .send(Message::Binary(
                        Packets::System(SystemPacket::Ping).encode().unwrap(),
                    ))
                    .await;
                ping_timer = Instant::now();
            }

            sleep(Duration::from_millis(100)).await;
        }
    });

    let u = user.clone();
    let a_s = app_state.clone();
    let reader_handle = tokio::spawn(async move {
        loop {
            if let Some(Ok(data)) = reader.next().await {
                let Ok(packet) = Packets::decode(&data.clone().into_data()) else {
                    if let Message::Close(Some(c)) = data {
                        tracing::info!("{} disconnected: {}", u.email(), c.reason);
                        break;
                    }
                    tracing::error!("Invalid packet received {:#?}", data);
                    sleep(Duration::from_millis(100)).await;
                    continue;
                };

                match packet {
                    Packets::System(system_packet) => {
                        process_system_packet(a_s.clone(), &u, system_packet).await;
                    }
                    Packets::SimConnect(sim_connect_packet) => {
                        process_message(a_s.api_service().clone(), &u, sim_connect_packet).await;
                    }
                }
            }
            sleep(Duration::from_millis(100)).await;
        }
    });

    tokio::select! {
        _ = reader_handle => {
            tracing::debug!("reader closed");
        }
        _ = writer_handle => {
            tracing::debug!("writer closed");
        }
    };

    // remove the user/channel from AppState
    app_state.api_service().remove_connection(&user);

    // returning from the handler closes the websocket connection
    tracing::debug!("Websocket context {who} destroyed");
}

async fn process_system_packet<S>(
    app_state: AppState<S>,
    user: &User,
    packet: SystemPacket,
) -> ControlFlow<(), ()>
where
    S: ApiService,
{
    tracing::trace!("< [{}]: {:?}", user.email(), packet);
    match packet {
        SystemPacket::Ping => {
            // send back pong
            let _ = app_state
                .api_service()
                .send(user, Packets::System(SystemPacket::Pong))
                .await;
        }
        SystemPacket::Pong => {
            // noop
        }
        SystemPacket::Close { reason } => {
            tracing::debug!("{} disconnect: {}", user.email(), reason);
            return ControlFlow::Break(()); // whatever this means
        }
    };

    ControlFlow::Continue(())
}

async fn process_message<S>(
    api_service: Arc<S>,
    user: &User,
    packet: SimConnectPacket,
) -> ControlFlow<(), ()>
where
    S: ApiService,
{
    // This needs to be abstracted away from here like actually
    match packet {
        SimConnectPacket::AtcID(id) => info!("[{}] Registration to {}", user.email(), id),
        SimConnectPacket::Fuel(fuel) => info!("[{}] Fuel state: {}", user.email(), fuel),
        SimConnectPacket::Airplane(airplane) => info!("[{}] Airplane: {}", user.email(), airplane),
        SimConnectPacket::Gps(gps) => info!("[{}] Location: {}", user.email(), gps),
        SimConnectPacket::OnGround(on_ground) => {
            info!("[{}] On Ground: {}", user.email(), on_ground)
        }
        SimConnectPacket::IsParked(is_parked) => {
            info!("[{}] Is Parked: {}", user.email(), is_parked)
        }
    }

    // The airplane variant checks if the player is in the correct airplane for the selected job etc
    // Send error if not and don't handle the flight

    // tracing::info!("< [{}]: {:?}", user.email(), packet);

    // On incoming packets, we use the internal api to store stuff to the database
    // We'll use Server Side Events (SSE) to keep the dashboard up to date with the state of the database

    ControlFlow::Continue(())
}
