use std::collections::{BTreeMap, BTreeSet, VecDeque};

pub enum StringContextElem {
    Opaque { path: String },
    DrvDeep { drv_path: String },
    Built { drv_path: String, output: String },
}

impl StringContextElem {
    pub fn decode(encoded: &str) -> Self {
        if let Some(drv_path) = encoded.strip_prefix('=') {
            StringContextElem::DrvDeep {
                drv_path: drv_path.to_string(),
            }
        } else if let Some(rest) = encoded.strip_prefix('!') {
            if let Some(second_bang) = rest.find('!') {
                let output = rest[..second_bang].to_string();
                let drv_path = rest[second_bang + 1..].to_string();
                StringContextElem::Built { drv_path, output }
            } else {
                StringContextElem::Opaque {
                    path: encoded.to_string(),
                }
            }
        } else {
            StringContextElem::Opaque {
                path: encoded.to_string(),
            }
        }
    }
}

pub type InputDrvs = BTreeMap<String, BTreeSet<String>>;
pub type Srcs = BTreeSet<String>;
pub fn extract_input_drvs_and_srcs(context: &[String]) -> Result<(InputDrvs, Srcs), String> {
    let mut input_drvs: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
    let mut input_srcs: BTreeSet<String> = BTreeSet::new();

    for encoded in context {
        match StringContextElem::decode(encoded) {
            StringContextElem::Opaque { path } => {
                input_srcs.insert(path);
            }
            StringContextElem::DrvDeep { drv_path } => {
                compute_fs_closure(&drv_path, &mut input_drvs, &mut input_srcs)?;
            }
            StringContextElem::Built { drv_path, output } => {
                input_drvs.entry(drv_path).or_default().insert(output);
            }
        }
    }

    Ok((input_drvs, input_srcs))
}

fn compute_fs_closure(
    drv_path: &str,
    input_drvs: &mut BTreeMap<String, BTreeSet<String>>,
    input_srcs: &mut BTreeSet<String>,
) -> Result<(), String> {
    let mut queue: VecDeque<String> = VecDeque::new();
    let mut visited: BTreeSet<String> = BTreeSet::new();

    queue.push_back(drv_path.to_string());

    while let Some(current_path) = queue.pop_front() {
        if visited.contains(&current_path) {
            continue;
        }
        visited.insert(current_path.clone());
        input_srcs.insert(current_path.clone());

        if !current_path.ends_with(".drv") {
            continue;
        }

        let content = std::fs::read_to_string(&current_path)
            .map_err(|e| format!("failed to read derivation {}: {}", current_path, e))?;

        let inputs = parse_derivation_inputs(&content)
            .ok_or_else(|| format!("failed to parse derivation {}", current_path))?;

        for src in inputs.input_srcs {
            input_srcs.insert(src.clone());
            if !visited.contains(&src) {
                queue.push_back(src);
            }
        }

        for (dep_drv, outputs) in inputs.input_drvs {
            input_srcs.insert(dep_drv.clone());
            let entry = input_drvs.entry(dep_drv.clone()).or_default();
            for output in outputs {
                entry.insert(output);
            }
            if !visited.contains(&dep_drv) {
                queue.push_back(dep_drv);
            }
        }
    }

    Ok(())
}

struct DerivationInputs {
    input_drvs: Vec<(String, Vec<String>)>,
    input_srcs: Vec<String>,
}

fn parse_derivation_inputs(aterm: &str) -> Option<DerivationInputs> {
    let aterm = aterm.strip_prefix("Derive([")?;

    let mut bracket_count: i32 = 1;
    let mut pos = 0;
    let bytes = aterm.as_bytes();
    while pos < bytes.len() && bracket_count > 0 {
        match bytes[pos] {
            b'[' => bracket_count += 1,
            b']' => bracket_count -= 1,
            _ => {}
        }
        pos += 1;
    }
    if bracket_count != 0 {
        return None;
    }

    let rest = &aterm[pos..];
    let rest = rest.strip_prefix(",[")?;

    let mut input_drvs = Vec::new();
    let mut bracket_count: i32 = 1;
    let mut start = 0;
    pos = 0;
    let bytes = rest.as_bytes();

    while pos < bytes.len() && bracket_count > 0 {
        match bytes[pos] {
            b'[' => bracket_count += 1,
            b']' => bracket_count -= 1,
            b'(' if bracket_count == 1 => {
                start = pos;
            }
            b')' if bracket_count == 1 => {
                let entry = &rest[start + 1..pos];
                if let Some((drv_path, outputs)) = parse_input_drv_entry(entry) {
                    input_drvs.push((drv_path, outputs));
                }
            }
            _ => {}
        }
        pos += 1;
    }

    let rest = &rest[pos..];
    let rest = rest.strip_prefix(",[")?;

    let mut input_srcs = Vec::new();
    bracket_count = 1;
    pos = 0;
    let bytes = rest.as_bytes();

    while pos < bytes.len() && bracket_count > 0 {
        match bytes[pos] {
            b'[' => bracket_count += 1,
            b']' => bracket_count -= 1,
            b'"' if bracket_count == 1 => {
                pos += 1;
                let src_start = pos;
                while pos < bytes.len() && bytes[pos] != b'"' {
                    if bytes[pos] == b'\\' && pos + 1 < bytes.len() {
                        pos += 2;
                    } else {
                        pos += 1;
                    }
                }
                let src = std::str::from_utf8(&bytes[src_start..pos]).ok()?;
                input_srcs.push(src.to_string());
            }
            _ => {}
        }
        pos += 1;
    }

    Some(DerivationInputs {
        input_drvs,
        input_srcs,
    })
}

fn parse_input_drv_entry(entry: &str) -> Option<(String, Vec<String>)> {
    let entry = entry.strip_prefix('"')?;
    let quote_end = entry.find('"')?;
    let drv_path = entry[..quote_end].to_string();

    let rest = &entry[quote_end + 1..];
    let rest = rest.strip_prefix(",[")?;
    let rest = rest.strip_suffix(']')?;

    let mut outputs = Vec::new();
    for part in rest.split(',') {
        let part = part.trim();
        if let Some(name) = part.strip_prefix('"').and_then(|s| s.strip_suffix('"')) {
            outputs.push(name.to_string());
        }
    }

    Some((drv_path, outputs))
}