mod argument;
mod storage;

use anyhow::Result;
use argument::Argument;

use native_tls::{HandshakeError, Identity, TlsAcceptor, TlsStream};
use std::{
    fs::File,
    io::{Read, Write},
    net::{SocketAddr, TcpListener, TcpStream},
    os::unix::fs::FileExt,
    sync::Arc,
    thread,
    time::{SystemTime, UNIX_EPOCH},
};

fn main() -> Result<()> {
    use clap::Parser;
    let argument = Arc::new(Argument::parse());

    // https://geminiprotocol.net/docs/protocol-specification.gmi#the-use-of-tls
    let acceptor = TlsAcceptor::new(Identity::from_pkcs12(
        &{
            let mut buffer = vec![];
            File::open(&argument.identity)?.read_to_end(&mut buffer)?;
            buffer
        },
        &argument.password,
    )?)?;

    let listener = TcpListener::bind(&argument.bind)?;

    println!("[{}] [info] Server started on {}", now(), argument.bind);

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                thread::spawn({
                    let argument = argument.clone();
                    let peer = stream.peer_addr()?;
                    let connection = acceptor.accept(stream);
                    move || handle(argument, peer, connection)
                });
            }
            Err(e) => println!("[{}] [error] {e}", now()),
        }
    }
    Ok(())
}

fn handle(
    argument: Arc<Argument>,
    peer: SocketAddr,
    connection: Result<TlsStream<TcpStream>, HandshakeError<TcpStream>>,
) {
    use titanite::*;
    println!("[{}] [info] [{peer}] New peer connected..", now());
    match connection {
        Ok(mut stream) => {
            // server should work with large files without memory overload,
            // because of that incoming data read partially, using chunks;
            // collect header bytes first to route the request
            let mut header = Vec::with_capacity(HEADER_MAX_LEN);
            loop {
                let mut buffer = vec![0];
                match stream.read(&mut buffer) {
                    Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection.", now()),
                    Ok(_) => {
                        header.push(buffer[0]);
                        // header bytes collected
                        if header.len() > HEADER_MAX_LEN || buffer[0] == b'\n' {
                            // detect controller for the request by parse its header bytes
                            return match Request::from_bytes(&header) {
                                Ok(request) => match request {
                                    Request::Gemini(this) => {
                                        gemini(this, &argument, &peer, &mut stream)
                                    }
                                    Request::Titan(this) => {
                                        titan(this, &argument, &peer, &mut stream)
                                    }
                                },
                                Err(e) => send(
                                    &response::failure::temporary::General {
                                        message: Some("Internal server error".to_string()),
                                    }
                                    .into_bytes(),
                                    &mut stream,
                                    |result| match result {
                                        Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                    },
                                ),
                            };
                        }
                    }
                    Err(e) => {
                        return send(
                            &response::failure::temporary::General {
                                message: Some("Internal server error".to_string()),
                            }
                            .into_bytes(),
                            &mut stream,
                            |result| match result {
                                Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                            },
                        )
                    }
                }
            }
        }
        Err(e) => println!("[{}] [warning] [{peer}] Handshake issue: {e}", now()),
    }
}

fn gemini(
    gemini: titanite::request::Gemini,
    argument: &Argument,
    peer: &SocketAddr,
    stream: &mut TlsStream<TcpStream>,
) {
    use titanite::*;
    println!("[{}] [info] [{peer}] Request: {}", now(), gemini.url);
    // file could be large,
    // to not overflow the memory pool, build the response using chunks
    match storage::Item::from_url(gemini.url.as_str(), &argument.directory) {
        Ok(item) => {
            let mut read: usize = 0;
            // create header packet
            match stream
                .write_all(&response::success::default::Meta { mime: item.mime }.into_bytes())
            {
                // chunk begin
                Ok(()) => loop {
                    let mut data = vec![0; argument.chunk];
                    match item.file.read_at(&mut data, read as u64) {
                        Ok(l) => match stream.write_all(&data[..l]) {
                            Ok(()) => {
                                // EOF
                                if l == 0 {
                                    println!("[{}] [info] [{peer}] Response: {read} bytes", now());
                                    match close(stream) {
                                        Ok(()) => println!(
                                            "[{}] [info] [{peer}] Connection closed by server.",
                                            now()
                                        ),
                                        Err(e) => println!("[{}] [warning] [{peer}] {e}", now()),
                                    }
                                    break;
                                }
                                read += l;
                            }
                            Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                        },
                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                    }
                },
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            }
        }
        Err(e) => send(
            &response::failure::permanent::NotFound {
                message: Some("Not found".to_string()),
            }
            .into_bytes(),
            stream,
            |result| match result {
                Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            },
        ),
    }
}

fn titan(
    titan: titanite::request::Titan,
    argument: &Argument,
    peer: &SocketAddr,
    stream: &mut TlsStream<TcpStream>,
) {
    use titanite::*;
    println!("[{}] [info] [{peer}] Request: {}", now(), titan.meta.url);
    // require content type for application,
    // even MIME value is optional by Titan specification
    let mime = match titan.meta.mime {
        Some(mime) => mime,
        None => {
            const MESSAGE: &str = "Content type is required";
            return send(
                &response::failure::permanent::BadRequest {
                    message: Some(MESSAGE.to_string()),
                }
                .into_bytes(),
                stream,
                |result| match result {
                    Ok(()) => println!("[{}] [warning] [{peer}] {MESSAGE}", now()),
                    Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                },
            );
        }
    };
    // validate total bytes
    let mut total = 0;
    // create new destination file
    match storage::Item::create(&argument.directory, mime) {
        Ok(mut tmp) => loop {
            let mut input = vec![0; argument.chunk];
            match stream.read(&mut input) {
                Ok(0) => {
                    println!("[{}] [warning] [{peer}] Peer closed connection.", now());
                    if let Err(e) = tmp.delete() {
                        println!("[{}] [error] [{peer}] {e}", now());
                    }
                    break;
                }
                Ok(read) => {
                    total += read;
                    // validate server-side limits
                    if argument.size.is_some_and(|limit| total > limit) {
                        if let Err(e) = tmp.delete() {
                            println!("[{}] [error] [{peer}] {e}", now());
                        }
                        const MESSAGE: &str = "Allowed max length limit reached";
                        return send(
                            &response::failure::permanent::BadRequest {
                                message: Some(MESSAGE.to_string()),
                            }
                            .into_bytes(),
                            stream,
                            |result| match result {
                                Ok(()) => println!("[{}] [warning] [{peer}] {MESSAGE}", now()),
                                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                            },
                        );
                    }

                    // validate client-side limits (from header)
                    if total > titan.meta.size {
                        if let Err(e) = tmp.delete() {
                            println!("[{}] [error] [{peer}] {e}", now());
                        }
                        const MESSAGE: &str = "Data size mismatch header declaration";
                        return send(
                            &response::failure::permanent::BadRequest {
                                message: Some(MESSAGE.to_string()),
                            }
                            .into_bytes(),
                            stream,
                            |result| match result {
                                Ok(()) => println!("[{}] [warning] [{peer}] {MESSAGE}", now()),
                                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                            },
                        );
                    }
                    // begin chunk recording into the temporary file
                    match tmp.file.write(&input[..read]) {
                        Ok(write) => {
                            // validate file bytes recorded match stream bytes received
                            if write != read {
                                if let Err(e) = tmp.delete() {
                                    println!("[{}] [error] [{peer}] {e}", now());
                                }
                                const MESSAGE: &str = "File size mismatch";
                                return send(
                                    &response::failure::temporary::General {
                                        message: Some(MESSAGE.to_string()),
                                    }
                                    .into_bytes(),
                                    stream,
                                    |result| match result {
                                        Ok(()) => {
                                            println!("[{}] [error] [{peer}] {MESSAGE}", now())
                                        }
                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                    },
                                );
                            }

                            // just to make sure
                            if total > titan.meta.size {
                                panic!()
                            }

                            // all data received
                            if titan.meta.size == total {
                                return match tmp.commit() {
                                    Ok(pmt) => send(
                                        &response::redirect::Permanent {
                                            target: match argument.redirect {
                                                Some(ref target) => format!(
                                                    "{}/{}",
                                                    target.trim_end_matches("/"),
                                                    pmt.to_uri(&argument.directory)
                                                ),
                                                None => format!(
                                                    "gemini://{}/{}",
                                                    argument.bind,
                                                    pmt.to_uri(&argument.directory)
                                                ),
                                            },
                                        }
                                        .into_bytes(),
                                        stream,
                                        |result| match result {
                                            Ok(()) => println!(
                                                "[{}] [info] [{peer}] Save to {}",
                                                now(),
                                                pmt.path.to_string_lossy()
                                            ),
                                            Err(e) => {
                                                println!("[{}] [error] [{peer}] {e}", now());
                                                if let Err(e) = pmt.delete() {
                                                    println!("[{}] [error] [{peer}] {e}", now());
                                                }
                                            }
                                        },
                                    ),
                                    Err((tmp, e)) => send(
                                        &response::failure::temporary::General {
                                            message: Some("Internal server error".to_string()),
                                        }
                                        .into_bytes(),
                                        stream,
                                        |result| {
                                            match result {
                                                Ok(()) => {
                                                    println!("[{}] [warning] [{peer}] {e}", now())
                                                }
                                                Err(e) => {
                                                    println!("[{}] [error] [{peer}] {e}", now())
                                                }
                                            };
                                            if let Err(e) = tmp.delete() {
                                                println!("[{}] [error] [{peer}] {e}", now());
                                            }
                                        },
                                    ),
                                };
                            }
                        }
                        Err(e) => {
                            return send(
                                &response::failure::temporary::General {
                                    message: Some("Internal server error".to_string()),
                                }
                                .into_bytes(),
                                stream,
                                |result| {
                                    match result {
                                        Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                    };
                                    if let Err(e) = tmp.delete() {
                                        println!("[{}] [error] [{peer}] {e}", now());
                                    }
                                },
                            )
                        }
                    }
                }
                Err(e) => {
                    if let Err(e) = tmp.delete() {
                        println!("[{}] [error] [{peer}] {e}", now());
                    }
                    return send(
                        &response::failure::temporary::General {
                            message: Some("Internal server error".to_string()),
                        }
                        .into_bytes(),
                        stream,
                        |result| match result {
                            Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                            Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                        },
                    );
                }
            }
        },
        Err(e) => send(
            &response::failure::temporary::General {
                message: Some("Internal server error".to_string()),
            }
            .into_bytes(),
            stream,
            |result| match result {
                Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            },
        ),
    }
}

fn close(stream: &mut TlsStream<TcpStream>) -> Result<()> {
    stream.flush()?;
    // close connection gracefully
    // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections
    stream.shutdown()?;
    Ok(())
}

fn send(data: &[u8], stream: &mut TlsStream<TcpStream>, callback: impl FnOnce(Result<()>)) {
    callback((|| {
        stream.write_all(data)?;
        close(stream)?;
        Ok(())
    })());
}

fn now() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis()
}