#![deny(warnings, rust_2018_idioms)]
// Updates flatbuffers: fbs/*.fbs and src/fbs/*.rs. Executed via update-fbs.sh.

use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, ExitStatus};
use std::{env, fs, io};

use futures::future;
use hyper::Method;
use log::{error, info};

use kapiti::filter::fetcher::Fetcher;
use kapiti::logging;

macro_rules! io_err {
    ($($arg:tt)*) => (Err(io::Error::new(io::ErrorKind::InvalidData, format!($($arg)+))))
}

#[tokio::main]
async fn main() -> Result<(), io::Error> {
    logging::setup(None)?;

    let fbs_specs_dir = Path::new("fbs");
    let csv_dir = fbs_specs_dir.join("enum-specs");
    let sections_csv_path = csv_dir.join("sections.csv");
    let sections = csv::Reader::from_path(&sections_csv_path)
        .expect(format!("opening {:?}", sections_csv_path).as_str())
        .deserialize()
        .map(|v| v.expect("deserializing sections.csv entry"))
        .collect();

    match env::args().nth(1) {
        Some(val) => {
            if val == "--fetch" {
                info!("Fetching csvs (--fetch):");
                download_csvs(&csv_dir, &sections).await?;
            } else {
                panic!("Unrecognized arg: {}", val);
            }
        }
        None => info!("Skipping csv download (no --fetch)"),
    }

    // Set up path for generated rust code (both ours and flatc output)
    let rs_dir = Path::new("src").join("fbs");
    if fs::metadata(&rs_dir).is_ok() {
        // Clean any existing content.
        fs::remove_dir_all(&rs_dir).expect(format!("removing output dir: {:?}", rs_dir).as_str());
    }
    // Not needed by flatc, but we want it for writing dns_enums_conv.rs
    fs::create_dir_all(&rs_dir).expect(format!("creating output dir: {:?}", rs_dir).as_str());

    let enums_conv_rs = rs_dir.join("dns_enums_conv.rs");

    // Generate ROOT/fbs/dns_enums.fbs and ROOT/src/fbs/dns_enums_conv.rs from CSV specs in ROOT/fbs/enum-specs/*
    generate_enums_fbs(fbs_specs_dir, &enums_conv_rs, &csv_dir, &sections)?;

    // Find all .fbs files in ROOT/fbs/, to be used for codegen and then mod.rs afterwards
    let mut fbs_files = fs::read_dir(fbs_specs_dir)
        .expect(format!("iterating input dir: {:?}", fbs_specs_dir).as_str())
        .map(|entry| {
            entry
                .expect(format!("resolving entry in input dir: {:?}", fbs_specs_dir).as_str())
                .path()
        })
        .filter(|pathbuf| pathbuf.extension().map_or(false, |ext| ext == "fbs"))
        .collect::<Vec<PathBuf>>();
    fbs_files.sort();

    // Generate ROOT/src/fbs/*.rs from ROOT/fbs/*.fbs
    let status = flatc(&fbs_files, &rs_dir)?;
    if !status.success() {
        return io_err!("Failed to execute flatc");
    }

    // Generate ROOT/src/fbs/mod.rs
    generate_modrs(&rs_dir, &fbs_files, enums_conv_rs)?;

    Ok(())
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct SectionsRecord {
    Filename: String,
    EnumName: String,
    DataType: String,
}

/// Redownloads CSV specs from IANA. This is then consumed by this tool.
async fn download_csvs(csv_dir: &Path, sections: &Vec<SectionsRecord>) -> Result<(), io::Error> {
    let fetcher = Fetcher::new(1024 * 1024, Some("text/csv".to_string()));
    let client = Fetcher::external_client();

    let mut responses = future::try_join_all(sections.into_iter().map(|section| {
        let url = format!(
            "https://www.iana.org/assignments/dns-parameters/{}",
            section.Filename.as_str()
        );
        let path = csv_dir.join(&section.Filename);
        info!("- Downloading {:?} to {:?}", url, path);

        let req = fetcher.build_request(Method::GET, &url);
        client.request(req)
    }))
    .await
    .expect("Failed to query sections");

    let request_count = sections.len();
    for i in 0..request_count {
        let mut resp = responses
            .get_mut(i)
            .expect(format!("Missing response index {}", i).as_str());
        let section = sections
            .get(i)
            .expect(format!("Missing section index {}", i).as_str());
        let path = csv_dir.join(&section.Filename);

        info!(
            "[{}/{}] {:?} status: {}",
            i + 1,
            request_count,
            section.Filename,
            resp.status()
        );
        fetcher
            .handle_response(&section.Filename, path.as_path(), &mut resp)
            .await?;
    }

    Ok(())
}

/// Generates dns_enums.fbs from IANA csv spec files. This is then consumed by flatc.
fn generate_enums_fbs(
    fbs_specs_dir: &Path,
    enums_conv_rs: &Path,
    csv_dir: &Path,
    sections: &Vec<SectionsRecord>,
) -> Result<(), io::Error> {
    let fbspath = fbs_specs_dir.join("dns_enums.fbs");
    let mut fbsfile = fs::File::create(&fbspath).expect("creating dns_enums.fbs file");
    let mut convfile = fs::File::create(&enums_conv_rs).expect("creating dns_enums_conv.rs file");
    info!("Generating {:?} and {:?}", &fbspath, enums_conv_rs);
    fbsfile.write(b"// This file is autogenerated by update_fbs.rs. Don't touch.\n")?;

    convfile.write(b"// This file is autogenerated by update_fbs.rs. Don't touch.\n\n")?;
    convfile.write(b"use super::dns_enums_generated;\n")?;

    // Iterate over entries in input, open their respective files
    for section in sections {
        let enum_csv_path = csv_dir.join(&section.Filename);
        info!("- Parsing {:?} from {:?}", section.EnumName, enum_csv_path);
        let entries = match section.EnumName.as_str() {
            // Rather than trying to make everything perfectly generic,
            // just have per-file/entry tweaks/cleanup via separate functions.
            "ResourceClass" => generate_resourceclass_enum(enum_csv_path)?,
            "ResourceType" => generate_resourcetype_enum(enum_csv_path)?,
            "OpCode" => generate_opcode_enum(enum_csv_path)?,
            "ResponseCode" => generate_responsecode_enum(enum_csv_path)?,
            "OPTOptionCode" => generate_optoptioncode_enum(enum_csv_path)?,
            &_ => panic!("Unsupported enum: {}", section.EnumName),
        };

        // Write enum entries to enums.fbs
        fbsfile
            .write(format!("\nenum {} : {} {{\n", section.EnumName, section.DataType).as_bytes())?;
        for entry in &entries {
            if entry.doc {
                fbsfile.write(format!("\n  /// {}\n", entry.comment1).as_bytes())?;
                if let Some(c2) = &entry.comment2 {
                    fbsfile.write(format!("  /// {}\n", c2).as_bytes())?;
                }
            } else {
                fbsfile.write(format!("\n  // {}\n", entry.comment1).as_bytes())?;
                if let Some(c2) = &entry.comment2 {
                    fbsfile.write(format!("  // {}\n", c2).as_bytes())?;
                }
            }
            if let Some(n) = &entry.name {
                if let Some(i) = &entry.intval {
                    fbsfile.write(format!("  {} = {},\n", n, i).as_bytes())?;
                } else {
                    panic!("Missing entry.intval for entry.name={}", n)
                }
            }
        }
        fbsfile.write(b"\n}\n")?;

        // Write enum entries to enum_conv.rs

        // string -> enum
        convfile.write(
            format!(
                "\npub fn {}_str(s: String) -> Option<dns_enums_generated::{}> {{\n",
                section.EnumName.to_lowercase(),
                section.EnumName
            )
            .as_bytes(),
        )?;
        convfile.write(b"    match s.as_str() {\n")?;
        for entry in &entries {
            if let Some(n) = &entry.name {
                convfile.write(
                    format!(
                        "        \"{}\" => Some(dns_enums_generated::{}::{}),\n",
                        n, section.EnumName, n
                    )
                    .as_bytes(),
                )?;
            }
        }
        convfile.write(b"        _ => None\n")?;
        convfile.write(b"    }\n")?;
        convfile.write(b"}\n")?;

        // usize -> enum
        convfile.write(
            format!(
                "\npub fn {}_int(i: usize) -> Option<dns_enums_generated::{}> {{\n",
                section.EnumName.to_lowercase(),
                section.EnumName
            )
            .as_bytes(),
        )?;
        convfile.write(b"    match i {\n")?;
        for entry in &entries {
            if let Some(n) = &entry.name {
                if let Some(i) = &entry.intval {
                    convfile.write(
                        format!(
                            "        {} => Some(dns_enums_generated::{}::{}),\n",
                            i, section.EnumName, n
                        )
                        .as_bytes(),
                    )?;
                } else {
                    panic!("Missing entry.intval for entry.name={}", n)
                }
            }
        }
        convfile.write(b"        _ => None\n")?;
        convfile.write(b"    }\n")?;
        convfile.write(b"}\n")?;
    }
    Ok(())
}

struct EnumEntry {
    /// Whether to use 2 slashes (false) or 3 slashes (true) on comment1 and comment2
    doc: bool,
    /// First comment line
    comment1: String,
    /// Second comment line
    comment2: Option<String>,
    /// The name of the enum, or None (comment-only)
    name: Option<String>,
    /// The int value of the enum, ignored if name=None
    intval: Option<String>,
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct ResourceClassRecord {
    Decimal: String,
    Name: String,
    Reference: String,
}

fn generate_resourceclass_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>, io::Error> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .expect(format!("opening {:?}", enum_csv_path).as_str())
        .deserialize()
    {
        let record: ResourceClassRecord =
            record.expect(format!("reading/deserializing from {:?}", enum_csv_path).as_str());
        let ignored = record.Decimal.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or reference
        let mut comment1 = record.Name.clone();
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: Some(record.Decimal),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Name.as_str() {
                // When the value is reserved, include the number in the enum name
                "Reserved" => format!("RESERVED_{}", record.Decimal),
                // Manually make these names usable
                "QCLASS NONE" => "NONE".to_string(),
                "QCLASS * (ANY)" => "ANY".to_string(),
                // All others: Use the first word, capitalized
                _ => match record.Name.split_ascii_whitespace().next() {
                    Some(first_word) => first_word.to_uppercase(),
                    None => panic!("Missing name '{}' in {:?}", record.Name, enum_csv_path),
                },
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: None,
                name: Some(format!("CLASS_{}", name)),
                intval: Some(record.Decimal),
            });
        }
    }
    Ok(entries)
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct ResourceTypeRecord {
    Value: String,
    TYPE: String,
    Meaning: String,
    Reference: String,
}

fn generate_resourcetype_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>, io::Error> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .expect(format!("opening {:?}", enum_csv_path).as_str())
        .deserialize()
    {
        let record: ResourceTypeRecord =
            record.expect(format!("reading/deserializing from {:?}", enum_csv_path).as_str());
        let ignored = record.Value.contains('-') || record.TYPE == "Unassigned";

        // Enum comment: Type or Reference+Meaning
        let mut comment1: String;
        if record.Reference.is_empty() && record.Meaning.is_empty() {
            // No additional info, just write the "TYPE" (usually something like 'Unassigned')
            comment1 = record.TYPE.clone();
        } else {
            // Instead of the TYPE, write the Meaning and/or Reference
            comment1 = String::new();
            if !record.Meaning.is_empty() {
                comment1.push_str(record.Meaning.replace('\n', " ").as_str());
            }
            if !record.Reference.is_empty() {
                if !record.Meaning.is_empty() {
                    comment1.push(' ');
                }
                comment1.push_str(record.Reference.replace('\n', " ").as_str())
            }
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: Some(record.Value),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.TYPE.as_str() {
                // When the value is reserved, include the number in the enum name
                "Reserved" => format!("RESERVED_{}", record.Value),
                // Manually make these names usable
                "*" => "ANY".to_string(),
                // All others: Use the type as-is
                _ => match record.Meaning.contains("OBSOLETE") {
                    true => record.TYPE.replace('-', "_") + "_OBSOLETE",
                    false => record.TYPE.replace('-', "_"),
                },
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: None,
                name: Some(format!("TYPE_{}", name)),
                intval: Some(record.Value),
            });
        }
    }
    Ok(entries)
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct OpCodeRecord {
    OpCode: String,
    Name: String,
    Reference: String,
}

fn generate_opcode_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>, io::Error> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .expect(format!("opening {:?}", enum_csv_path).as_str())
        .deserialize()
    {
        let record: OpCodeRecord =
            record.expect(format!("reading/deserializing from {:?}", enum_csv_path).as_str());
        let ignored = record.OpCode.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or reference
        let mut comment1 = record.Name.clone();
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: Some(record.OpCode),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Name.as_str() {
                // Manually make these names usable
                "DNS Stateful Operations (DSO)" => "DSO".to_string(),
                "IQuery  (Inverse Query, OBSOLETE)" => "IQUERY_OBSOLETE".to_string(),
                // All others: Use the name as-is
                _ => record.Name.to_uppercase(),
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: None,
                name: Some(format!("OP_{}", name)),
                intval: Some(record.OpCode),
            });
        }
    }
    Ok(entries)
}

// RCODE,Name,Description,Reference
#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct ResponseCodeRecord {
    RCODE: String,
    Name: String,
    Description: String,
    Reference: String,
}

fn generate_responsecode_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>, io::Error> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .expect(format!("opening {:?}", enum_csv_path).as_str())
        .deserialize()
    {
        let record: ResponseCodeRecord =
            record.expect(format!("reading/deserializing from {:?}", enum_csv_path).as_str());
        let ignored = record.RCODE.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or description/reference
        let mut comment1 = record.Name.clone();
        if !record.Description.is_empty() {
            comment1.push_str(format!(": {}", record.Description.replace('\n', " ")).as_str());
        }
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: Some(record.RCODE),
                name: None,
                intval: None,
            });
        } else if record.Description == "Not Authorized" {
            // Don't print out a duplicate RCODE=9 for both versions of "NotAuth". Let the first one handle it.
            continue;
        } else if record.Name == "BADSIG" {
            // Don't print out a duplicate RCODE=16 value for both BADVERS+BADSIG. Let BADVERS handle it.
            continue;
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Name.as_str() {
                // Manually make these names usable
                "Reserved, can be allocated by Standards Action" => "RESERVED_65535".to_string(),
                // BADVERS and BADSIG share RCODE=16, so include both in the name. We skip BADSIG above.
                "BADVERS" => "BADVERS_BADSIG".to_string(),
                // All others: Use the name as-is
                _ => record.Name.to_uppercase(),
            };
            let comment2 = match record.Name.as_str() {
                "NotAuth" => Some(String::from("NotAuth: Not Authorized [RFC2845]")),
                "BADVERS" => Some(String::from("BADSIG: TSIG Signature Failure [RFC2845]")),
                _ => None,
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: comment2,
                name: Some(format!("RESPONSE_{}", name)),
                intval: Some(record.RCODE),
            });
        }
    }
    Ok(entries)
}

// RCODE,Name,Description,Reference
#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct OPTOptionCodeRecord {
    Value: String,
    Name: String,
    Status: String,
    Reference: String,
}

fn generate_optoptioncode_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>, io::Error> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .expect(format!("opening {:?}", enum_csv_path).as_str())
        .deserialize()
    {
        let record: OPTOptionCodeRecord =
            record.expect(format!("reading/deserializing from {:?}", enum_csv_path).as_str());
        let ignored = record.Value.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or status/reference
        let mut comment1 = record.Name.clone();
        if !record.Status.is_empty() {
            comment1.push_str(format!(": {}", record.Status.replace('\n', " ")).as_str());
        }
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: Some(record.Value),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Value.as_str() {
                // Override the csv-defined name of "Reserved" for the zero value
                // "EXTENDED-RCODE value 0 indicates that an unextended RCODE is in use"
                "0" => "NONE".to_string(),
                _ => match record.Name.as_str() {
                    // When the value is reserved, include the number in the enum name
                    "Reserved" => format!("RESERVED_{}", record.Value),
                    // All others: Use the type as-is (but with '-' and ' ' cleaned up)
                    _ => record
                        .Name
                        .to_uppercase()
                        .replace('-', "_")
                        .replace(' ', "_"),
                },
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1: comment1,
                comment2: None,
                name: Some(format!("OPTOPTION_{}", name)),
                intval: Some(record.Value),
            });
        }
    }
    Ok(entries)
}

/// Generates rust framebuffer code from <spec_files> to <output_dir>/*.rs.
fn flatc(spec_files: &Vec<PathBuf>, output_dir: &Path) -> Result<ExitStatus, io::Error> {
    info!("Running flatc: {:?} => {:?}", spec_files, output_dir);

    // Before starting, check that the user has flatc installed
    // As a side effect the flatbuffers version is written to stdout, but that may be handy for debugging someday.
    if let Err(_) = Command::new("flatc").arg("--version").status() {
        error!("Unable to run 'flatc --version'. Ensure that flatbuffers is installed.");
        return io_err!("Failed to run flatc: install flatbuffers");
    }

    Command::new("flatc")
        .arg("-o")
        .arg(
            output_dir
                .to_str()
                .expect("Unable to convert output dir to str"),
        )
        .arg("--rust")
        .args(spec_files.iter().map(|p| {
            p.to_str()
                .expect("failed to convert flatc input path to string")
        }))
        .status()
}

/// Generates output/mod.rs so that the output package can be consumed by the rest of the project.
fn generate_modrs(
    output: &Path,
    fbs_files: &Vec<PathBuf>,
    enums_conv_rs: PathBuf,
) -> Result<(), io::Error> {
    let modrs_path = output.join("mod.rs");
    info!("Generating {:?}", modrs_path);
    let mut modrs = fs::File::create(&modrs_path).expect("creating mod.rs file");
    modrs.write(b"// This file is autogenerated by update_fbs.rs. Don't touch.\n\n")?;

    for path in fbs_files {
        if let Some(filestem) = path.file_stem() {
            let filestem = filestem
                .to_str()
                .expect(format!("formatting path as utf-8: {:?}", path).as_str());
            modrs.write(
                format!(
                    "#[allow(unused_imports, dead_code)]\npub mod {}_generated;\n",
                    filestem
                )
                .as_bytes(),
            )?;
        }
    }

    modrs.write(b"\n")?;

    // Omit the #[allow] and the "_generated" suffix when writing the dns_enums_conv entry
    if let Some(filestem) = enums_conv_rs.file_stem() {
        let filestem = filestem
            .to_str()
            .expect(format!("formatting path as utf-8: {:?}", enums_conv_rs).as_str());
        modrs.write(format!("pub mod {};\n", filestem).as_bytes())?;
    }

    Ok(())
}