use std::{net::SocketAddr, path::{Path, PathBuf}, ffi::{OsStr, OsString}, fs::File, io::{BufReader, Seek, Cursor, BufRead}, env::args_os, os::unix::prelude::OsStrExt, cmp::Reverse, sync::{Arc, RwLock}, collections::{HashMap, HashSet}, time::{Instant, Duration}};

use anyhow::{Context, Result, anyhow};
use axum::{Router, routing::{get, post}, response::{IntoResponse, Redirect, Response}, http::{StatusCode, header}, extract::{self, State}, Form, handler::Handler};
use axum_sessions::{async_session::{MemoryStore, SessionStore}, SessionLayer, extractors::{ReadableSession, WritableSession}};
use axum_template::{RenderHtml, engine::Engine};
use chacha20poly1305::{XChaCha20Poly1305, KeyInit, AeadCore, aead::{OsRng, rand_core::RngCore, Aead}, XNonce};
use chrono::{DateTime, Utc};
use error::AppError;
use exif::Exif;
use image::{imageops::FilterType, DynamicImage};
use serde::{Serialize, Deserialize};
use tokio::{task, sync::Semaphore};
use tower_http::{trace::{self, TraceLayer}, compression::CompressionLayer};
use tracing::Level;
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};

mod config;
mod error;

#[global_allocator]
static GLOBAL: jemallocator::Jemalloc = jemallocator::Jemalloc;

// TODO
// * Polish the css
// * Brute-force protection / rate-limiting
// * Support for headlines
// * Image cache cleanup
// * Limit session time (inactivity timeout)

type TemplateEngine = Engine<minijinja::Environment<'static>>;
type ImageDir = PathBuf;
type SecretKey = [u8; 64];
type ImageListCache = Arc<RwLock<Vec<ImageInfo>>>;
type ImageCache = Arc<RwLock<HashMap<OsString, (DateTime<Utc>, Vec<u8>)>>>;
type Config = Arc<RwLock<config::Config>>;
type CpuTaskLimit = Arc<Semaphore>;

#[derive(Clone, extract::FromRef)]
struct ApplicationState {
    config: Config,
    engine: TemplateEngine,
    image_cache: ImageCache,
    image_dir: ImageDir,
    image_list_cache: ImageListCache,
    secret_key: SecretKey,
    cpu_task_limit: CpuTaskLimit,
}

#[tokio::main]
async fn main() {
    let args = args_os()
        .skip(1)
        .collect::<Vec<_>>();
    let args = args
        .windows(2)
        .next()
        .expect("Usage: image-gallery IMAGE_DIRECTORY PORT");

    let image_path = &args[0];
    let port = args[1].to_string_lossy().parse::<u16>().expect("Invalid port specified");

    let mut secret_key: SecretKey = [0; 64];
    OsRng.fill_bytes(&mut secret_key);

    let default_tracing = "image_gallery=debug,tower_http=error".into();
    let tracing_filter = EnvFilter::try_from_default_env().unwrap_or(default_tracing);
    let tracing_formatter = tracing_subscriber::fmt::layer()
        .with_target(false)
        .compact();
    tracing_subscriber::registry()
        .with(tracing_filter)
        .with(tracing_formatter)
        .init();

    let sessions = MemoryStore::new();
    let image_dir = PathBuf::from(image_path);
    let config = Config::default();
    let image_cache = ImageCache::default();
    let image_list_cache = ImageListCache::default();

    tokio::spawn(update_config_and_image_list_cache_job(
        image_dir.clone(),
        config.clone(),
        image_list_cache.clone(),
        sessions.clone(),
    ));

    let mut jinja = minijinja::Environment::new();
    jinja.set_auto_escape_callback(|_| minijinja::AutoEscape::Html);
    jinja.add_template("base", include_str!("../templates/base.html")).unwrap();
    jinja.add_template("index", include_str!("../templates/index.html")).unwrap();
    jinja.add_template("login", include_str!("../templates/login.html")).unwrap();

    let session_layer = SessionLayer::new(sessions, &secret_key)
        .with_cookie_name("session")
        .with_session_ttl(None);

    let app = Router::new()
        .route("/", get(index.layer(CompressionLayer::new())))
        .route("/config.yaml", get(example_config))
        .route("/authenticate", post(authenticate))
        .route("/image/:encrypted_filename_hex", get(converted_image))
        .layer(TraceLayer::new_for_http()
            .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
            .on_response(trace::DefaultOnResponse::new().level(Level::INFO))
        )
        .layer(session_layer)
        .with_state(ApplicationState {
            config,
            engine: Engine::from(jinja),
            image_cache,
            image_dir,
            image_list_cache,
            secret_key,
            cpu_task_limit: Arc::new(Semaphore::new(4)),
        });

    let addr = SocketAddr::from(([127, 0, 0, 1], port));
    tracing::info!("listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn update_config_and_image_list_cache_job(
    image_dir: ImageDir,
    config: Config,
    image_metadata_cache: ImageListCache,
    sessions: MemoryStore,
) {
    let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(10));
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    let mut last_images = HashSet::new();
    let mut last_image_dir_change = DateTime::<Utc>::UNIX_EPOCH;
    let mut last_forced_update = Instant::now();

    loop {
        interval.tick().await;

        // Read config
        let config_file = image_dir.join("config.yaml");
        let config = config.clone();
        let invalidate_sessions = if config_file.exists() {
            task::spawn_blocking(move || {
                match config::read_from_file(&config_file) {
                    Ok(new_config) => {
                        let mut config = config.write().unwrap();
                        if new_config != *config {
                            tracing::debug!("New config: {:?}", new_config);
                            *config = new_config;
                            true
                        } else {
                            false
                        }
                    },
                    Err(error) => {
                        tracing::error!("Could not read config, using default: {:#}", error);
                        *config.write().unwrap() = config::Config::default();
                        true
                    }
                }
            }).await.unwrap()
        } else {
            *config.write().unwrap() = config::Config::default();
            true
        };
        if invalidate_sessions {
            sessions.clear_store().await.ok();
        }

        // Read the modification time of the image dir
        // Note that some fuse implementations (e.g. seaf_fuse) will update
        // the mtime of the directory when a file inside the directory is
        // modified. It might also be very expensive to stat files in such
        // a case, so we only record the modification time of the directory.
        let image_dir_change = match image_dir.metadata().and_then(|m| m.modified()) {
            Ok(modified) => modified.into(),
            Err(error) => {
                tracing::error!("Could not read modification time of image dir: {:#}", error);
                continue
            },
        };

        // Read all image files from the image_dir
        let images: HashSet<_> = match image_dir.read_dir() {
            Ok(images) => {
                // flatten ignores io error while traversing the iterator
                images.flatten().filter_map(|entry| {
                    let path = entry.path();
                    if path.extension().unwrap_or_default().to_ascii_lowercase() == "jpg" {
                        Some(path)
                    } else {
                        None
                    }
                }).collect()
            },
            Err(error) => {
                tracing::error!("Could not read images: {:#}", error);
                continue
            },
        };

        let update_image_metadata = if image_dir_change != last_image_dir_change {
            tracing::debug!("Update image list because image dir was modified");
            true
        } else if images != last_images {
            // TODO: Maybe clear removed files from the image cache here?
            tracing::debug!("Update image list because list of images changed");
            true
        } else if last_forced_update.elapsed() > Duration::from_secs(60*60) {
            last_forced_update = Instant::now();
            tracing::debug!("Update image list because one hour elapsed");
            true
        } else {
            false
        };

        // Update image list
        if update_image_metadata {
            let images = images.clone();
            let image_metadata = task::spawn_blocking(move || {
                read_image_metadata(&images)
            }).await.unwrap_or_else(|error| {
                tracing::error!("Could not read images due to panic: {:#}", error);
                Vec::new()
            });

            tracing::debug!("{} images in the image list cache", image_metadata.len());
            *image_metadata_cache.write().unwrap() = image_metadata;
        }

        last_images = images;
        last_image_dir_change = image_dir_change;

    }
}

#[derive(Debug, Serialize, Default, Clone)]
pub struct ImageInfo {
    width: u32,
    height: u32,
    created: DateTime<Utc>,
    #[serde(with = "hex")]
    encrypted_name: Vec<u8>,
    #[serde(skip_serializing)]
    name: OsString,
}

#[derive(Debug, Serialize)]
pub struct LoginTemplate {
    title: String,
}

#[derive(Debug, Serialize)]
pub struct IndexTemplate {
    title: String,
    top_message: Option<String>,
    images: Vec<ImageInfo>,
}

async fn index(
    engine: TemplateEngine,
    State(config): State<Config>,
    State(image_list_cache): State<ImageListCache>,
    State(secret_key): State<SecretKey>,
    session: ReadableSession,
) -> Result<Response, AppError> {
    let logged_in = session.get::<()>("logged_in").is_some();

    let config = config.read().unwrap();

    if logged_in {
        let images = image_list_cache.read().unwrap();

        // Encrypt image names
        let chacha_key = chacha20poly1305::Key::from_slice(&secret_key[0..32]);
        let chacha = XChaCha20Poly1305::new(chacha_key);
        let mut images = images.iter().cloned().map(|mut image_info| {
            let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);
            let mut ciphertext = chacha.encrypt(&nonce, image_info.name.as_bytes())?;

            image_info.encrypted_name = nonce.as_slice().to_vec();
            image_info.encrypted_name.append(&mut ciphertext);

            Ok(image_info)
        }).collect::<Result<Vec<_>>>()?;

        images.sort_by_key(|entry| Reverse(entry.created));

        Ok(RenderHtml("index", engine, IndexTemplate {
            title: config.title.clone(),
            top_message: config.top_message.clone(),
            images,
        }).into_response())
    } else {
        Ok(RenderHtml("login", engine, LoginTemplate {
            title: config.title.clone(),
        }).into_response())
    }
}

async fn example_config() -> &'static str {
    include_str!("../config.yaml.example")
}

#[derive(Deserialize)]
struct AuthenticateForm {
    password: String,
}

async fn authenticate(
    mut session: WritableSession,
    State(config): State<Config>,
    Form(form): Form<AuthenticateForm>,
) -> Redirect {
    if config.read().unwrap().passwords.contains(&form.password) {
        session.insert("logged_in", ()).ok();
    }
    Redirect::to("/")
}

async fn converted_image(
    extract::Path(encrypted_filename_hex): extract::Path<String>,
    State(image_cache): State<ImageCache>,
    State(image_dir): State<ImageDir>,
    State(secret_key): State<SecretKey>,
    session: ReadableSession,
    State(cpu_task_limit): State<CpuTaskLimit>,
) -> Result<impl IntoResponse, AppError> {
    session.get::<()>("logged_in")
        .ok_or(anyhow!("Trying to load image while not logged in!"))
        .context(StatusCode::FORBIDDEN)?;

    // Decrypt image name
    let image_name = {
        let encrypted_filename = hex::decode(&encrypted_filename_hex)?;
        let nonce = XNonce::from_slice(&encrypted_filename[0..24]);
        let ciphertext = &encrypted_filename[24..];
        let chacha_key = chacha20poly1305::Key::from_slice(&secret_key[0..32]);
        let chacha = XChaCha20Poly1305::new(chacha_key);
        chacha.decrypt(nonce, ciphertext)
    }.with_context(|| format!("Could not decrypt filename {}", encrypted_filename_hex))?;
    let image_name = OsStr::from_bytes(&image_name);

    let image_path = image_dir.join(image_name);

    image_path.exists()
        .then_some(())
        .ok_or(anyhow!("Requested image not found!"))
        .context(StatusCode::NOT_FOUND)?;

    // Check if we have the file already in cache
    let image_modified: Option<DateTime<Utc>> = std::fs::metadata(&image_path)
            .and_then(|m| m.modified())
            .map(DateTime::from)
            .ok();
    let cached_buffer = if let Some(image_modified) = image_modified {
        image_cache
            .read().unwrap()
            .get(image_name)
            .filter(|(cache_modified, _)| *cache_modified == image_modified)
            .map(|(_, image_buffer)| image_buffer)
            .cloned()
    } else { None };

    let image_buffer = match cached_buffer {
        Some(image_buffer) =>  {
            tracing::debug!("Read {:?} from cache", image_name);
            image_buffer
        }
        None => {
            let _cpu_task_permit = cpu_task_limit.clone().acquire_owned().await?;
            let image_buffer = task::spawn_blocking(move || {
                convert_image(&image_path)
                    .with_context(|| format!("Could not convert image {:?}", image_path))
            }).await??;

            if let Some(image_modified) = image_modified {
                tracing::debug!("Add {:?} ({}k) to cache", image_name, image_buffer.len() / 1024);
                image_cache
                    .write().unwrap()
                    .insert(image_name.to_owned(), (image_modified, image_buffer.clone()));
            }

            image_buffer
        }
    };

    Ok((
        [
            (header::CONTENT_TYPE, "image/webp"),
        ],
        image_buffer
    ))
}

fn convert_image(image_path: &Path) -> Result<Vec<u8>> {
    let file = File::open(image_path)?;
    let mut file = BufReader::new(file);

    // Read exif data first, as it only reads a tiny amount of data
    let exif = read_exif_data(&mut file);

    file.rewind()?;
    let mut image = image::io::Reader::new(&mut file)
        .with_guessed_format()?
        .decode()?;

    // Fix image orientation
    if let Ok(exif) = exif {
        image = fix_image_orientation(image, &exif);
    }

    let small_image = image.resize(600, 600, FilterType::Lanczos3);

    let mut image_buffer = Cursor::new(Vec::new());
    small_image.write_to(&mut image_buffer, image::ImageOutputFormat::WebP)?;

    Ok(image_buffer.into_inner())
}

fn read_exif_data(mut file: impl BufRead+Seek) -> Result<Exif> {
    file.rewind()?;
    let exifreader = exif::Reader::new();
    Ok(exifreader.read_from_container(&mut file)?)
}

fn read_image_size(mut file: impl BufRead+Seek, exif: Option<&Exif>) -> Result<(u32, u32)> {
    file.rewind()?;

    let (mut width, mut height) = image::io::Reader::new(&mut file)
        .with_guessed_format()?
        .into_dimensions()?;

    if let Some(exif) = exif {
        if let Some(orientation) = exif.get_field(exif::Tag::Orientation, exif::In::PRIMARY) {
            if let Some(5 | 6 | 7 | 8) = orientation.value.get_uint(0) {
                std::mem::swap(&mut width, &mut height);
            };
        };
    };

    Ok((width, height))
}

// How many degrees to rotate clockwise, does not support flipped images
fn fix_image_orientation(image: DynamicImage, exif: &Exif) -> DynamicImage {
    match exif.get_field(exif::Tag::Orientation, exif::In::PRIMARY) {
        Some(orientation) =>
            match orientation.value.get_uint(0) {
                Some(1) => image,
                Some(2) => image.fliph(),
                Some(3) => image.rotate180(),
                Some(4) => image.rotate180().fliph(),
                Some(5) => image.rotate90().fliph(),
                Some(6) => image.rotate90(),
                Some(7) => image.rotate270().fliph(),
                Some(8) => image.rotate270(),
                _ => image,
            },
        None => image,
    }
}

fn read_image_metadata(images: &HashSet<PathBuf>) -> Vec<ImageInfo> {
    let mut files = vec![];

    for path in images {
        let image_info = match read_image_info(path) {
            Ok(image_info) => image_info,
            Err(error) => {
                tracing::warn!("Skipping {:?} due to error: {:#}", path, error);
                continue;
            }
        };
        files.push(image_info);
    }

    files
}

fn extract_exif_string(field: &exif::Field) -> Option<String> {
    match field.value {
        exif::Value::Ascii(ref value) if !value.is_empty() => {
            String::from_utf8(value[0].clone()).ok()
        }
        _ => None,
    }
}

fn read_image_info(path: &Path) -> Result<ImageInfo> {
    let file = File::open(path)?;
    let mut file = BufReader::new(file);

    let exif = read_exif_data(&mut file).ok();

    // Check if we need to flip the coordinates
    let (width, height) = read_image_size(&mut file, exif.as_ref())
        .context("Could not read image size")?;

    let datetime: Option<DateTime<Utc>> = exif.and_then(|exif| {
        // First, try to read creation date from EXIF data
        let datetime_without_timezone = exif
            .get_field(exif::Tag::DateTimeOriginal, exif::In::PRIMARY)
            .and_then(extract_exif_string);
        let timezone = exif
            .get_field(exif::Tag::OffsetTimeOriginal, exif::In::PRIMARY)
            .and_then(extract_exif_string);

        match (datetime_without_timezone, timezone) {
            (Some(datetime), Some(timezone)) => Some(datetime + &timezone),
            (Some(datetime), None) => Some(datetime + "+00:00"),
            _ => None,
        }.and_then(|datetime| {
            DateTime::parse_from_str(&datetime, "%Y:%m:%d %T%:z")
                .ok().map(|d| d.with_timezone(&Utc))
        })
    }).or_else(|| {
        // If that doesn't work, fall back to the file modification time
        std::fs::metadata(path)
            .and_then(|m| m.modified())
            .map(DateTime::from)
            .ok()
    });

    if datetime.is_none() {
        tracing::warn!("Could not determine original datetime for {:?}", path);
    }

    Ok(ImageInfo {
        width,
        height,
        created: datetime.unwrap_or_default(),
        name: path.file_name().expect("invalid file path").to_owned(),
        ..Default::default()
    })
}