/*
 * Takes the first n lines of input, passes to stdout.
 */

use std::error::Error;
use std::io::{self, Read, Write, BufRead};
use structopt::StructOpt;
use std::path::PathBuf;

#[derive(StructOpt)]
#[structopt(rename_all = "kebab-case", about = "take first N rows of input (AFT-aware)" )]
struct Opt {
    #[structopt(long = "--lines", short = "-n", default_value = "10")]
    lines: u64,

    #[structopt(parse(from_os_str))]
    input: Option<PathBuf>,
}

fn main() -> Result<(), Box<dyn Error>> {

    let opt = Opt::from_args();

    // treat "-" as "use stdin"
    let input = match opt.input {
        Some(s) if s.to_string_lossy() == "-" => None,
        other => other,
    };
    let stdin = io::stdin();
    let stdin = stdin.lock();
    let mut input: Box<dyn BufRead> = match input {
        None => {
            Box::new(io::BufReader::new(stdin))
        },
        Some(path) => {
            let f = std::fs::File::open(path)?;
            Box::new(io::BufReader::new(f))
        },
    };
    let mut stdout = io::stdout();

    // read first byte of input to check if we are getting AFT; if not just pass through
    let mut first_byte = [0; 1];
    let got = input.read(&mut first_byte)?;
    if got != 1 {
        panic!("couldn't read a byte from input");
    };

    let is_aft = first_byte[0] == 0x01;
    if !is_aft {
        // simple implementation of lines
        stdout.write(&first_byte)?;
        for l in input.lines().take(opt.lines as usize) {
            writeln!(stdout, "{}", l?)?;
        }
        return Ok(())
    };

    let mut header_bytes = vec![];
    // TODO: sanity limit on length
    input.read_until(0x02, &mut header_bytes)?;
    if header_bytes[header_bytes.len()-1] != 0x02 {
        panic!("expected an AFT header");
    };
    stdout.write(&first_byte)?;
    stdout.write(&header_bytes)?;
    for l in input.lines().take(opt.lines as usize) {
        writeln!(stdout, "{}", l?)?;
    }
    Ok(())
}