use std::ops::{ControlFlow, Range};

use anyhow::{anyhow, Context};
use codespan_reporting::diagnostic::Diagnostic;
use tracing::{error, info};

use crate::{
    parse::{self, parse_toml_with_diagnostics, parse_tree_with_diagnostics},
    state::{report_diagnostics, FileId, Source, Treehouse},
    tree::ast::Branch,
    vfs::{self, Content, Dir, Edit, EditPath, VPath},
};

use super::{FixAllArgs, FixArgs};

struct Fix {
    range: Range<usize>,
    replacement: String,
}

#[derive(Default)]
struct State {
    fixes: Vec<Fix>,
}

fn dfs_fix_branch(treehouse: &mut Treehouse, file_id: FileId, state: &mut State, branch: &Branch) {
    let mut rng = rand::thread_rng();
    let ulid = ulid::Generator::new()
        .generate_with_source(&mut rng)
        .expect("failed to generate ulid for block"); // (wtf moment. you know how big the 80-bit combination space is?)

    let indent = " ".repeat(branch.indent_level);
    if let Some(attributes) = branch.attributes.clone() {
        // Scenario: Attributes need to be parsed as TOML and the id attribute has to be added into
        // the top-level table. Then we also need to pretty-print everything to match the right
        // indentation level.
        if let Ok(mut toml) =
            parse_toml_with_diagnostics(treehouse, file_id, attributes.data.clone())
        {
            if !toml.contains_key("id") {
                toml["id"] = toml_edit::value(ulid.to_string());
                toml.key_decor_mut("id")
                    .unwrap()
                    .set_prefix(" ".repeat(branch.indent_level + 2));
            }
            let mut toml_string = toml.to_string();

            // This is incredibly janky and barely works.
            let leading_spaces: usize = toml_string.chars().take_while(|&c| c == ' ').count();
            match leading_spaces {
                0 => toml_string.insert(0, ' '),
                1 => (),
                _ => toml_string.replace_range(0..leading_spaces - 1, ""),
            }

            let toml_string = fix_indent_in_generated_toml(&toml_string, branch.indent_level);

            state.fixes.push(Fix {
                range: attributes.data.clone(),
                replacement: toml_string,
            })
        }
    } else {
        // Scenario: No attributes at all.
        // In this case we can do a fast path where we generate the `% id = "whatever"` string
        // directly, not going through toml_edit.
        state.fixes.push(Fix {
            range: branch.kind_span.start..branch.kind_span.start,
            replacement: format!("% id = \"{ulid}\"\n{indent}"),
        });
    }

    // Then we fix child branches.
    for child in &branch.children {
        dfs_fix_branch(treehouse, file_id, state, child);
    }
}

fn fix_indent_in_generated_toml(toml: &str, min_indent_level: usize) -> String {
    let toml = toml.trim_end();

    let mut result = String::with_capacity(toml.len());

    for (i, line) in toml.lines().enumerate() {
        if line.is_empty() {
            result.push('\n');
        } else {
            let desired_line_indent_level = if i == 0 { 1 } else { min_indent_level + 2 };
            let leading_spaces: usize = line.chars().take_while(|&c| c == ' ').count();
            let needed_indentation = desired_line_indent_level.saturating_sub(leading_spaces);
            for _ in 0..needed_indentation {
                result.push(' ');
            }
            result.push_str(line);
            result.push('\n');
        }
    }

    for _ in 0..min_indent_level {
        result.push(' ');
    }

    result
}

pub fn fix_file(
    treehouse: &mut Treehouse,
    diagnostics: &mut Vec<Diagnostic<FileId>>,
    file_id: FileId,
) -> Result<String, parse::ErrorsEmitted> {
    let source = treehouse.source(file_id).input();
    parse_tree_with_diagnostics(file_id, source)
        .map(|roots| {
            let mut source = treehouse.source(file_id).input().to_owned();
            let mut state = State::default();

            for branch in &roots.branches {
                dfs_fix_branch(treehouse, file_id, &mut state, branch);
            }

            // Doing a depth-first search of the branches yields fixes from the beginning of the file
            // to its end. The most efficient way to apply all the fixes then is to reverse their order,
            // which lets us modify the source string in place because the fix ranges always stay
            // correct.
            for fix in state.fixes.iter().rev() {
                source.replace_range(fix.range.clone(), &fix.replacement);
            }

            source
        })
        .map_err(|mut new| {
            diagnostics.append(&mut new);
            parse::ErrorsEmitted
        })
}

pub fn fix_file_cli(fix_args: FixArgs, root: &dyn Dir) -> anyhow::Result<Edit> {
    let file = if &*fix_args.file == VPath::new("-") {
        std::io::read_to_string(std::io::stdin().lock()).context("cannot read file from stdin")?
    } else {
        vfs::query::<Content>(root, &fix_args.file)
            .ok_or_else(|| anyhow!("cannot read file to fix"))?
            .string()?
    };

    let mut treehouse = Treehouse::new();
    let mut diagnostics = vec![];
    let file_id = treehouse.add_file(fix_args.file.clone(), Source::Other(file));

    Ok(
        if let Ok(fixed) = fix_file(&mut treehouse, &mut diagnostics, file_id) {
            if fix_args.apply {
                let edit_path = vfs::query::<EditPath>(root, &fix_args.file).ok_or_else(|| {
                    anyhow!(
                        "{} is not an editable file (perhaps it is not in a persistent path?)",
                        fix_args.file
                    )
                })?;

                // Try to write the backup first. If writing that fails, bail out without overwriting
                // the source file.
                if let Some(backup_path) = fix_args.backup {
                    let backup_edit_path =
                        vfs::query::<EditPath>(root, &backup_path).ok_or_else(|| {
                            anyhow!("backup file {backup_path} is not an editable file")
                        })?;
                    Edit::Seq(vec![
                        Edit::Write(
                            backup_edit_path,
                            treehouse.source(file_id).input().to_owned().into(),
                        ),
                        Edit::Write(edit_path, fixed.into()),
                    ])
                } else {
                    Edit::Write(edit_path, fixed.into())
                }
            } else {
                println!("{fixed}");
                Edit::NoOp
            }
        } else {
            report_diagnostics(&treehouse, &diagnostics)?;
            Edit::NoOp
        },
    )
}

pub fn fix_all_cli(fix_all_args: FixAllArgs, dir: &dyn Dir) -> anyhow::Result<Edit> {
    let mut edits = vec![];

    fn fix_one(dir: &dyn Dir, path: &VPath) -> anyhow::Result<Edit> {
        if path.extension() == Some("tree") {
            let Some(content) = vfs::query::<Content>(dir, path).map(Content::bytes) else {
                return Ok(Edit::NoOp);
            };
            let content = String::from_utf8(content).context("file is not valid UTF-8")?;

            let mut treehouse = Treehouse::new();
            let mut diagnostics = vec![];
            let file_id = treehouse.add_file(path.to_owned(), Source::Other(content));
            let edit_path = vfs::query::<EditPath>(dir, path).context("path is not editable")?;

            if let Ok(fixed) = fix_file(&mut treehouse, &mut diagnostics, file_id) {
                if fixed != treehouse.source(file_id).input() {
                    return Ok(Edit::Write(edit_path, fixed.into()));
                }
            } else {
                report_diagnostics(&treehouse, &diagnostics)?;
            }
        }

        Ok(Edit::NoOp)
    }

    info!("gathering edits");
    vfs::walk_dir_rec(dir, VPath::ROOT, &mut |path| {
        match fix_one(dir, path) {
            Ok(Edit::NoOp) => (),
            Ok(edit) => edits.push(edit),
            Err(err) => error!("cannot fix {path}: {err:?}"),
        }

        ControlFlow::Continue(())
    });

    // NOTE: This number may be higher than you expect, because NoOp edits also count!
    info!("{} edits to apply", edits.len());

    if !fix_all_args.apply {
        info!("dry run; add `--apply` to apply changes");
        Ok(Edit::Dry(Box::new(Edit::All(edits))))
    } else {
        Ok(Edit::All(edits))
    }
}
