mod api;
mod config;
mod model;
mod state;
mod storage_manager;

use crate::config::PicOxConfig;
use crate::model::album::Album;
use crate::model::api_key::{ApiKey, ApiPermissions};
use crate::model::image::{Image, ImageData};
use crate::state::Context;
use crate::storage_manager::{StorageManager, StoreError};
use axum::body::Bytes;
use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, Request, State};
use axum::http::{HeaderMap, StatusCode};
use axum::middleware::Next;
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{middleware, Json, Router};
use axum_macros::FromRequest;
use base64::Engine;
use j_db::database::Database;
use j_db::model::JdbModel;
use log::info;
use rand::seq::SliceRandom;
use rand::thread_rng;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use structopt::StructOpt;
use tokio::sync::RwLock;

type PicContext = Arc<Context>;

#[derive(StructOpt, Debug, Clone)]
#[structopt(about = "PicOx Commands")]
enum SubCommands {
    /// Start the PicOx server
    Start,
    /// Create a Picox API Key
    CreateKey {
        /// Key description
        description: String,
        /// API Key permissions (WRITE, DELETE)
        permissions: ApiPermissions,
    },
    /// Dump the database state to json
    Dump {
        /// Output path
        out_path: PathBuf,
    },
    /// Import the database state from json
    Import {
        /// Path to json containing DB state
        db_file: PathBuf,
    },
}

#[derive(Debug, Clone, StructOpt)]
struct Args {
    /// Path to the config file
    #[structopt(short, long, env = "PICOX_CONFIG")]
    pub config: PathBuf,
    #[structopt(subcommand)]
    pub command: SubCommands,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CreateAlbum {
    pub album_name: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AlbumQuery {
    pub album_name: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AddImage {
    pub album: AlbumQuery,
    pub tags: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub enum ImageSort {
    Random,
    DateAscending,
    DateDescending,
    #[default]
    None,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ImageQuery {
    pub album: Option<String>,
    #[serde(default)]
    pub tags: Vec<String>,
    #[serde(default)]
    pub order: ImageSort,
    #[serde(default)]
    pub limit: usize,
}

#[derive(FromRequest)]
#[from_request(via(axum::Json), rejection(PicOxError))]
struct Response<T>(T);

impl<T> IntoResponse for Response<T>
where
    Json<T>: IntoResponse,
{
    fn into_response(self) -> axum::response::Response {
        Json(self.0).into_response()
    }
}

#[allow(dead_code)]
pub enum PicOxError {
    StoreError(StoreError),
    DbError(j_db::error::JDbError),
    AlbumNotFound,
    ImageNotFound,
    TokenInvalid,
    NoUserInHeader,
}

impl From<StoreError> for PicOxError {
    fn from(value: StoreError) -> Self {
        Self::StoreError(value)
    }
}

impl From<j_db::error::JDbError> for PicOxError {
    fn from(value: j_db::error::JDbError) -> Self {
        Self::DbError(value)
    }
}

#[derive(Serialize)]
struct ErrorResponse {
    pub message: String,
}

impl IntoResponse for PicOxError {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            PicOxError::StoreError(err) => match err {
                StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()),
                StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()),
                StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()),
                StoreError::IOError(_) => (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "IO Error Has Occurred!".to_string(),
                ),
            },
            PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
            PicOxError::AlbumNotFound => (
                StatusCode::INTERNAL_SERVER_ERROR,
                "No album found!".to_string(),
            ),
            PicOxError::ImageNotFound => (
                StatusCode::INTERNAL_SERVER_ERROR,
                "Image not found".to_string(),
            ),
            PicOxError::TokenInvalid => (StatusCode::UNAUTHORIZED, "Token is invalid".to_string()),
            PicOxError::NoUserInHeader => (
                StatusCode::INTERNAL_SERVER_ERROR,
                "User not found in header".to_string(),
            ),
        };

        (status, Response(ErrorResponse { message })).into_response()
    }
}

async fn create_album(
    State(context): State<PicContext>,
    headers: HeaderMap,
    Json(album): Json<CreateAlbum>,
) -> Result<Response<Album>, PicOxError> {
    let user_id = get_user_id_from_headers(&headers)?;
    let new_album = Album::new(&album.album_name, Vec::new(), user_id);

    info!(
        "Creating new album '{}' for user {}",
        album.album_name, user_id
    );
    let new_album = context.db.insert(new_album)?;

    Ok(Response(new_album))
}

async fn get_album(
    album_id: Path<u64>,
    State(context): State<PicContext>,
) -> Result<Response<Album>, PicOxError> {
    let album = context.db.get::<Album>(*album_id)?;

    Ok(Response(album))
}

async fn query_images(
    image_query: Query<ImageQuery>,
    State(context): State<PicContext>,
) -> Result<Response<Vec<Image>>, PicOxError> {
    let album_id = if let Some(album) = &image_query.album {
        Some(
            Album::find_album_by_query(
                &context.db,
                AlbumQuery {
                    album_name: Some(album.to_string()),
                },
            )
            .ok_or(PicOxError::AlbumNotFound)?
            .id()
            .unwrap(),
        )
    } else {
        None
    };

    let mut images: Vec<Image> = context
        .db
        .filter(|_, img: &Image| {
            if let Some(album_id) = album_id {
                if img.album != album_id {
                    return false;
                }
            }

            if !image_query.tags.is_empty() {
                let mut found = false;
                for tag in &image_query.tags {
                    if img.tags.contains(tag) {
                        found = true;
                        break;
                    }
                }

                if !found {
                    return false;
                }
            }

            true
        })?
        .collect();

    match image_query.order {
        ImageSort::Random => {
            images.shuffle(&mut thread_rng());
        }
        ImageSort::DateAscending => {
            images.sort_by(|img_a, img_b| img_a.create_date.cmp(&img_b.create_date))
        }
        ImageSort::DateDescending => {
            images.sort_by(|img_a, img_b| img_b.create_date.cmp(&img_a.create_date))
        }
        ImageSort::None => {}
    }

    if images.len() > image_query.limit {
        images.drain(image_query.limit..);
    }

    Ok(Response(images))
}

fn get_user_id_from_headers(headers: &HeaderMap) -> Result<u64, PicOxError> {
    let user = headers.get("user").ok_or(PicOxError::NoUserInHeader)?;
    let user_str = user.to_str().unwrap();
    let user: u64 = user_str.parse().unwrap();

    Ok(user)
}

async fn add_image(
    State(context): State<PicContext>,
    headers: HeaderMap,
    mut img_data: Multipart,
) -> Result<Response<Image>, PicOxError> {
    let mut data: Vec<u8> = Vec::new();
    let mut metadata: Option<AddImage> = None;
    let mut file_name = None;
    while let Some(field) = img_data.next_field().await.unwrap() {
        let field_name = field.name();
        if let Some(field_name) = field_name {
            if field_name == "metadata" {
                let metadata_json = field.text().await.unwrap();
                metadata = Some(serde_json::from_str(&metadata_json).unwrap())
            } else if field_name == "img_data" {
                file_name = Some(field.file_name().unwrap().to_string());
                let file_segment = field.bytes().await.unwrap_or(Bytes::new());
                data.extend_from_slice(file_segment.as_ref());
            }
        }
    }

    let user = headers.get("user").unwrap().to_str().unwrap();
    let user: u64 = user.parse().unwrap();

    let mut album =
        Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap();

    let mut store_manager = context.store_manager.write().await;
    let img = store_manager
        .store_img(
            &context.db,
            None,
            ImageData::Bytes(data),
            &file_name.unwrap(),
            user,
            album.id().unwrap(),
        )
        .await?;

    album.images.push(img.id().unwrap());

    context.db.insert::<Album>(album)?;

    Ok(Response(img))
}

async fn query_album(
    album_query: Query<AlbumQuery>,
    State(context): State<PicContext>,
) -> Result<Response<Option<Album>>, PicOxError> {
    let resp = Album::find_album_by_query(&context.db, album_query.0);
    Ok(Response(resp))
}

async fn check_token_header(
    State(context): State<PicContext>,
    mut request: Request,
    next: Next,
) -> Result<impl IntoResponse, PicOxError> {
    let headers = request.headers();

    if let Some(token) = headers.get("token") {
        if let Some(api_key) = ApiKey::find_api_key_by_token(&context.db, token.to_str().unwrap())?
        {
            request
                .headers_mut()
                .insert("user", api_key.id().unwrap().into());
            return Ok(next.run(request).await);
        }
    }

    Err(PicOxError::TokenInvalid)
}

async fn run_picox(db: Database, config: PicOxConfig) {
    let store_manager = StorageManager::new(config.storage_config.clone());

    let context = Context {
        db,
        config: config.clone(),
        store_manager: RwLock::new(store_manager),
    };

    let context = Arc::new(context);

    let app = Router::new()
        .route("/api/image/", post(add_image))
        .layer(DefaultBodyLimit::max(1024 * 1024 * 1024))
        .route("/api/album/create", post(create_album))
        .layer(middleware::from_fn_with_state(
            context.clone(),
            check_token_header,
        ))
        .route("/api/album/:id", get(get_album))
        .route("/api/album/", get(query_album))
        .route("/api/image/", get(query_images))
        .with_state(context);

    let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap();
    info!("Serving at {}", config.host);
    axum::serve(listener, app).await.unwrap();
}

#[tokio::main]
async fn main() {
    let args = Args::from_args();

    // initialize tracing
    tracing_subscriber::fmt::init();

    let config = PicOxConfig::new(args.config);

    let db = Database::new(&config.db_path).unwrap();

    match args.command {
        SubCommands::Start => {
            run_picox(db, config).await;
        }
        SubCommands::CreateKey {
            description,
            permissions,
        } => {
            let (token, api_key) = ApiKey::create_new_key(&db, description, permissions).unwrap();

            info!("New Key info: {:?}", api_key);
            info!("Token: {}", base64::prelude::BASE64_STANDARD.encode(token));
            db.db.flush().unwrap();
        }
        SubCommands::Dump { out_path } => {
            info!("Dumping database state to {:?}", out_path);

            tokio::fs::write(out_path, db.dump_db().unwrap().pretty(4))
                .await
                .unwrap();
        }
        SubCommands::Import { db_file } => {
            info!("Importing database state from {:?}", db_file);
            let db_state = tokio::fs::read(db_file).await.unwrap();
            let db_state = String::from_utf8(db_state).unwrap();
            let db_state = json::parse(&db_state).unwrap();

            db.import_db(db_state).unwrap();
            db.db.flush().unwrap();
        }
    }
}