#![cfg_attr(nightly_build,feature(proc_macro_diagnostic))]

#[cfg(nightly_build)]
use proc_macro::Diagnostic;

extern crate proc_macro;
use std::collections::HashMap;

use litrs::{IntegerLit, StringLit};

use syn::*;
use syn::spanned::*;
use quote::*;

use std::io::Write;
use std::fmt::Write as OtherWrite;

// see https://github.com/dtolnay/syn/blob/master/examples/heapsize/heapsize_derive/src/lib.rs for
// mini example on how to use quote

// FIXME: implement: flags_after_position_arguments() 
// FIXME: generate fish completions?
// FIXME: make 'Option < .. >' stripping more robust
// FIXME: convert input_args to (usize, String), so improved debug info can be given (output full
// string, and an arrow below the faulty argument etc.

fn split_flags_from_other(input: proc_macro2::TokenStream) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
    let mut flags = proc_macro2::TokenStream::new();
    let mut other = proc_macro2::TokenStream::new();
    for token in input {
        if let proc_macro2::TokenTree::Group(group) = token {
            for token in group.stream() {
                if let proc_macro2::TokenTree::Literal(lit) = &token {
                    if lit.to_string().starts_with("\"-") {
                        flags.extend(proc_macro2::TokenStream::from(token));
                        flags.extend(quote!{,});
                    } else {
                        other.extend(proc_macro2::TokenStream::from(token));
                        other.extend(quote!{,});
                    }
                }
            }
        }
    }
    (flags,other)
}
// strip away '(' and ')'
fn unwrap_group(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
    let mut retval = proc_macro2::TokenStream::new();
    for token in input {
        if let proc_macro2::TokenTree::Group(group) = token {
            for token in group.stream() {
                retval.extend(proc_macro2::TokenStream::from(token));
            }
            break;
        }
    }
    retval
}

// strip away assignment operator
fn unwrap_assignment(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
    let mut retval = proc_macro2::TokenStream::new();
    for (i, token) in input.into_iter().enumerate() {
        if i != 0 {
            retval.extend(proc_macro2::TokenStream::from(token));
        }
    }
    retval
}

fn lit_string(string: &str) -> proc_macro2::TokenStream {
    proc_macro2::TokenStream::from(proc_macro2::TokenTree::from(proc_macro2::Literal::string(string)))
}

fn parse_string(string: &str) -> proc_macro2::TokenStream {
    string.parse::<proc_macro2::TokenStream>().unwrap()
}

// warning is not supported yet: https://github.com/rust-lang/rust/issues/54140
fn compile_error(string: &str) -> proc_macro::TokenStream {
    let expanded = quote!(
        compile_error!(#string);
        );
    proc_macro::TokenStream::from(expanded)
}

#[cfg(nightly_build)]
fn compile_warning(span: proc_macro2::Span, message: &str) {
    Diagnostic::spanned(span.unwrap(), proc_macro::Level::Warning, message)
    .emit();
}

#[cfg(not(nightly_build))]
fn compile_warning(_span: proc_macro2::Span, _message: &str) {
}

fn first_string(tokens: proc_macro2::TokenStream) -> String {
    for i in tokens {
        if let Ok(lit) = StringLit::try_from(i) {
            return lit.value().to_string();
        }
    }
    String::new()
}
fn strings(tokens: proc_macro2::TokenStream) -> Vec<String> {
    let mut retval = Vec::new();
    for i in tokens {
        if let Ok(lit) = StringLit::try_from(i) {
            retval.push(lit.value().to_string());
        }
    }
    retval
}
fn first_number(tokens: proc_macro2::TokenStream) -> usize {
    for i in tokens {
        if let Ok(lit) = IntegerLit::try_from(i) {
            return lit.value().unwrap();
        }
    }
    0
}
fn numbers(tokens: proc_macro2::TokenStream) -> Vec<usize> {
    let mut retval = Vec::new();
    for i in tokens {
        if let Ok(lit) = IntegerLit::try_from(i) {
            retval.push(lit.value().unwrap());
        }
    }
    retval
}
#[derive(Debug,Clone)]
struct Flag {
    option_names: Vec<String>,
    argument_descriptions: Vec<String>,
    description: String,
}

#[derive(Debug,Clone)]
struct Arguments {
    option_name: String,
    #[allow(dead_code)]
    min: usize,
    #[allow(dead_code)]
    max: usize,
}

#[derive(Debug,Clone)]
struct Command {
    command_name: String,
    short_description: String, // one line (without trailing .)
    flags: Vec<Flag>,
    positional: Option<Arguments>,
    subcommands: Option<String>,
    related: Vec<(String, usize)>, // only for programs
    long_description: String, // only for programs
}

static mut COMMANDS : Option<HashMap<String,Command>> = None;
static mut SUBCOMMANDS : Option<HashMap<String,Vec<String>>> = None;

#[proc_macro_derive(Command,attributes(related,description,command,flag_option,program,hide,subcommands,positional,flag))]
pub fn derive_command(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;

    match input.data {
        Data::Struct(ref data) => {
            let mut long_description : Option<String> = None;
            let mut related : Vec<(String,usize)> = Vec::new();
            let mut description = None;
            let mut command_name = None;
            let mut program_name = None;
            for a in &input.attrs {
                if a.path.get_ident().map_or(false, |x| x == "command") {
                    command_name = Some(unwrap_group(a.tokens.clone()));
                }
                if a.path.get_ident().map_or(false, |x| x == "description") {
                    description = Some(unwrap_group(a.tokens.clone()));
                }
                if a.path.get_ident().map_or(false, |x| x == "doc") {
                    let mut strings = strings(a.tokens.clone()).concat();
                    strings = strings.strip_prefix(' ').unwrap_or("").to_string();
                    if let Some(desc) = &mut long_description {
                        desc.push('\n');
                        desc.push_str(&strings);
                    } else {
                        long_description = Some(strings);
                    }
                }
                if a.path.get_ident().map_or(false, |x| x == "program") {
                    program_name = Some(unwrap_group(a.tokens.clone()));
                }
                if a.path.get_ident().map_or(false, |x| x == "related") {
                    related.push((first_string(unwrap_group(a.tokens.clone())), first_number(unwrap_group(a.tokens.clone()))));
                }
            }
            if !(command_name.is_some() ^ program_name.is_some()) {
                return compile_error("either command() or program() attribute is needed for this struct");
            }
            if let Some(name) = &program_name {
                if name.is_empty() {
                    let bin_name = r#"env!("CARGO_BIN_NAME")"#;
                    program_name = Some(bin_name.parse::<proc_macro2::TokenStream>().unwrap());
                }
            }
            // if description.is_none() && program_name.is_some() {
            //     let pkg_description = r#"env!("CARGO_PKG_DESCRIPTION")"#;
            //     description = Some(pkg_description.parse::<proc_macro2::TokenStream>().unwrap());
            // }
            if description.is_none() {
                return compile_error("description() attribute is needed for this struct");
            }
            // if long_description.is_none() {
            //     return compile_error("description rustdoc (///) is needed for this struct");
            // }

            let mut flags_vec = Vec::new();
            let mut flags = proc_macro2::TokenStream::new();
            let mut positional = None;
            let mut subcommands = None;
            let mut command_type = parse_string("commandy::TypeOfCommand::Other()");
            match data.fields {
                Fields::Named(ref fields) => {
                    fields.named.iter().for_each(|f| {
                        let name = &f.ident;
                        let mut flag_names = None;
                        let mut arguments = None;
                        let mut description = None;
                        let mut is_some = false;
                        for a in &f.attrs {
                            if a.path.get_ident().map_or(false, |x| x == "flag") {
                                let (parsed_flags, parsed_arguments) = split_flags_from_other(a.tokens.clone());
                                flag_names = Some(parsed_flags);
                                arguments = Some(parsed_arguments);
                                is_some = true;
                            }
                            if a.path.get_ident().map_or(false, |x| x == "doc") {
                                description = Some(unwrap_assignment(a.tokens.clone()));
                            }
                            if a.path.get_ident().map_or(false, |x| x == "positional") {
                                positional = Some((name, a.tokens.clone(), a.span()));
                                is_some = true;
                            }
                            if a.path.get_ident().map_or(false, |x| x == "subcommands") {
                                let mut subtype = f.ty.to_token_stream().to_string();
                                // FIXME: hacky, implement more robust conversion. See https://stackoverflow.com/questions/55271857/how-can-i-get-the-t-from-an-optiont-when-using-syn
                                if subtype.starts_with("Option < ") {
                                    subtype = subtype[9..subtype.len()-2].to_string();
                                }
                                subcommands = Some((name,subtype,a.span()));
                            }
                        }
                        if is_some && description.is_none() {
                            compile_warning(f.span(), "no description for flag given, add with rustdoc (///) comment");

                            description = Some(lit_string(""));
                        }
                        if let Some(flag_names) = flag_names {
                            flags.extend(quote_spanned! {f.span() =>
                                 commandy::Flag{
                                    option_names: &[#flag_names],
                                    argument_descriptions: &[#arguments],
                                    field: &mut self.#name,
                                    description: #description,
                                },
                            });
                            let argument_descriptions : Vec<String> = if let Some(arguments) = arguments { strings(arguments) } else { Vec::new() };
                            let option_names : Vec<String> = strings(flag_names);
                            flags_vec.push(Flag{
                                option_names,
                                argument_descriptions,
                                description: description.map_or(String::new(), |x| strings(x).concat().trim().into()),
                            });
                        }
                    });
                }
                Fields::Unnamed(_) | Fields::Unit => {
                    return compile_error("only named structs are supported for the Command macro");
                },
            }
            let mut command_check = None;
            if command_name.is_some() {
                command_check = Some(quote! {
                    if input_args.first().map(|x| x.as_str()) != Some(command_name) {
                        return Ok(false);
                    }
                });
            } else {
                command_name = program_name.clone();
            }
            if subcommands.is_some() {
                command_type = parse_string("commandy::TypeOfCommand::WithSubCommands()");
            }
            unsafe {
                if COMMANDS.is_none() {
                    COMMANDS = Some(HashMap::new());
                }
                if let Some(commands) = &mut COMMANDS {
                    let key = if let Some(pn) = &program_name { first_string(pn.clone()) } else { name.to_string()};
                    let mut pos = None;
                    if let Some((_name, tokens, _span)) = &positional {
                        let numbers = numbers(unwrap_group(tokens.clone()));
                        let mut strings = strings(unwrap_group(tokens.clone()));
                        pos = Some(Arguments {
                            option_name: strings.remove(0),
                            min: numbers[0],
                            max: numbers[1],
                        });
                    }
                    commands.insert(key, Command {
                        command_name: command_name.clone().map_or(String::new(), |x| strings(x).concat().trim().into()),
                        short_description: description.clone().map_or(String::new(), |x| strings(x).concat().trim().into()),
                        flags: flags_vec,
                        positional: pos,
                        subcommands: subcommands.as_ref().map(|x| x.1.clone()),
                        related,
                        long_description: long_description.unwrap_or_default(),
                    });
                    //eprintln!("COMMANDS = {:?}", commands);
                }
                if let Some(program_name) = program_name {
                    print_mage_page(first_string(program_name));
                }
            }
            let mut call_subcommands = proc_macro2::TokenStream::new();
            if let Some((subname,subcommands,span)) = subcommands {
                let subtype : proc_macro2::TokenStream = subcommands.parse().unwrap();
                call_subcommands = quote_spanned! {span =>
                    self.#subname = #subtype ::parse_subcommands(input_args, &format!("{}{} ", command_prefix, command_name), action)?;
                    if action == commandy::ParseArgumentAction::Parse && self.subcommand.is_none() {
                        let mut names = Vec::new();
                        #subtype::list_subcommands(&mut names);
                        eprintln!("expected a subcommand: {:?}\nbut could not parse subcommand from remaining command line arguments: {:?}", names, input_args);
                        std::process::exit(-1);
                    }
                };
            }
            let mut positional_args = proc_macro2::TokenStream::new();
            if let Some((posname, tokens, span)) = positional {
                positional_args = quote_spanned! {span =>
                    .positional(#tokens, &mut self.#posname)
                };
            }
            let expanded = quote! {
                impl commandy::ArgumentParser for #name {
                    fn insert_name(&self, names: &mut Vec<&'static str>) {
                        names.push(#command_name);
                    }
                    fn parse_arguments(&mut self, input_args: &mut Vec<String>, command_prefix: &str, action: commandy::ParseArgumentAction) -> std::result::Result<bool,String> {
                        let command_name = #command_name;
                        let flags = &mut [
                            #flags
                        ];
                        let mut command = commandy::Command {
                            command_name,
                            short_description: #description,
                            flags,
                            flags_after_position_arguments: false,
                            .. commandy::Command::default()
                        }
                        #positional_args
                        ;
                        if let commandy::ParseArgumentAction::ShowSmallOverview = action {
                            println!();
                            command.print_small_overview(command_prefix, &#command_type);
                            return Ok(false);
                        }
                        #command_check
                        input_args.remove(0);
                        let action = command.parse_args(input_args, command_prefix, &#command_type)?;
                        #call_subcommands
                        if let commandy::ParseArgumentAction::ShowSmallOverview = action {
                            std::process::exit(-1);
                        }
                        // if subcommand and is empty, generate error
                        Ok(true)
                    }
                }
            };
            //println!("RESULT\n{}", expanded);
            proc_macro::TokenStream::from(expanded)
        }
        Data::Enum(ref e) => {
            let mut indirections = Vec::new();
            let mut options = proc_macro2::TokenStream::new();
            let mut names = proc_macro2::TokenStream::new();
            for variant in &e.variants {
                let subident = &variant.ident;
                if variant.fields.len() != 1 {
                    return compile_error("only enums with a single field per option are supported");
                }
                for field in &variant.fields {
                    let ty = &field.ty;
                    indirections.push(ty.to_token_stream().to_string());
                    options.extend(quote_spanned! {field.span() =>
                        let mut option = #ty::default();
                        if (&mut option as &mut dyn commandy::ArgumentParser).parse_arguments(input_args, command_prefix, action)? {
                            return Ok(Some(Self::#subident(option)));
                        }
                    });
                    names.extend(quote_spanned! {field.span() =>
                        let option = #ty::default();
                        (&option as & dyn commandy::ArgumentParser).insert_name(names);
                    });
                }
            }
            unsafe {
                if SUBCOMMANDS.is_none() {
                    SUBCOMMANDS = Some(HashMap::new());
                }
                if let Some(subcommands) = &mut SUBCOMMANDS {
                    subcommands.insert(name.to_string(), indirections);
                }
            }
            let expanded = quote!(
                impl #name {
                    pub fn parse_subcommands(input_args: &mut Vec<String>, command_prefix: &str, action: commandy::ParseArgumentAction) -> Result<Option<Self>,String> {
                        #options
                        return Ok(None);
                    }
                    pub fn list_subcommands(names: &mut Vec<&'static str>) {
                        #names
                    }
                }
            );
            proc_macro::TokenStream::from(expanded)
        },
        Data::Union(_) => unimplemented!(),
    }
}

fn print_mage_page(program_name: String) {
    let out_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
    // skip man page generation on non release builds, to LSP is not slowed down etc
    if cfg!(debug_assertions) || cfg!(target_os = "windows") {
        //eprintln!("skipping man page generation of {}/{}.1 because compiling in debug mode", out_dir, program_name);
        return;
    }
    const TERM_BRIGHT_YELLOW   : &str = "\x1b[93m";
    const TERM_RESET           : &str = "\x1b[0m";
    const TERM_BOLD            : &str = "\x1b[1m";
    eprintln!("  {}{}Generating{} man page {}.1 ({})", TERM_BOLD, TERM_BRIGHT_YELLOW, TERM_RESET, program_name, out_dir);
	// see mdoc for more
    let mut contents = String::new();
    let command;
    unsafe {
        if let Some(commands) = &COMMANDS {
            command = &commands[&program_name];
        } else {
            eprintln!("NOT FOUND: {}", program_name);
            panic!();
        }
    }
    writeln!(contents, ".Dd {}", std::env::var("CARGO_PKG_VERSION").unwrap()).unwrap();
    writeln!(contents, ".Dt {program_name} 1").unwrap();
    writeln!(contents, ".Sh NAME").unwrap();
    writeln!(contents, ".Nm {program_name}").unwrap();
    writeln!(contents, ".Nd {}", command.short_description).unwrap();
    writeln!(contents, ".Sh SYNOPSIS").unwrap();
    writeln!(contents, ".Nm {program_name}").unwrap();

    // overview of usage
    print_man_flags_overview(&mut contents, command);

    // print long description
	if !command.long_description.is_empty() {
        writeln!(contents, ".Sh DESCRIPTION").unwrap();
        writeln!(contents, "{}", command.long_description).unwrap();
	}
    
    // print flags in detail
    print_man_flags(&mut contents, &command.flags);

    // print subcommands
    if let Some(subcommands) = &command.subcommands {
        print_subcommands(&mut contents, subcommands, &program_name);
    }


    if !command.related.is_empty() {
        writeln!(contents, ".Sh SEE ALSO").unwrap();
        for (name,section) in &command.related {
            writeln!(contents, ".Xr {} {}", name, section).unwrap();
        }
    }
    writeln!(contents, ".Sh AUTHORS").unwrap();
    writeln!(contents, "{}", std::env::var("CARGO_PKG_AUTHORS").unwrap()).unwrap();

    let file = format!("{}/{}.1", out_dir, program_name);
    let path = std::path::Path::new(&file);
    let mut file = std::fs::File::create(&path).unwrap();
    file.write_all(contents.as_bytes()).unwrap();
}

fn print_subcommands(contents: &mut String, subcommands_name: &String, command_prefix: &str) {
    let subs;
    unsafe {
        if let Some(subcommands) = &SUBCOMMANDS {
            subs = &subcommands[subcommands_name];
        } else {
            eprintln!("SUBCOMMANDS NOT FOUND: {}", subcommands_name);
            panic!();
        }
    }
    for sub in subs {
        let command;
        unsafe {
            if let Some(commands) = &COMMANDS {
                command = &commands[sub];
            } else {
                eprintln!("COMMAND NOT FOUND: {}", sub);
                panic!();
            }
        }
        writeln!(contents, ".Sh SUBCOMMAND {} {}", command_prefix, command.command_name).unwrap();
        writeln!(contents, ".Nm {} {}", command_prefix, command.command_name).unwrap();
        print_man_flags_overview(contents, command);
        writeln!(contents).unwrap();
        writeln!(contents, "{}", command.short_description).unwrap();
        print_man_flags(contents, &command.flags);
        if let Some(subcommands) = &command.subcommands {
            print_subcommands(contents, subcommands, &format!("{} {}", command_prefix, command.command_name));
        }
    }
}

fn print_man_flags_overview(contents: &mut String, command: &Command) {
    for flag in &command.flags {
        if flag.option_names.is_empty() {
            continue;
        }
        write!(contents, ".Op Fl {}", flag.option_names.first().unwrap().strip_prefix('-').unwrap()).unwrap();
        for arg in &flag.argument_descriptions {
            write!(contents, " Ar {}", arg).unwrap();
        }
        writeln!(contents).unwrap();
    }
    if let Some(positional) = &command.positional {
        writeln!(contents, ".Ar {}", positional.option_name).unwrap();
    }
    if command.subcommands.is_some() {
        writeln!(contents, ".Ar subcommand").unwrap();
    }
}

fn print_man_flags(contents: &mut String, flags: &[Flag]) {
    if flags.is_empty() {
        return;
    }
    writeln!(contents, ".Sh FLAGS").unwrap();
    writeln!(contents, ".Bl -tag -width Ds").unwrap();
    for flag in flags {
        //.Op Fl o Ar file Ar file2
        write!(contents, ".It Fl ").unwrap();
        write!(contents, "{}", flag.option_names.first().unwrap().strip_prefix('-').unwrap()).unwrap();
        for name in flag.option_names.iter().skip(1) {
            write!(contents, ",{}", name).unwrap();
        }
        for arg in &flag.argument_descriptions {
            write!(contents, " Ar {}", arg).unwrap();
        }
        writeln!(contents).unwrap();
        writeln!(contents, "{}", flag.description).unwrap();
    }
    writeln!(contents, ".El").unwrap();
}