use axum::extract::ws::{Message as AxumMessage, WebSocket, WebSocketUpgrade};
use axum::extract::{FromRequest, Path, Request, State};
use axum::http::{header, HeaderMap, StatusCode, Uri};
use axum::response::{IntoResponse, Response};
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message as WsMessage;

use crate::computer;
use crate::state::AppState;

pub async fn view_root(
    State(state): State<AppState>,
    Path(bot_id): Path<String>,
    req: Request,
) -> Response {
    proxy(state, bot_id, String::new(), req).await
}

pub async fn view_path(
    State(state): State<AppState>,
    Path((bot_id, rest)): Path<(String, String)>,
    req: Request,
) -> Response {
    proxy(state, bot_id, rest, req).await
}

async fn proxy(state: AppState, bot_id: String, rest: String, req: Request) -> Response {
    let upgrade = req
        .headers()
        .get(header::UPGRADE)
        .and_then(|value| value.to_str().ok())
        .map(|value| value.eq_ignore_ascii_case("websocket"))
        .unwrap_or(false);
    let ensure = upgrade || is_viewer_page(&rest) || rest.contains("websockify");
    let port = match upstream_port(&state, &bot_id, ensure).await {
        Ok(port) => port,
        Err(status) => return status.into_response(),
    };
    if upgrade {
        return match WebSocketUpgrade::from_request(req, &state).await {
            Ok(ws) => ws
                .on_upgrade(move |socket| proxy_socket(socket, port, rest))
                .into_response(),
            Err(error) => error.into_response(),
        };
    }
    if is_viewer_page(&rest) {
        return viewer_page().await;
    }
    http_proxy(port, &rest, req).await
}

fn is_viewer_page(rest: &str) -> bool {
    rest.is_empty()
        || rest == "vnc_lite.html"
        || rest == "vnc.html"
        || rest == "index.html"
        || rest == "embed.html"
}

async fn viewer_page() -> Response {
    let dir = std::env::var("LAZYBOY_WEB_DIR").unwrap_or_else(|_| "apps/web".into());
    let path = std::path::Path::new(&dir).join("vnc.html");
    let html = match tokio::fs::read_to_string(&path).await {
        Ok(body) => body,
        Err(_) => include_str!("../../../apps/web/vnc.html").to_string(),
    };
    let mut headers = HeaderMap::new();
    headers.insert(header::CONTENT_TYPE, "text/html; charset=utf-8".parse().unwrap());
    headers.insert(header::CACHE_CONTROL, "no-store".parse().unwrap());
    headers.insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, "*".parse().unwrap());
    (StatusCode::OK, headers, html).into_response()
}

async fn upstream_port(state: &AppState, bot_id: &str, ensure: bool) -> Result<u16, StatusCode> {
    let actor = state
        .bootstrap()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let bot = state
        .db
        .get_bot(&actor, bot_id)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;
    let computer = state
        .db
        .get_computer(bot.computer_id.as_deref().unwrap_or(""))
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;
    let computer_ref = computer::computer_ref(&computer).ok_or(StatusCode::NOT_FOUND)?;
    let screen = if ensure {
        match computer::ensure_bot_screen(state, &actor, bot_id, &computer, None).await {
            Ok(bound) => bound.row,
            Err(_) => state
                .db
                .get_screen(&computer.id, bot_id)
                .await
                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
        }
    } else {
        state
            .db
            .get_screen(&computer.id, bot_id)
            .await
            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
    };
    let session = state
        .sandbox
        .connect_screen(
            &computer_ref,
            true,
            &computer::adapter_context_for(&actor, bot_id, "view", screen.as_ref(), None),
        )
        .await
        .map_err(|_| StatusCode::BAD_GATEWAY)?;
    let url = session.url.ok_or(StatusCode::NOT_FOUND)?;
    let rewritten = rewrite_upstream(&url);
    let uri: Uri = rewritten.parse().map_err(|_| StatusCode::BAD_GATEWAY)?;
    uri.port_u16().ok_or(StatusCode::BAD_GATEWAY)
}

fn rewrite_upstream(url: &str) -> String {
    let host = std::env::var("LAZYBOY_SCREEN_UPSTREAM").unwrap_or_else(|_| "127.0.0.1".into());
    url.replace("127.0.0.1", &host)
        .replace("localhost", &host)
}

async fn http_proxy(port: u16, rest: &str, req: Request) -> Response {
    let host = std::env::var("LAZYBOY_SCREEN_UPSTREAM").unwrap_or_else(|_| "127.0.0.1".into());
    let path = if rest.is_empty() { "vnc_lite.html" } else { rest };
    let query = req.uri().query().unwrap_or_default();
    let url = if query.is_empty() {
        format!("http://{host}:{port}/{path}")
    } else {
        format!("http://{host}:{port}/{path}?{query}")
    };
    let client = reqwest::Client::new();
    let method = reqwest::Method::from_bytes(req.method().as_str().as_bytes()).unwrap_or(reqwest::Method::GET);
    match client.request(method, url).send().await {
        Ok(upstream) => {
            let status = StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
            let content_type = upstream
                .headers()
                .get(reqwest::header::CONTENT_TYPE)
                .and_then(|value| value.to_str().ok())
                .unwrap_or("application/octet-stream")
                .to_string();
            match upstream.bytes().await {
                Ok(bytes) => {
                    let mut headers = HeaderMap::new();
                    if let Ok(value) = content_type.parse() {
                        headers.insert(header::CONTENT_TYPE, value);
                    }
                    headers.insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, "*".parse().unwrap());
                    headers.insert(
                        header::HeaderName::from_static("cross-origin-resource-policy"),
                        "cross-origin".parse().unwrap(),
                    );
                    (status, headers, bytes).into_response()
                }
                Err(_) => StatusCode::BAD_GATEWAY.into_response(),
            }
        }
        Err(_) => StatusCode::BAD_GATEWAY.into_response(),
    }
}

async fn proxy_socket(mut client: WebSocket, port: u16, rest: String) {
    let host = std::env::var("LAZYBOY_SCREEN_UPSTREAM").unwrap_or_else(|_| "127.0.0.1".into());
    let path = if rest.is_empty() { "websockify".into() } else { rest };
    let url = format!("ws://{host}:{port}/{path}");
    let Ok((upstream, _)) = tokio_tungstenite::connect_async(url).await else {
        let _ = client.send(AxumMessage::Close(None)).await;
        return;
    };
    let (mut up_write, mut up_read) = upstream.split();
    loop {
        tokio::select! {
            incoming = client.recv() => {
                match incoming {
                    Some(Ok(AxumMessage::Binary(data))) => {
                        if up_write.send(WsMessage::Binary(data)).await.is_err() { break; }
                    }
                    Some(Ok(AxumMessage::Text(text))) => {
                        if up_write.send(WsMessage::Text(text.to_string().into())).await.is_err() { break; }
                    }
                    Some(Ok(AxumMessage::Ping(data))) => {
                        if up_write.send(WsMessage::Ping(data)).await.is_err() { break; }
                    }
                    Some(Ok(AxumMessage::Pong(data))) => {
                        if up_write.send(WsMessage::Pong(data)).await.is_err() { break; }
                    }
                    Some(Ok(AxumMessage::Close(_))) | None => break,
                    Some(Err(_)) => break,
                }
            }
            incoming = up_read.next() => {
                match incoming {
                    Some(Ok(WsMessage::Binary(data))) => {
                        if client.send(AxumMessage::Binary(data)).await.is_err() { break; }
                    }
                    Some(Ok(WsMessage::Text(text))) => {
                        if client.send(AxumMessage::Text(text.to_string().into())).await.is_err() { break; }
                    }
                    Some(Ok(WsMessage::Ping(data))) => {
                        if client.send(AxumMessage::Ping(data.into())).await.is_err() { break; }
                    }
                    Some(Ok(WsMessage::Pong(data))) => {
                        if client.send(AxumMessage::Pong(data.into())).await.is_err() { break; }
                    }
                    Some(Ok(WsMessage::Close(_))) | Some(Ok(WsMessage::Frame(_))) | None => break,
                    Some(Err(_)) => break,
                }
            }
        }
    }
}