// Copyright © 2020 Jakob L. Kreuze <zerodaysfordays@sdf.org>
//
// This file is part of Kona.
//
// Kona is free software; you can redistribute it and/or modify it under
// the terms of the GNU Affero General Public License as published by
// the Free Software Foundation; either version 3 of the License, or (at
// your option) any later version.
//
// Kona is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public
// License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with Kona. If not, see <http://www.gnu.org/licenses/>.

#![feature(proc_macro_hygiene, decl_macro)]

#[macro_use]
extern crate anyhow;
extern crate image;
#[macro_use]
extern crate rocket;
extern crate rocket_contrib;
#[macro_use]
extern crate serde_derive;

use std::convert::AsRef;
use std::fs;
use std::io::Write;
use std::os::unix::fs::symlink;
use std::path::Path;
use std::sync::Mutex;

use anyhow::Result;
use image::ImageFormat;
use multipart::{MultipartFormData, MultipartFormDataField, MultipartFormDataOptions};
use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};
use rocket::http::ContentType;
use rocket::request::Form;
use rocket::{Data, State};
use rocket_contrib::json::Json;
use rocket_contrib::serve::StaticFiles;
use rocket_contrib::templates::Template;

use kona::{Keyset, TagDatabase};

/// Maximum number of results returned by any API endpoint.
const RESULTS_PER_QUERY: i64 = 50;

/// Shorthand for the state that is passed to all handlers.
type SiteState<'a> = State<'a, Mutex<DatabaseConnection>>;

/// Bundle of everything pertaining to database access.
///
/// SQLite handles concurrent access reasonably well, but Kona is designed for
/// the single-user use case. If this turns out to be a problem, the
/// `TagDatabase` connection will be nixed.
struct DatabaseConnection {
    tb: TagDatabase,
    file_dir: String,
}

impl DatabaseConnection {
    fn new() -> Result<Self> {
        // FIXME: It is somewhat of an arbitrary choice that the directory
        // containing symbolic links and thumbnails be relative to the binary.
        let bin = std::env::current_exe()?.canonicalize()?;
        let bin_dir = bin.parent().unwrap();
        let file_dir = bin_dir.join("files/").to_str().unwrap().to_string();
        let tb = TagDatabase::new(&bin_dir.join("kona.db"))?;
        Ok(DatabaseConnection { tb, file_dir })
    }
}

/// Serialization protocol for a file entry, used by the 'posts' endpoint.
#[derive(Debug, Serialize)]
struct FileResult {
    id: i64,
    filename: String,
    thumb_filename: String,
    orig_filename: String,
    tags: Vec<String>,
}

/// Serialization protocol for a tag entry, used by the 'tags' endpoint.
type TagResult = (i64, String);

/// Return whether or not `filename` ends with one of `extensions`.
fn has_extension<T, S>(filename: T, extensions: &[S]) -> bool
where
    T: AsRef<Path>,
    S: AsRef<str>,
{
    filename
        .as_ref()
        .extension()
        .and_then(|ext| ext.to_str())
        .map(|ext| extensions.iter().any(|other| ext == other.as_ref()))
        .unwrap_or(false)
}

/// Return whether or not `filename` ends with the extension of a known video file type.
fn is_video<T: AsRef<Path>>(filename: T) -> bool {
    has_extension(filename, &vec!["mp4", "mkv", "webm"])
}

/// Return whether or not `filename` ends with the extension of a known image file type.
fn is_image<T: AsRef<Path>>(filename: T) -> bool {
    has_extension(filename, &vec!["png", "jpg", "jpeg", "gif"])
}

/// Return the store file name of `filename`.
fn store_filename<T: AsRef<Path>>(id: i64, path: T) -> String {
    if let Some(filename) = path.as_ref().file_name() {
        let filename = filename.to_string_lossy();
        if let Some(n) = filename.rfind('.') {
            format!("{id}.{ext}", id = id, ext = &filename[n + 1..])
        } else {
            id.to_string()
        }
    } else {
        id.to_string()
    }
}

/// Return the store file name of the thumbnail of `filename`.
fn thumb_filename<T: AsRef<Path>>(id: i64, path: T) -> String {
    if let Some(ext) = path.as_ref().extension() {
        if let Some(ext) = ext.to_str() {
            let ext = if is_video(path.as_ref()) { "png" } else { ext };
            format!("{id}_thumb.{ext}", id = id, ext = ext)
        } else {
            format!("{}_thumb", id)
        }
    } else {
        format!("{}_thumb", id)
    }
}

/// Macro for obtaining a reference to the `tb` field of `SiteState`.
///
/// Ownership boundaries are quite difficult with mutexes, so I've opted to just
/// expand the call for locking the database.
macro_rules! lock_database {
    ($conn:expr) => {{
        $conn
            .inner()
            .lock()
            .map_err(|_| anyhow!("Could not lock database."))?
    }};
}

#[get("/tags")]
fn get_tags(conn: SiteState) -> Result<Json<Vec<TagResult>>> {
    let tb = &lock_database!(conn).tb;
    Ok(Json(tb.top_tags(None)?))
}

#[get("/posts?<query>&<last>")]
fn get_posts(conn: SiteState, query: String, last: Option<i64>) -> Result<Json<Vec<FileResult>>> {
    let tb = &lock_database!(conn).tb;

    let keyset = if let Some(id) = last {
        Keyset {
            last_id: id,
            how_many: RESULTS_PER_QUERY,
        }
    } else {
        Keyset {
            last_id: i64::MAX,
            how_many: RESULTS_PER_QUERY,
        }
    };

    Ok(Json(
        tb.query(&kona::parse_query(query)?, Some(keyset))?
            .iter()
            .map(|file| FileResult {
                id: file.id,
                filename: store_filename(file.id, &file.path),
                thumb_filename: thumb_filename(file.id, &file.path),
                orig_filename: file.path.clone(),
                tags: file.tags.clone(),
            })
            .collect(),
    ))
}

#[get("/posts/<id>")]
fn get_post(conn: SiteState, id: i64) -> Result<Json<FileResult>> {
    let tb = &lock_database!(conn).tb;
    let file = tb.file_by_id(id)?;

    Ok(Json(FileResult {
        id: file.id,
        filename: store_filename(file.id, &file.path),
        thumb_filename: thumb_filename(file.id, &file.path),
        orig_filename: file.path.clone(),
        tags: file.tags,
    }))
}

#[derive(FromForm)]
struct UpdateTags {
    tags: String,
}

#[post("/posts/<id>", data = "<form>")]
fn update_post(conn: SiteState, id: i64, form: Form<UpdateTags>) -> Result<()> {
    let tb = &lock_database!(conn).tb;

    let mut adding = Vec::new();
    let mut removing = Vec::new();

    for atom in form.tags.split(',') {
        if atom.starts_with("-") {
            removing.push(&atom[1..]);
        } else {
            adding.push(atom);
        }
    }

    if !adding.is_empty() {
        tb.tag_file(id, &adding[..])?;
    }
    if !removing.is_empty() {
        tb.untag_file(id, &removing[..])?;
    }

    Ok(())
}

#[post("/posts", data = "<data>")]
fn put_post(conn: SiteState, content_type: &ContentType, data: Data) -> Result<()> {
    let conn = lock_database!(conn);

    let options = MultipartFormDataOptions::with_multipart_form_data_fields(vec![
        MultipartFormDataField::text("tags"),
        MultipartFormDataField::raw("image")
            .size_limit(32 * 1024 * 1024)
            .content_type_by_string(Some(multipart::mime::IMAGE_STAR))?,
    ]);

    let mut data = MultipartFormData::parse(content_type, data, options)?;
    let raw = if let Some(mut data) = data.raw.remove("image") {
        data.remove(0)
    } else {
        bail!("No `image` field.")
    };

    let filename = raw
        .file_name
        .unwrap_or_else(|| thread_rng().sample_iter(&Alphanumeric).take(30).collect());
    let path = Path::new(&conn.file_dir).join(&filename);
    fs::File::create(&path)?.write_all(&raw.raw)?;

    let tags = if let Some(mut vec) = data.texts.remove("tags") {
        let text = vec.remove(0).text;
        text.split(',').map(String::from).collect()
    } else {
        vec![]
    };
    let file = conn.tb.add_file(&path, &tags)?;
    add_to_store(&file, &conn.file_dir);

    Ok(())
}

#[get("/")]
fn index(conn: SiteState) -> Result<Template> {
    display_posts(conn, String::from(""), None)
}

#[get("/upload")]
fn upload() -> Template {
    Template::render("upload", None::<()>)
}

#[get("/posts?<query>&<last>")]
fn display_posts(conn: SiteState, query: String, last: Option<i64>) -> Result<Template> {
    #[derive(Serialize)]
    struct TagResult {
        name: String,
        count: i64,
    }

    #[derive(Serialize)]
    struct Context {
        tags: Vec<TagResult>,
        images: Vec<FileResult>,
        next_page: String,
    }

    // This is done in a block because `tb` needs to go out of scope before
    // calling out to `get_posts` to perform the query. Otherwise, we deadlock
    // because the mutex is locked for the entirety of `display_posts`.
    let tag_results = {
        let tb = &lock_database!(conn).tb;
        tb.top_tags(Some(20))?
            .iter()
            .map(|tup| TagResult {
                name: tup.1.clone(),
                count: tup.0,
            })
            .collect()
    };

    let images = get_posts(conn, query.clone(), last)?.into_inner();

    Ok(Template::render(
        "index",
        Context {
            tags: tag_results,
            next_page: if images.len() as i64 == RESULTS_PER_QUERY {
                let last_image = images.last().unwrap();
                format!("/posts?query={}&last={}", query, last_image.id)
            } else {
                String::from("")
            },
            images,
        },
    ))
}

#[get("/posts/<id>")]
fn display_post(conn: SiteState, id: i64) -> Result<Template> {
    #[derive(Serialize)]
    struct Context {
        id: i64,
        tags: Vec<String>,
        filename: String,
        orig_filename: String,
        is_image: bool,
        is_video: bool,
        is_unknown: bool,
    }

    let tb = &lock_database!(conn).tb;

    let file = tb.file_by_id(id)?;

    let is_image = is_image(&file.path);
    let is_video = is_video(&file.path);
    let orig_filename = String::from(
        Path::new(&file.path)
            .file_name()
            .ok_or_else(|| anyhow!("Directory indexed and returned by query"))?
            .to_str()
            .ok_or_else(|| anyhow!("Could not parse filename"))?,
    );

    Ok(Template::render(
        "image",
        Context {
            id,
            orig_filename,
            tags: file.tags,
            filename: store_filename(file.id, &file.path),
            is_image,
            is_video,
            is_unknown: !is_image && !is_video,
        },
    ))
}

/// Create a thumbnail for `file` in the store at `file_dir`.
fn add_to_store(file: &kona::File, file_dir: &str) {
    let store_path = store_filename(file.id, &file.path);
    let store_path = Path::new(file_dir).join(store_path);
    symlink(&file.path, &store_path).ok();
    if !store_path.exists() {
        panic!("could not create symlink {:?}!", store_path);
    }

    let thumb_path = thumb_filename(file.id, &file.path);
    let thumb_path = Path::new(file_dir).join(&thumb_path);
    if !thumb_path.exists() {
        if let Ok(im) = image::open(&store_path) {
            let out = &mut fs::File::create(&thumb_path).unwrap();
            im.thumbnail(100, 100)
                .write_to(out, ImageFormat::Png)
                .unwrap();
        } else if is_video(&file.path) {
            // FIXME: Can we do this... natively?
            std::process::Command::new("ffmpeg")
                .arg("-i")
                .arg(&store_path)
                .arg("-vframes")
                .arg("1")
                .arg("-s")
                .arg("100x100")
                .arg(&thumb_path)
                .output()
                .expect("failed to execute process");
        } else {
            let src = Path::new("./static/img/missing_thumb.png");
            symlink(src.canonicalize().unwrap(), &thumb_path).unwrap();
            if !thumb_path.exists() {
                panic!("could not create symlink {:?}!", thumb_path);
            }
        }
    }
}

/// Maybe initialize the directory for storing image symlinks and thumbnails.
fn create_image_directory(conn: &DatabaseConnection) {
    fs::create_dir(&conn.file_dir).ok();
    let query = kona::parse_query("").unwrap();
    for file in conn.tb.query(&query, None).unwrap() {
        add_to_store(&file, &conn.file_dir);
    }
}

fn main() {
    let conn = DatabaseConnection::new().unwrap();
    create_image_directory(&conn);
    rocket::ignite()
        .mount("/public", StaticFiles::from("./static/"))
        .mount("/image", StaticFiles::from(&conn.file_dir))
        .mount("/", routes![index, upload, display_posts, display_post])
        .mount(
            "/api",
            routes![get_posts, get_tags, get_post, put_post, update_post],
        )
        .attach(Template::fairing())
        .manage(Mutex::new(conn))
        .launch();
}