//! Source-to-source Luau emission.
//!
//! The current implementation is a CST-preserving bootstrap toward a full
//! pretty-printer: [`strip_types`] keeps parseable source text stable and
//! blanks typed-only syntax when callers request untyped output.

use crate::{
    parse::{Error, ParseConfig, parse_file_with},
    syntax::Expr,
    visit::{self, Visitor, WalkControl},
};

/// Options for [`strip_types`].
#[derive(Clone, Copy, Debug, Eq, PartialEq, Default)]
pub struct StripOptions {
    /// Parser configuration used before stripping.
    pub parse: ParseConfig,
    /// Preserve typed Luau syntax when true, returning the source unchanged.
    pub with_types: bool,
    /// Return code even when parsing reports recoverable errors.
    pub ignore_parse_errors: bool,
}

/// Blanks type annotations out of Luau source, preserving line and column
/// positions of the remaining code.
///
/// # Errors
/// Returns the parse errors when the source does not parse and
/// `ignore_parse_errors` is unset. Each error renders "line:col: message"
/// via its `Display` impl.
pub fn strip_types(source: &str, mut options: StripOptions) -> Result<String, Vec<Error>> {
    options.parse.store_cst_data = true;

    let parsed = parse_file_with(source, &options.parse);
    if !parsed.errors.is_empty() && !options.ignore_parse_errors {
        return Err(parsed.errors);
    }

    let mut code = if options.with_types {
        source.to_owned()
    } else {
        strip_untyped_syntax(source)
    };
    if options.ignore_parse_errors {
        insert_error_expr_markers(&mut code, source, &parsed.root);
    }

    Ok(code)
}

fn insert_error_expr_markers(code: &mut String, source: &str, root: &crate::syntax::Stat) {
    let mut visitor = ErrorExprMarkerVisitor {
        source,
        offsets: Vec::new(),
    };
    visit::walk_stat(root, &mut visitor);
    visitor.offsets.sort_unstable();
    visitor.offsets.dedup();

    for offset in visitor.offsets.into_iter().rev() {
        code.insert_str(offset, "(error-expr)");
    }
}

struct ErrorExprMarkerVisitor<'source> {
    source: &'source str,
    offsets: Vec<usize>,
}

impl Visitor<'_> for ErrorExprMarkerVisitor<'_> {
    fn visit_expr(&mut self, expr: &Expr) -> WalkControl {
        let Expr::Error {
            location,
            expressions,
            ..
        } = expr
        else {
            return WalkControl::Continue;
        };

        if !expressions.is_empty() {
            return WalkControl::Continue;
        }

        if let Some(location) = location
            && location.begin == location.end
            && let Some(offset) = location.begin.byte_offset(self.source)
        {
            self.offsets.push(offset);
        }

        WalkControl::Continue
    }
}

fn strip_untyped_syntax(source: &str) -> String {
    let bytes = source.as_bytes();
    let mut output = bytes.to_vec();
    let mut index = 0;

    while index < bytes.len() {
        if let Some(next) = skip_comment_or_string(bytes, index) {
            index = next;
        } else if starts_with(bytes, index, b"::") {
            let start = trim_horizontal_space_before(bytes, index);
            let end = type_suffix_end(bytes, index + 2);
            blank_ascii_range(&mut output, start, end);
            index = end;
        } else if bytes.get(index) == Some(&b':') && !starts_with(bytes, index + 1, b":") {
            if let Some(end) = annotation_end_before_assignment(bytes, index + 1) {
                blank_ascii_range(&mut output, index, end);
                index = end;
            } else {
                index += 1;
            }
        } else if starts_with(bytes, index, b"<<") {
            if let Some(end) = explicit_type_instantiation_end(bytes, index + 2) {
                blank_ascii_range(&mut output, index, end);
                index = end;
            } else {
                index += 1;
            }
        } else {
            index += 1;
        }
    }

    String::from_utf8(output).expect("ASCII blanking preserves valid UTF-8")
}

fn blank_ascii_range(output: &mut [u8], start: usize, end: usize) {
    if let Some(slice) = output.get_mut(start..end) {
        for byte in slice {
            if !matches!(*byte, b'\n' | b'\r') {
                *byte = b' ';
            }
        }
    }
}

fn skip_comment_or_string(bytes: &[u8], index: usize) -> Option<usize> {
    match bytes.get(index).copied()? {
        b'\'' | b'"' => Some(quoted_string_end(bytes, index + 1, bytes[index])),
        b'-' if starts_with(bytes, index + 1, b"-") => Some(comment_end(bytes, index + 2)),
        b'[' => long_bracket_end(bytes, index).map(|end| end.max(index + 1)),
        _ => None,
    }
}

fn quoted_string_end(bytes: &[u8], mut index: usize, quote: u8) -> usize {
    while index < bytes.len() {
        match bytes.get(index) {
            Some(&b'\\') => index = (index + 2).min(bytes.len()),
            Some(&q) if q == quote => return index + 1,
            _ => index += 1,
        }
    }
    index
}

fn comment_end(bytes: &[u8], index: usize) -> usize {
    if let Some(end) = long_bracket_end(bytes, index) {
        return end;
    }

    bytes
        .get(index..)
        .and_then(|slice| slice.iter().position(|&byte| byte == b'\n'))
        .map(|offset| index + offset)
        .unwrap_or(bytes.len())
}

fn long_bracket_end(bytes: &[u8], index: usize) -> Option<usize> {
    let equals = long_bracket_equals(bytes, index)?;
    let content_start = index + 2 + equals;
    let mut cursor = content_start;
    while cursor < bytes.len() {
        if bytes.get(cursor) == Some(&b']')
            && bytes
                .get(cursor + 1..cursor + 1 + equals)
                .is_some_and(|slice| slice.iter().all(|&byte| byte == b'='))
            && bytes.get(cursor + 1 + equals) == Some(&b']')
        {
            return Some(cursor + 2 + equals);
        }
        cursor += 1;
    }
    Some(bytes.len())
}

fn long_bracket_equals(bytes: &[u8], index: usize) -> Option<usize> {
    if bytes.get(index) != Some(&b'[') {
        return None;
    }

    let mut cursor = index + 1;
    while bytes.get(cursor) == Some(&b'=') {
        cursor += 1;
    }
    (bytes.get(cursor) == Some(&b'[')).then_some(cursor - index - 1)
}

fn trim_horizontal_space_before(bytes: &[u8], index: usize) -> usize {
    let mut start = index;
    while start > 0
        && bytes
            .get(start - 1)
            .is_some_and(|&b| matches!(b, b' ' | b'\t'))
    {
        start -= 1;
    }
    start
}

fn find_matching_generic_gt(bytes: &[u8], start: usize) -> Option<usize> {
    let mut depth = 1usize;
    let mut index = start;
    while index < bytes.len() {
        if let Some(next) = skip_comment_or_string(bytes, index) {
            index = next;
            continue;
        }
        match bytes.get(index).copied() {
            Some(b'(' | b'[' | b'{' | b'<') => depth += 1,
            Some(b')' | b']' | b'}' | b'>') => {
                depth -= 1;
                if depth == 0 {
                    return Some(index);
                }
            }
            Some(b';') if depth == 1 => return None,
            _ => {}
        }
        if depth == 1
            && (starts_with(bytes, index, b"==")
                || starts_with(bytes, index, b">=")
                || starts_with(bytes, index, b"<=")
                || starts_with(bytes, index, b"~="))
        {
            return None;
        }
        index += 1;
    }
    None
}

fn type_suffix_end(bytes: &[u8], mut index: usize) -> usize {
    let mut depth = 0usize;
    while index < bytes.len() {
        if let Some(next) = skip_comment_or_string(bytes, index) {
            index = next;
            continue;
        }
        if depth == 0 && (bytes.get(index) == Some(&b'<') || bytes.get(index) == Some(&b'>')) {
            let gt_index = (bytes.get(index) == Some(&b'<'))
                .then(|| find_matching_generic_gt(bytes, index + 1))
                .flatten();
            if let Some(gt_index) = gt_index {
                index = gt_index + 1;
                continue;
            }
            break;
        }
        match bytes.get(index).copied() {
            Some(b'(' | b'[' | b'{') => depth += 1,
            Some(b')' | b']' | b'}') if depth > 0 => depth -= 1,
            Some(
                b',' | b';' | b'=' | b'+' | b'-' | b'*' | b'/' | b'%' | b'^' | b')' | b']' | b'}',
            ) if depth == 0 => {
                break;
            }
            _ => {}
        }
        index += 1;
    }
    index
}

fn annotation_end_before_assignment(bytes: &[u8], index: usize) -> Option<usize> {
    let end = type_suffix_end(bytes, index);
    (bytes.get(end) == Some(&b'=')).then_some(end)
}

fn explicit_type_instantiation_end(bytes: &[u8], mut index: usize) -> Option<usize> {
    let mut depth = 1usize;
    while index < bytes.len() {
        if let Some(next) = skip_comment_or_string(bytes, index) {
            index = next;
            continue;
        }
        if starts_with(bytes, index, b"<<") {
            depth += 1;
            index += 2;
        } else if starts_with(bytes, index, b">>") {
            depth -= 1;
            index += 2;
            if depth == 0 {
                return Some(index);
            }
        } else {
            index += 1;
        }
    }
    None
}

fn starts_with(bytes: &[u8], index: usize, needle: &[u8]) -> bool {
    bytes
        .get(index..)
        .is_some_and(|slice| slice.starts_with(needle))
}

#[cfg(any())]
mod tests {
    use super::{StripOptions, strip_types};

    #[test]
    fn preserves_parseable_source_in_bootstrap_slice() {
        let source = " local x = 1 ";

        let printed = strip_types(source, StripOptions::default());

        assert_eq!(printed.expect("parses"), source);
    }

    #[test]
    fn reports_parse_error_when_not_ignored() {
        let printed = strip_types("local x =", StripOptions::default());

        let errors = printed.expect_err("parse error reported");
        assert!(!errors.is_empty());
        // The Display rendering carries location and message.
        assert!(errors[0].to_string().contains(':'));
    }

    #[test]
    fn renders_empty_error_expressions_when_errors_are_ignored() {
        let printed = strip_types(
            "\nrepeat\n    print(\"hello world\")\n",
            StripOptions {
                with_types: true,
                ignore_parse_errors: true,
                ..StripOptions::default()
            },
        );

        assert_eq!(
            printed.expect("ignored parse error"),
            "\nrepeat\n    print(\"hello world\")\n(error-expr)"
        );
    }

    #[test]
    fn strips_types_from_untyped_output_while_preserving_columns() {
        let printed = strip_types(
            " local s: string= f<<A, B>>() :: any+ g() :: number ",
            StripOptions::default(),
        );

        assert_eq!(
            printed.expect("parses"),
            " local s        = f        ()       + g()           "
        );
    }

    #[test]
    fn strips_types_with_operators_correctly() {
        let printed = strip_types("local x = y :: A > z", StripOptions::default());
        assert_eq!(printed.expect("parses"), "local x = y      > z");

        let printed = strip_types(
            "local x = y :: A < z",
            StripOptions {
                ignore_parse_errors: true,
                ..StripOptions::default()
            },
        );
        assert_eq!(printed.expect("parses"), "local x = y      < z");
    }

    #[test]
    fn strips_types_with_strings_containing_equals_correctly() {
        let printed = strip_types(
            "local x: \"foo=bar\" = \"foo=bar\"",
            StripOptions::default(),
        );
        assert_eq!(printed.expect("parses"), "local x            = \"foo=bar\"");
    }

    #[test]
    fn strips_types_with_comments_correctly() {
        let printed = strip_types(
            "local s: { a: string } -- comment\n = f()",
            StripOptions::default(),
        );
        assert_eq!(
            printed.expect("parses"),
            "local s                          \n = f()"
        );
    }
}
