mod icon;

use std::{collections::HashMap, fs, io::ErrorKind, path::PathBuf, str::FromStr};

use icon::extract_icon;
use mlua::{Lua, Table, Value as LuaValue};
use once_cell::sync::{Lazy, OnceCell};

use axum::{
    extract::Request,
    handler::HandlerWithoutStateExt,
    http::{HeaderName, HeaderValue, StatusCode},
    response::{Html, IntoResponse},
    routing::get,
};
use base64::prelude::*;
use serde::Serialize;
use tera::{to_value, Context, Result as TeraResult, Tera, Value};
use time::OffsetDateTime;
use tokio::time::Instant;
use tokio_stream::{wrappers::ReadDirStream, StreamExt};
use tower_http::services::ServeDir;

pub type AError = Box<dyn std::error::Error>;
pub type AResult<T> = Result<T, AError>;

pub static TERA: Lazy<Tera> = Lazy::new(|| {
    let mut tera = Tera::default();
    tera.add_raw_template(
        "index.html.jinja",
        include_str!("../templates/index.html.jinja"),
    )
    .unwrap();
    tera.register_filter("iconize", iconize);
    tera.register_filter("from_ico", from_ico);
    tera.register_filter("size", size);
    tera.register_filter("time", time);
    tera.register_filter("md", md);

    tera
});

pub static ICONS: Lazy<HashMap<&str, String>> = Lazy::new(|| {
    let mut map = HashMap::new();
    map.insert(
        "unknown",
        BASE64_STANDARD.encode(include_bytes!("../icons/unknown.png")),
    );
    map.insert(
        "back",
        BASE64_STANDARD.encode(include_bytes!("../icons/back.png")),
    );
    map.insert(
        "dir",
        BASE64_STANDARD.encode(include_bytes!("../icons/dir.png")),
    );
    map.insert(
        "binary",
        BASE64_STANDARD.encode(include_bytes!("../icons/binary.png")),
    );
    map.insert(
        "image",
        BASE64_STANDARD.encode(include_bytes!("../icons/image2.png")),
    );
    map.insert(
        "compressed",
        BASE64_STANDARD.encode(include_bytes!("../icons/compressed.png")),
    );

    map
});

fn from_ico(value: &Value, _args: &HashMap<String, Value>) -> TeraResult<Value> {
    let Some(icon) = value.as_array() else {
        return iconize(
            &Value::String(String::from("")),
            &std::default::Default::default(),
        );
    };

    let icon = icon
        .iter()
        .map(|n| n.as_u64().unwrap() as u8)
        .collect::<Vec<_>>();

    let data = BASE64_STANDARD.encode(icon);

    let str = format!("<img src=\"data:image/png;base64,{}\" alt=\"\"/>", data);

    Ok(to_value(str).unwrap())
}

fn iconize(value: &Value, _args: &HashMap<String, Value>) -> TeraResult<Value> {
    let icon = value.as_str().unwrap_or("unknown");

    let icon = match icon {
        "DIR" | "BACK" => "dir",
        "image/png" | "image/jpg" | "image/webp" => "image",
        "application/zip" => "compressed",
        "application/octet-stream" => "binary",
        _ => "unknown",
    };

    let icons = &*ICONS;
    let icon = if icons.contains_key(icon) {
        &icons[icon]
    } else {
        println!("Unknown filetype: {}", icon);
        &icons["unknown"]
    };

    let str = format!("<img src=\"data:image/png;base64,{}\" alt=\"\" />", icon);

    Ok(to_value(str).unwrap())
}

fn size(value: &Value, _args: &HashMap<String, Value>) -> TeraResult<Value> {
    let Some(value) = value.as_f64() else {
        return Ok(value.clone());
    };

    Ok(to_value(human_bytes::human_bytes(value)).unwrap())
}

fn time(value: &Value, _args: &HashMap<String, Value>) -> TeraResult<Value> {
    let time_format: std::vec::Vec<time::format_description::FormatItem<'_>> =
        time::format_description::parse("[day]-[month]-[year] [hour]:[minute]:[second]").unwrap();

    let Some(time) = value.as_array() else {
        return Ok(value.clone());
    };

    let time: Vec<i64> = time.iter().map(|v| v.as_i64().unwrap()).collect();

    let new_time: OffsetDateTime = OffsetDateTime::now_utc();
    let new_time = new_time.replace_year(time[0] as i32).unwrap();
    let new_time = new_time.replace_ordinal(time[1] as u16).unwrap();
    let new_time = new_time.replace_hour(time[2] as u8).unwrap();
    let new_time = new_time.replace_minute(time[3] as u8).unwrap();
    let new_time = new_time.replace_second(time[4] as u8).unwrap();
    let new_time = new_time.replace_nanosecond(time[5] as u32).unwrap();

    Ok(to_value(new_time.format(&time_format).unwrap()).unwrap())
}

fn md(value: &Value, _args: &HashMap<String, Value>) -> TeraResult<Value> {
    // parse value as markdown and return html
    if let Some(value) = value.as_str() {
        if let Ok(md) = markdown::to_html_with_options(value, &markdown::Options::gfm()) {
            return Ok(to_value(md).unwrap());
        }
    }

    Ok(to_value(value).unwrap())
}

pub static BASE_DIR: OnceCell<PathBuf> = OnceCell::new();

#[derive(Debug, Serialize)]
struct FileInfo {
    filename: String,
    size: usize,
    format: String,
    icon: Option<Vec<u8>>,
    // created: OffsetDateTime,
    modified: OffsetDateTime,
    description: Option<String>,
}

#[tokio::main]
async fn main() -> AResult<()> {
    let listener = tokio::net::TcpListener::bind("[::]:3000").await.unwrap();
    Ok(axum::serve(listener, file_handler.into_make_service()).await?)
}

async fn file_handler(request: Request) -> impl IntoResponse {
    let base_dir = BASE_DIR.get_or_init(|| {
        PathBuf::from_str(&std::env::var("BASE_DIR").unwrap_or(String::from("./public"))).unwrap()
    });

    if !request.uri().path().ends_with(".lua") {
        return ServeDir::new(base_dir)
            .fallback(get(handler))
            .try_call(request)
            .await
            .map(|e| e.into_response())
            .map_err(|_| StatusCode::NOT_FOUND);
    };

    let lua = Lua::new();
    let globals = lua.globals();

    use mlua::{ExternalResult, LuaSerdeExt};

    let fetch_json = lua
        .create_function(|lua, uri: String| {
            let resp = reqwest::blocking::get(&uri)
                .and_then(|resp| resp.error_for_status())
                .into_lua_err()?;
            let json = resp.json::<serde_json::Value>().into_lua_err()?;
            lua.to_value(&json)
        })
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let fetch_b64 = lua
        .create_function(|lua, uri: String| {
            let resp = reqwest::blocking::get(&uri)
                .and_then(|resp| resp.error_for_status())
                .into_lua_err()?;
            let bytes = resp.bytes().into_lua_err()?;
            let b64 = BASE64_STANDARD.encode(bytes);
            lua.to_value(&b64)
        })
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let dbg = lua
        .create_function(|_, value: LuaValue| {
            dbg!(value);
            Ok(())
        })
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    globals.set("fetch", fetch_json).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    globals.set("fetch_b64", fetch_b64).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    globals.set("dbg", dbg).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let request_table = lua.create_table().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let headers_table = lua
        .create_table_from(
            request
                .headers()
                .into_iter()
                .map(|(k, v)| (k.to_string(), v.to_str().unwrap())),
        )
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    request_table.set("headers", headers_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    request_table
        .set("uri", request.uri().to_string())
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    request_table
        .set("method", request.method().to_string())
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    // Inject functions to change headers and such
    globals.set("request", request_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let response_table = lua.create_table().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    response_table
        .set(
            "headers",
            lua.create_table().map_err(|e| {
                eprintln!("Lua Error: {:?}", e);
                StatusCode::INTERNAL_SERVER_ERROR
            })?,
        )
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    globals.set("response", response_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let mut path = base_dir.clone();

    let uri = urlencoding::decode(&request.uri().path()[1..]).map_err(|e| {
        println!("{:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    path.push(&*uri);
    let full_path = fs::canonicalize(&path).map_err(|e| {
        println!("{:?}", e);
        StatusCode::NOT_FOUND
    })?;

    let full_base_path = fs::canonicalize(base_dir).map_err(|e| {
        println!("{:?}", e);
        StatusCode::NOT_FOUND
    })?;

    if !full_path.starts_with(full_base_path) {
        return Err(StatusCode::BAD_REQUEST);
    }

    let script = fs::read_to_string(full_path).map_err(|e| {
        eprintln!("Lua Read Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let script = lua.load(script).eval::<LuaValue>().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let result = lua.load("return response").eval::<Table>().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let script = if let LuaValue::String(script) = script {
        script.to_string_lossy().to_string()
    } else {
        serde_json::to_string(&script).map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?
    };

    let mut response = Html(script).into_response();

    if let Ok(headers) = result.get::<&str, Table>("headers") {
        let pairs = headers.pairs::<String, String>();

        for pair in pairs {
            let (k, v) = pair.map_err(|e| {
                eprintln!("Lua Error: {:?}", e);
                StatusCode::INTERNAL_SERVER_ERROR
            })?;
            response.headers_mut().insert(
                HeaderName::from_str(&k).map_err(|e| {
                    eprintln!("Lua Error: {:?}", e);
                    StatusCode::INTERNAL_SERVER_ERROR
                })?,
                HeaderValue::from_str(&v).map_err(|e| {
                    eprintln!("Lua Error: {:?}", e);
                    StatusCode::INTERNAL_SERVER_ERROR
                })?,
            );
        }
    };

    Ok(response)
}

async fn handler(request: Request) -> Result<Html<String>, StatusCode> {
    let start = Instant::now();
    let base_dir = BASE_DIR.get().unwrap();
    let mut path = base_dir.clone();

    let uri = urlencoding::decode(&request.uri().path()[1..]).map_err(|e| {
        println!("{:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    path.push(&*uri);
    let full_path = fs::canonicalize(&path).map_err(|e| {
        println!("{:?}", e);
        StatusCode::NOT_FOUND
    })?;

    let full_base_path = fs::canonicalize(base_dir).map_err(|e| {
        println!("{:?}", e);
        StatusCode::NOT_FOUND
    })?;

    if !full_path.starts_with(full_base_path) {
        return Err(StatusCode::BAD_REQUEST);
    }

    let body = match render_dir(&path, start).await {
        Ok(body) => Ok(body),
        Err(e) => {
            let ioerror: Option<&std::io::Error> = e.downcast_ref();
            if let Some(ioerror) = ioerror {
                if ioerror.kind() == ErrorKind::NotFound {
                    Err(StatusCode::NOT_FOUND)
                } else {
                    println!("{:?}", ioerror);
                    Err(StatusCode::INTERNAL_SERVER_ERROR)
                }
            } else {
                println!("{:?}", e);
                Err(StatusCode::INTERNAL_SERVER_ERROR)
            }
        }
    };

    Ok(Html(body?))
}

async fn render_dir(path: &PathBuf, time: Instant) -> AResult<String> {
    let base_dir = BASE_DIR.get().unwrap();
    let time2 = Instant::now();
    let base = base_dir.display().to_string();
    let dirname = path
        .display()
        .to_string()
        .replace(&base, "")
        .replace('\\', "/");

    let filename = path.file_name().unwrap().to_string_lossy().to_string();
    let description = PathBuf::from_str(&format!(
        "{}/.{}.info",
        path.parent().unwrap().to_string_lossy(),
        filename
    ))?;
    let description = if description.exists() {
        Some(fs::read_to_string(description)?)
    } else {
        None
    };

    let directories = get_directories(path).await?;
    let files = get_files(path).await?;

    let tera = &*TERA;

    let mut context = Context::new();
    context.insert("description", &description);
    context.insert("dirname", &dirname);
    context.insert("directories", &directories);
    context.insert("files", &files);

    let loading = time2.elapsed() + time.elapsed();
    context.insert("loading", &loading.as_millis());
    Ok(tera.render("index.html.jinja", &context)?)
}

async fn get_directories(path: &PathBuf) -> AResult<Vec<FileInfo>> {
    let mut contents = ReadDirStream::new(tokio::fs::read_dir(&path).await?);

    let mut files = vec![];

    while let Some(file) = contents.next().await {
        let file = file?;
        let metadata = file.metadata().await?;
        if !metadata.is_dir() {
            continue;
        }
        let filename = file.file_name().to_string_lossy().to_string();
        if filename.starts_with('.') {
            continue;
        }

        // let created: OffsetDateTime = metadata.created()?.into();
        let modified: OffsetDateTime = metadata.modified()?.into();

        let description = PathBuf::from_str(&format!(
            "{}/.{}.info",
            file.path().parent().unwrap().to_string_lossy(),
            filename
        ))?;

        let description = if description.exists() {
            Some(fs::read_to_string(description)?)
        } else {
            None
        };

        files.push(FileInfo {
            filename,
            format: "DIR".to_string(),
            size: 0,
            icon: None,
            // created,
            modified,
            description,
        });
    }

    Ok(files)
}

async fn get_files(path: &PathBuf) -> AResult<Vec<FileInfo>> {
    let mut contents = ReadDirStream::new(tokio::fs::read_dir(&path).await?);

    let mut files = vec![];

    while let Some(file) = contents.next().await {
        let file = file?;
        let metadata = file.metadata().await?;
        if metadata.is_dir() {
            continue;
        }
        let filename = file.file_name().to_string_lossy().to_string();
        if filename.starts_with('.') {
            continue;
        }
        let format = file_format::FileFormat::from_file(file.path())?;
        let format = format.media_type();

        let icon = if format == "application/x-dosexec" {
            // ..
            extract_icon(&file.path()).ok().flatten()
        } else {
            None
        };

        // let created: OffsetDateTime = metadata.created()?.into();
        let modified: OffsetDateTime = metadata.modified()?.into();

        let description = PathBuf::from_str(&format!(
            "{}/.{}.info",
            file.path().parent().unwrap().to_string_lossy(),
            filename
        ))?;

        let description = if description.exists() {
            Some(fs::read_to_string(description)?)
        } else {
            None
        };

        files.push(FileInfo {
            filename,
            format: format.to_string(),
            size: metadata.len() as usize,
            icon,
            // created,
            modified,
            description,
        });
    }

    Ok(files)
}