use std::time::Duration;

use log::warn;
use rand::{rngs::OsRng, RngCore};
use rocket::{
    data::{Data, ToByteUnit},
    http::{CookieJar, Status},
    post,
    request::{FromRequest, Outcome},
    serde::json::{json, Json, Value},
    Request, State,
};
use sqlx::{MySql, Pool};

use crate::web::image_uploads::{prepare_image, ImageError, ImageUploads, MAX_UPLOAD_BYTES};

const UPLOAD_READ_TIMEOUT: Duration = Duration::from_secs(30);

pub fn mount(
    rocket: rocket::Rocket<rocket::Build>,
    uploads: ImageUploads,
) -> rocket::Rocket<rocket::Build> {
    let enabled = uploads.0.is_some();
    let rocket = rocket.manage(uploads);
    if enabled {
        rocket.mount("/dashboard", rocket::routes![upload_image])
    } else {
        rocket
    }
}

pub struct Uploader(u64);

#[rocket::async_trait]
impl<'r> FromRequest<'r> for Uploader {
    type Error = ();

    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        match uploader_id(request.cookies()) {
            Some(id) => Outcome::Success(Self(id)),
            None => Outcome::Error((Status::Unauthorized, ())),
        }
    }
}

fn uploader_id(cookies: &CookieJar<'_>) -> Option<u64> {
    cookies.get_private("userid")?.value().parse().ok().filter(|id| *id != 0)
}

type UploadResult = Result<Json<Value>, (Status, Json<Value>)>;

fn error(status: Status, message: &str) -> (Status, Json<Value>) {
    (status, Json(json!({ "error": message })))
}

// Requiring a non-form content type prevents cross-site HTML form submissions.
// Authorization is a request guard, so it runs before the body is read.
#[post("/api/images", format = "application/octet-stream", data = "<data>")]
pub async fn upload_image(
    uploader: Uploader,
    uploads: &State<ImageUploads>,
    pool: &State<Pool<MySql>>,
    data: Data<'_>,
) -> UploadResult {
    let service =
        uploads.0.as_ref().ok_or_else(|| error(Status::NotFound, "Image uploads disabled"))?;
    let _permit = service
        .slots
        .clone()
        .try_acquire_owned()
        .map_err(|_| error(Status::TooManyRequests, "Image uploads busy. Please try again."))?;
    // Bound the entire body read so stalled or trickling uploads release their slot.
    let bytes =
        tokio::time::timeout(UPLOAD_READ_TIMEOUT, data.open(MAX_UPLOAD_BYTES.bytes()).into_bytes())
            .await
            .map_err(|_| {
                error(Status::RequestTimeout, "Image upload timed out. Please try again.")
            })?
            .map_err(|_| error(Status::BadRequest, "Could not read image"))?;
    if !bytes.is_complete() {
        return Err(error(Status::PayloadTooLarge, "Images must be no larger than 8 MiB"));
    }
    let service_for_moderation = service.clone();
    let moderated = tokio::task::spawn_blocking(move || {
        prepare_image(service_for_moderation.model.as_ref(), &bytes.into_inner())
    })
    .await
    .map_err(|err| {
        warn!("Image moderation task failed: {}", err);
        error(Status::ServiceUnavailable, "Image moderation unavailable. Please try again.")
    })?;
    let image = moderated.map_err(|err| match err {
        ImageError::InvalidImage => error(
            Status::UnprocessableEntity,
            "Please upload a valid PNG, JPEG, GIF or WebP image (at most 16 million pixels)",
        ),
        ImageError::Nsfw => error(
            Status::UnavailableForLegalReasons,
            "This image was flagged as probably NSFW and cannot be uploaded.",
        ),
        ImageError::ModerationFailed => {
            error(Status::ServiceUnavailable, "Image moderation unavailable. Please try again.")
        }
    })?;

    let mut random = [0u8; 32];
    OsRng.fill_bytes(&mut random);
    let key = format!(
        "uploads/{}.png",
        random.iter().map(|byte| format!("{:02x}", byte)).collect::<String>()
    );
    let url = format!("{}/{}", service.public_url, key);
    let response =
        service.bucket.put_object_with_content_type(&key, &image, "image/png").await.map_err(
            |err| {
                warn!("Could not upload image to S3: {}", err);
                error(Status::BadGateway, "Could not store image. Please try again.")
            },
        )?;
    if !(200..300).contains(&response.status_code()) {
        warn!("S3 image upload returned {}", response.status_code());
        return Err(error(Status::BadGateway, "Could not store image. Please try again."));
    }
    if let Err(err) = sqlx::query("INSERT INTO image_uploads (uploader, path) VALUES (?, ?)")
        .bind(uploader.0)
        .bind(&key)
        .execute(pool.inner())
        .await
    {
        warn!("Could not record image upload: {}", err);
        match service.bucket.delete_object(&key).await {
            Ok(response) if (200..300).contains(&response.status_code()) => {}
            _ => warn!("Could not clean up untracked image object {}", key),
        }
        return Err(error(
            Status::InternalServerError,
            "Could not record image upload. Please try again.",
        ));
    }

    Ok(Json(json!({ "url": url })))
}

#[cfg(test)]
mod tests {
    use std::{io::Cursor, sync::Arc};

    use rocket::{http::ContentType, local::asynchronous::Client};
    use s3::{creds::Credentials, Bucket, Region};
    use tokio::sync::Semaphore;

    use super::*;
    use crate::web::image_uploads::ImageUploadService;

    fn service(unsafe_image: bool, endpoint: &str) -> ImageUploads {
        let model: &[u8] = if unsafe_image {
            include_bytes!("../../../../../tests/fixtures/image-upload-unsafe.onnx")
        } else {
            include_bytes!("../../../../../tests/fixtures/image-upload-safe.onnx")
        };
        let bucket = Bucket::new(
            "images",
            Region::Custom { region: "eu-west-2".into(), endpoint: endpoint.into() },
            Credentials::new(Some("test"), Some("test"), None, None, None).unwrap(),
        )
        .unwrap()
        .with_path_style();
        ImageUploads(Some(Arc::new(ImageUploadService {
            bucket,
            public_url: "https://images.example.com".into(),
            model: Some(nsfw::create_model(model).unwrap()),
            slots: Arc::new(Semaphore::new(2)),
        })))
    }

    fn image_bytes() -> Vec<u8> {
        let mut bytes = Cursor::new(Vec::new());
        image::DynamicImage::new_rgb8(2, 2).write_to(&mut bytes, image::ImageFormat::Png).unwrap();
        bytes.into_inner()
    }

    async fn client(uploads: ImageUploads) -> Client {
        let pool = sqlx::mysql::MySqlPoolOptions::new()
            .connect_lazy("mysql://unused:unused@127.0.0.1:1/unused")
            .unwrap();
        Client::tracked(mount(rocket::build().manage(pool), uploads)).await.unwrap()
    }

    #[rocket::get("/test-login")]
    fn login(cookies: &CookieJar<'_>) {
        cookies.add_private(("userid", "123"));
    }

    #[rocket::async_test]
    async fn disabled_upload_route_is_not_available() {
        let client = client(ImageUploads(None)).await;
        assert_eq!(
            client
                .post("/dashboard/api/images")
                .header(ContentType::Binary)
                .body(image_bytes())
                .dispatch()
                .await
                .status(),
            Status::NotFound
        );
    }

    #[rocket::async_test]
    async fn upload_requires_an_authenticated_private_cookie() {
        let client = client(service(false, "http://127.0.0.1:1")).await;
        for cookie in [None, Some("not-an-id"), Some("0")] {
            let mut request = client
                .post("/dashboard/api/images")
                .header(ContentType::Binary)
                .body(image_bytes());
            if let Some(cookie) = cookie {
                request = request.private_cookie(("userid", cookie));
            }
            assert_eq!(request.dispatch().await.status(), Status::Unauthorized);
        }
    }

    #[rocket::async_test]
    async fn unsafe_images_return_451_without_contacting_storage_or_database() {
        let client = client(service(true, "http://127.0.0.1:1")).await;
        let response = client
            .post("/dashboard/api/images")
            .header(ContentType::Binary)
            .private_cookie(("userid", "123456789012345678"))
            .body(image_bytes())
            .dispatch()
            .await;
        assert_eq!(response.status(), Status::UnavailableForLegalReasons);
        assert!(response.into_json::<Value>().await.unwrap()["error"]
            .as_str()
            .unwrap()
            .contains("NSFW"));
    }

    #[rocket::async_test]
    async fn rejects_invalid_oversized_and_form_uploads() {
        let client = client(service(false, "http://127.0.0.1:1")).await;
        for (body, status) in [
            (b"this is not an image".to_vec(), Status::UnprocessableEntity),
            (vec![0; MAX_UPLOAD_BYTES as usize + 1], Status::PayloadTooLarge),
        ] {
            let response = client
                .post("/dashboard/api/images")
                .header(ContentType::Binary)
                .private_cookie(("userid", "123456789012345678"))
                .body(body)
                .dispatch()
                .await;
            assert_eq!(response.status(), status);
        }
        let response = client
            .post("/dashboard/api/images")
            .header(ContentType::Form)
            .private_cookie(("userid", "123456789012345678"))
            .body("file=anything")
            .dispatch()
            .await;
        assert_eq!(response.status(), Status::NotFound);
    }

    #[rocket::async_test]
    async fn stalled_uploads_time_out_and_release_global_slots() {
        use rocket::{config::LogLevel, fairing::AdHoc, Config};
        use tokio::{
            io::{AsyncReadExt, AsyncWriteExt},
            net::TcpStream,
            sync::oneshot,
            time::{sleep, timeout},
        };

        let uploads = service(false, "http://127.0.0.1:1");
        let slots = uploads.0.as_ref().unwrap().slots.clone();
        let pool = sqlx::mysql::MySqlPoolOptions::new()
            .connect_lazy("mysql://unused:unused@127.0.0.1:1/unused")
            .unwrap();
        let config = Config {
            address: "127.0.0.1".parse().unwrap(),
            port: 0,
            log_level: LogLevel::Off,
            ..Config::debug_default()
        };
        let (send_port, receive_port) = oneshot::channel();
        let rocket =
            mount(rocket::custom(config).manage(pool).mount("/", rocket::routes![login]), uploads)
                .attach(AdHoc::on_liftoff("test address", |rocket| {
                    Box::pin(async move { send_port.send(rocket.config().port).unwrap() })
                }))
                .ignite()
                .await
                .unwrap();
        let shutdown = rocket.shutdown();
        let server = tokio::spawn(rocket.launch());
        let port = timeout(Duration::from_secs(5), receive_port).await.unwrap().unwrap();
        let base = format!("http://127.0.0.1:{port}");
        let http =
            reqwest::Client::builder().no_proxy().timeout(Duration::from_secs(5)).build().unwrap();
        let response = http.get(format!("{base}/test-login")).send().await.unwrap();
        let cookie = response.headers()["set-cookie"]
            .to_str()
            .unwrap()
            .split(';')
            .next()
            .unwrap()
            .to_string();
        let mut stalled = Vec::new();
        for _ in 0..2 {
            let mut socket = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
            socket.write_all(format!(
                "POST /dashboard/api/images HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/octet-stream\r\nContent-Length: 1024\r\nCookie: {cookie}\r\nConnection: close\r\n\r\nincomplete image"
            ).as_bytes()).await.unwrap();
            stalled.push(socket);
        }
        timeout(Duration::from_secs(5), async {
            while slots.available_permits() != 0 {
                sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .unwrap();

        let upload = || {
            http.post(format!("{base}/dashboard/api/images"))
                .header("cookie", &cookie)
                .header("content-type", "application/octet-stream")
                .body("not an image")
                .send()
        };
        assert_eq!(upload().await.unwrap().status(), Status::TooManyRequests.code);
        for mut socket in stalled {
            let mut response = String::new();
            timeout(
                UPLOAD_READ_TIMEOUT + Duration::from_secs(5),
                socket.read_to_string(&mut response),
            )
            .await
            .unwrap()
            .unwrap();
            assert!(response.starts_with("HTTP/1.1 408"), "{response}");
            assert!(response.contains("Image upload timed out"), "{response}");
        }
        assert_eq!(slots.available_permits(), 2);
        assert_eq!(upload().await.unwrap().status(), Status::UnprocessableEntity.code);

        drop(http);
        shutdown.notify();
        timeout(Duration::from_secs(5), server).await.unwrap().unwrap().unwrap();
    }

    #[sqlx::test]
    async fn stores_moderated_image_and_tracks_uploader(pool: Pool<MySql>) {
        use tokio::{
            io::{AsyncReadExt, AsyncWriteExt},
            net::TcpListener,
        };

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let endpoint = format!("http://{}", listener.local_addr().unwrap());
        let storage = tokio::spawn(async move {
            let (mut socket, _) = listener.accept().await.unwrap();
            let mut received = Vec::new();
            let mut buffer = [0; 4096];
            let (headers, body) = loop {
                let count = socket.read(&mut buffer).await.unwrap();
                assert!(count > 0);
                received.extend_from_slice(&buffer[..count]);
                if let Some(end) = received.windows(4).position(|window| window == b"\r\n\r\n") {
                    let headers = String::from_utf8(received[..end].to_vec()).unwrap();
                    let length: usize = headers
                        .lines()
                        .find_map(|line| {
                            line.to_lowercase()
                                .strip_prefix("content-length:")
                                .map(|length| length.trim().parse().unwrap())
                        })
                        .unwrap();
                    if received.len() >= end + 4 + length {
                        break (headers, received[end + 4..end + 4 + length].to_vec());
                    }
                }
            };
            socket
                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
                .await
                .unwrap();
            (headers, body)
        });
        let client =
            Client::tracked(mount(rocket::build().manage(pool.clone()), service(false, &endpoint)))
                .await
                .unwrap();
        let response = client
            .post("/dashboard/api/images")
            .header(ContentType::Binary)
            .private_cookie(("userid", "123456789012345678"))
            .body(image_bytes())
            .dispatch()
            .await;
        assert_eq!(response.status(), Status::Ok);
        let response = response.into_json::<Value>().await.unwrap();
        let url = response["url"].as_str().unwrap();
        let row: (u64, String) =
            sqlx::query_as("SELECT uploader, path COLLATE utf8mb4_unicode_ci FROM image_uploads")
                .fetch_one(&pool)
                .await
                .unwrap();
        assert_eq!(
            row,
            (
                123456789012345678,
                url.strip_prefix("https://images.example.com/").unwrap().to_string()
            )
        );
        let (headers, body) = storage.await.unwrap();
        assert!(headers.starts_with("PUT /images/uploads/"));
        assert!(headers.to_lowercase().contains("content-type: image/png"));
        assert_eq!(image::load_from_memory(&body).unwrap().width(), 2);
        assert!(url.starts_with("https://images.example.com/uploads/"));
        assert!(url.ends_with(".png"));
    }
}