#![feature(seek_stream_len, slice_split_once, slice_internals)]
#![allow(internal_features)]

use std::{
    fs::File,
    io::{self, stdout, Seek, Stdin, Write},
    path::{Path, PathBuf},
    process::ExitCode,
};

use clap::Parser;
use cross::mem::FileMapping;
use regex::bytes::Regex;

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Invalid expression: {0}")]
    Regex(#[from] regex::Error),
}

trait Input {
    fn next_line(&mut self) -> io::Result<Option<&[u8]>>;
}

struct FileInput {
    mmap: FileMapping,
    position: usize,
}

struct StdinInput {
    stdin: Stdin,
    buffer: String,
    eof: bool,
}

impl StdinInput {
    pub fn new() -> Self {
        Self {
            stdin: io::stdin(),
            buffer: String::new(),
            eof: false,
        }
    }
}

impl FileInput {
    pub fn open<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let mut file = File::open(path)?;
        let size = file.stream_len()?.try_into().unwrap();
        // TODO handle non-mmap()able files
        let mmap = FileMapping::map(file, size, false)?;
        Ok(Self { mmap, position: 0 })
    }
}

impl Input for StdinInput {
    fn next_line(&mut self) -> io::Result<Option<&[u8]>> {
        if self.eof {
            return Ok(None);
        }

        self.buffer.clear();
        if self.stdin.read_line(&mut self.buffer)? == 0 {
            self.eof = true;
            return Ok(None);
        }

        Ok(Some(self.buffer.trim_end_matches('\n').as_bytes()))
    }
}

impl Input for FileInput {
    fn next_line(&mut self) -> io::Result<Option<&[u8]>> {
        if self.position >= self.mmap.len() {
            return Ok(None);
        }

        if let Some(nl) = core::slice::memchr::memchr(b'\n', &self.mmap[self.position..]) {
            let line = &self.mmap[self.position..self.position + nl];
            self.position += line.len() + 1;
            Ok(Some(line))
        } else {
            let slice = &self.mmap[self.position..];
            self.position = self.mmap.len();
            Ok(Some(slice))
        }
    }
}

fn run_on<I: Input>(input: &mut I, regex: &Regex) -> Result<bool, Error> {
    let mut any = false;
    loop {
        let Some(line) = input.next_line()? else {
            break;
        };

        if regex.is_match(line) {
            any = true;
            stdout().write_all(line).ok();
            stdout().write_all(b"\n").ok();
        }
    }
    Ok(any)
}

#[derive(Debug, Parser)]
struct Args {
    expression: String,
    filename: Option<PathBuf>,
}

fn run(args: &Args) -> Result<bool, Error> {
    let regex = Regex::new(&args.expression)?;
    if let Some(filename) = args.filename.as_deref() {
        let mut input = FileInput::open(filename)?;
        run_on(&mut input, &regex)
    } else {
        let mut input = StdinInput::new();
        run_on(&mut input, &regex)
    }
}

fn main() -> ExitCode {
    let args = Args::parse();
    match run(&args) {
        Ok(true) => ExitCode::SUCCESS,
        Ok(false) => ExitCode::FAILURE,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}
