mod samplers;
mod invoking_llms;
mod db;
mod config;
mod web_file_uploads;
mod web_app_state;
pub mod dedicated_ai_server;

use axum::{
    extract::{DefaultBodyLimit, Form, Query, State},
    response::{Html, IntoResponse, Redirect, Response},
    routing::{get, post},
    Router,
};
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum_extra::extract::cookie::{Cookie, CookieJar, SameSite};
use serde::Deserialize;
use std::future::Future;
use std::net::SocketAddr;
use tera::{Tera};
use tower_http::services::{ServeDir, ServeFile};

use samplers::random_junk::random_string;
use config::load_config_default;
use db::{connect_db, init_database};
use web_file_uploads::{get_file, upload_get, upload_post};
use web_app_state::{AppState, AppStateInner, AuthenticatedUserId};
use dedicated_ai_server::connection::connect_to_dedicated_ai_server;
use dedicated_ai_server::api::MessageInChat;
use crate::dedicated_ai_server::connection::{MessagePiece, MessagePiecePayload};
use frontend_protocol::{
    DekuBytes,
    UserRequest,
    UserResponse,
    UserResponseChatCompletion,
};

async fn init_app_state() -> Result<(AppState, tokio::task::JoinHandle<()>), Box<dyn std::error::Error>> {
    let config = load_config_default()?;
    let db = connect_db(&config).await?;
    let tera = Tera::new("frontend/pages/**/*.html")?;
    let (dedicated_ai, dedicated_ai_task_handler) = connect_to_dedicated_ai_server(
        config.dedicated_ai_server_address.clone(),
        config.dedicated_ai_server_port,
        config.dedicated_ai_server_secret.clone(),
    );
    Ok((std::sync::Arc::new(AppStateInner { config, db, tera, dedicated_ai }), dedicated_ai_task_handler))
}

enum PasscodeAuthenticationResult{
    WrongPassword,
    Ok(AuthenticatedUserId)
}

enum WebsiteAuthenticationResult {
    NoCookie,
    SomeCookie(PasscodeAuthenticationResult),
}

async fn website_authentication_with_passcode(
    state: &AppStateInner, passcode: &str) -> PasscodeAuthenticationResult {
    let row = state.db.query_opt(
        "SELECT id, name FROM person WHERE passcode = $1 LIMIT 1",
        &[&passcode],
    ).await;

    match row {
        Ok(Some(row)) => PasscodeAuthenticationResult::Ok(AuthenticatedUserId {
            id: row.get::<_, i32>(0),
            name: row.get::<_, String>(1),
        }),
        Ok(None) => PasscodeAuthenticationResult::WrongPassword,
        Err(err) => {
            eprintln!("auth query failed: {err}");
            PasscodeAuthenticationResult::WrongPassword
        }
    }
}

async fn website_authentication_with_cookie(
    state: &AppStateInner,
    jar: &CookieJar,
) -> WebsiteAuthenticationResult {
    let auth_cookie = jar.get("auth");
    match auth_cookie {
        None => WebsiteAuthenticationResult::NoCookie,
        Some(x) => {
            let passcode = x.value();
            WebsiteAuthenticationResult::SomeCookie(website_authentication_with_passcode(state, passcode).await)
        },
    }
}


fn axum_handler_with_auth<H, Fut, Res, T>(
    handler: H,
    state: AppState,
) -> impl Fn(CookieJar, T) -> std::pin::Pin<Box<dyn Future<Output = Response> + Send>> + Clone + Send + 'static
where
    H: (Fn(T, AppState, AuthenticatedUserId) -> Fut) + Send + Clone + 'static,
    Fut: Future<Output = Res> + Send + 'static,
    Res: IntoResponse + 'static,
    T: Send + 'static,
{
    move |jar: CookieJar, args: T| {
        let state = state.clone();
        let handler = handler.clone();
        Box::pin(async move {
            let res = website_authentication_with_cookie(state.as_ref(), &jar).await;
            match res {
                WebsiteAuthenticationResult::SomeCookie(PasscodeAuthenticationResult::Ok(user)) => {
                    handler(args, state, user).await.into_response()
                }
                WebsiteAuthenticationResult::NoCookie => {
                    Redirect::to("/login").into_response()
                }
                WebsiteAuthenticationResult::SomeCookie(PasscodeAuthenticationResult::WrongPassword) => {
                    Redirect::to("/login?error=cookie").into_response()
                }
            }
        })
    }
}

#[derive(Deserialize)]
struct LoginPageQuery {
    error: Option<String>,
}

#[derive(Deserialize)]
struct LoginPageForm {
    passcode: String,
    csrf_token: String,
}

async fn login_get(
    State(state): State<AppState>,
    jar: CookieJar,
    Query(query): Query<LoginPageQuery>,
) -> (CookieJar, Html<String>) {
    let cur_auth: WebsiteAuthenticationResult =
        website_authentication_with_cookie(state.as_ref(), &jar).await;

    let csrf = random_string(32);
    let jar = jar.add(
        Cookie::build(("csrf", csrf.clone()))
            .path("/")
            .http_only(true)
            .same_site(SameSite::Strict)
            .build(),
    );

    let error = match query.error.as_deref() {
        Some("cookie") => Some("Incorrect session cookie"),
        Some("password") => Some("Invalid passcode"),
        Some("csrf") => Some("Implicit log in attempt aborted!"),
        _ => None,
    };

    let mut ctx = tera::Context::new();
    ctx.insert("csrf", &csrf);
    if let Some(msg) = error {
        ctx.insert("error", msg);
    }
    if let WebsiteAuthenticationResult::SomeCookie(PasscodeAuthenticationResult::Ok(cur_user)) = cur_auth {
        ctx.insert("logged_in_cur_user", &cur_user.name);
    }

    let body = state.tera.render("login.html", &ctx).expect("render index");
    (jar, Html(body))
}

async fn login_post(
    State(state): State<AppState>,
    jar: CookieJar,
    Form(form): Form<LoginPageForm>
) -> impl IntoResponse {
    let csrf_ok = jar
        .get("csrf")
        .map(|cookie| cookie.value())
        == Some(form.csrf_token.as_str());

    if !csrf_ok {
        return Redirect::to("/login?error=csrf").into_response();
    }

    let res = website_authentication_with_passcode(&state, &form.passcode);

    match res.await {
        PasscodeAuthenticationResult::Ok(_) => {
            let jar = jar.add(
                Cookie::build(("auth", form.passcode))
                    .path("/")
                    .http_only(true)
                    .same_site(SameSite::Strict)
                    .build(),
            );
            (jar, Redirect::to("/")).into_response()
        }
        PasscodeAuthenticationResult::WrongPassword =>
            Redirect::to("/login?error=password").into_response()
    }
}


async fn index(
    _: (),
    state: AppState,
    user: AuthenticatedUserId,
) -> impl IntoResponse {
    let mut ctx = tera::Context::new();
    ctx.insert("logged_in_cur_user", &user.name);
    let body = state.tera
        .render("index.html", &ctx)
        .expect("render index");
    Html(body)
}


async fn welcome(
    _: (),
    _state: AppState,
    _user: AuthenticatedUserId,
) -> impl IntoResponse {
    let body = _state.tera
        .render("welcome.html", &tera::Context::new())
        .expect("render welcome");
    Html(body)
}

async fn board(
    _: (),
    _state: AppState,
    _user: AuthenticatedUserId,
) -> impl IntoResponse {
    let body = _state.tera
        .render("board.html", &tera::Context::new())
        .expect("render welcome");
    Html(body)
}

async fn chat(
    ws: Option<WebSocketUpgrade>,
    state: AppState,
    user: AuthenticatedUserId,
) -> Response {
    if let Some(ws) = ws {
        return ws
            .on_upgrade(move |socket| handle_chat_socket(socket, state.clone()))
            .into_response();
    }

    let mut ctx = tera::Context::new();
    ctx.insert("logged_in_cur_user", &user.name);
    let body = state
        .tera
        .render("chat.html", &ctx)
        .expect("render chat");
    Html(body).into_response()
}

async fn handle_chat_socket(mut socket: WebSocket, state: AppState) {
    'outer: loop {
        let msg = match socket.recv().await {
            Some(Ok(msg)) => msg,
            Some(Err(_)) => break,
            None => break,
        };

        match msg {
            Message::Binary(bytes) => {
                let request = match UserRequest::from_bytes(&bytes) {
                    Ok(request) => request,
                    Err(err) => {
                        eprintln!("[chat] failed to decode request: {err:#}");
                        continue;
                    }
                };

                let payload = match request {
                    UserRequest::ChatCompletion(payload) => payload,
                    UserRequest::ChatCompletionCancellation => {
                        eprintln!("[chat] protocol error: unexpected cancellation without active request");
                        break;
                    }
                };

                let messages = payload
                    .messages
                    .into_iter()
                    .map(|msg| MessageInChat {
                        role: msg.role,
                        content: msg.content,
                    })
                    .collect();

                let (mut response_rx, request_id) = match state.dedicated_ai.
                    send_chat_completion(messages)
                {
                    Ok(rx) => rx,
                    Err(err) => {
                        eprintln!("[chat] failed to send request: {err:#}");
                        let response = UserResponse::ChatCompletionCancellation;
                        // todo: make to_bytes nofail
                        if let Ok(bytes) = response.to_bytes() {
                            let _ = socket.send(Message::Binary(bytes)).await;
                        }
                        continue;
                    }
                };

                loop {
                    tokio::select! {
                        piece = response_rx.recv() => {
                            let piece = match piece {
                                Some(piece) => piece,
                                None => break,
                            };
                            let (response, should_break) = match piece {
                                MessagePiece::Piece(MessagePiecePayload(text)) => (
                                    UserResponse::ChatCompletion(UserResponseChatCompletion {
                                        piece: text,
                                    }),
                                    false,
                                ),
                                MessagePiece::End => (
                                    UserResponse::ChatCompletionEnd,
                                    true,
                                ),
                                MessagePiece::Cancelled | MessagePiece::DedicatedServerDisconnected => (
                                    UserResponse::ChatCompletionCancellation,
                                    true,
                                ),
                            };

                            let bytes = match response.to_bytes() {
                                Ok(bytes) => bytes,
                                Err(err) => {
                                    eprintln!("[chat] failed to encode response: {err:#}");
                                    break 'outer;
                                }
                            };
                            if socket.send(Message::Binary(bytes)).await.is_err() {
                                break 'outer;
                            }
                            if should_break {
                                break;
                            }
                        }
                        msg = socket.recv() => {
                            let msg = match msg {
                                Some(Ok(msg)) => msg,
                                Some(Err(_)) => break 'outer,
                                None => break 'outer,
                            };

                            match msg {
                                Message::Binary(bytes) => {
                                    let request = match UserRequest::from_bytes(&bytes) {
                                        Ok(request) => request,
                                        Err(err) => {
                                            eprintln!("[chat] failed to decode request: {err:#}");
                                            break 'outer;
                                        }
                                    };

                                    match request {
                                        UserRequest::ChatCompletionCancellation => {
                                            if let Err(err) = state
                                                .dedicated_ai
                                                .send_chat_completion_cancellation(request_id)
                                            {
                                                eprintln!("[chat] failed to send cancellation: {err:#}");
                                                break 'outer;
                                            }
                                        }
                                        UserRequest::ChatCompletion(_) => {
                                            eprintln!("[chat] protocol error: chat completion while receiving");
                                            break 'outer;
                                        }
                                    }
                                }
                                Message::Close(_) => break 'outer,
                                _ => {
                                    eprintln!("[chat] protocol error: unexpected non-binary message");
                                    break 'outer;
                                }
                            }
                        }
                    }
                }
            }
            Message::Close(_) => {
                println!(" [debug] websocket closed");
                break;
            }
            _ => {
                eprintln!("[chat] protocol error: unexpected non-binary message");
                break;
            }
        }
    }
}


pub async fn run_server() -> Result<(), Box<dyn std::error::Error>> {
    let (state, mut dedicated_ai_task) = init_app_state().await.expect("lol");
    let app = Router::new()
        .route("/login", get(login_get))
        .route("/login", post(login_post))
        .route("/upload", get(axum_handler_with_auth(upload_get, state.clone())))
        .route("/upload", post(axum_handler_with_auth(upload_post, state.clone())))
        .route("/get/:id", get(axum_handler_with_auth(get_file, state.clone())))
        .route("/welcome", get(axum_handler_with_auth(welcome, state.clone())), )
        .route("/board", get(axum_handler_with_auth(board, state.clone())), )
        .route("/chat", get(axum_handler_with_auth(chat, state.clone())), )
        .route("/", get(axum_handler_with_auth(index, state.clone())), )
        .nest_service("/static", ServeDir::new("frontend/static"))
        .nest_service("/pkg/frontend.js", ServeFile::new("target_pkg/frontend.js"))
        .nest_service("/pkg/frontend_bg.wasm", ServeFile::new("target_pkg/frontend_bg.wasm"))
        .layer(DefaultBodyLimit::disable())
        .with_state(state);

    let addr: SocketAddr = "127.0.0.1:3000".parse().expect("valid socket addr");
    println!("listening on http://{addr}");
    let listener = tokio::net::TcpListener::bind(addr)
        .await
        .expect("bind failed");
    let mut server_task = tokio::spawn(async move { axum::serve(listener, app).await });

    tokio::select! {
        res = &mut server_task => {
            dedicated_ai_task.abort();
            res??;
        }
        res = &mut dedicated_ai_task => {
            server_task.abort();
            res.map_err(|err| {
                std::io::Error::new(
                    std::io::ErrorKind::Other,
                    format!("dedicated ai task failed: {err}"),
                )
            })?;
            return Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "dedicated ai task ended unexpectedly",
            )
            .into());
        }
    }
    Ok(())
}

pub async fn init_db() -> Result<(), Box<dyn std::error::Error>> {
    let config = load_config_default()?;
    init_database(&config).await?;
    Ok(())
}
