use std::collections::{BTreeSet, HashSet};
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

fn main() {
    println!("cargo:rerun-if-env-changed=GHC_LIBDIR");
    println!("cargo:rerun-if-env-changed=GHC_RTS_LIB");
    println!("cargo:rerun-if-changed=rust/build.rs");

    let libdir = env::var("GHC_LIBDIR").unwrap_or_else(|_| ghc_print_libdir());
    let explicit_rts = env::var("GHC_RTS_LIB").ok().map(PathBuf::from);
    let rts_path = explicit_rts.unwrap_or_else(|| find_rts_library(Path::new(&libdir)));
    let rts_dir = rts_path
        .parent()
        .unwrap_or_else(|| Path::new(&libdir))
        .to_path_buf();
    let rts_name = rts_path
        .file_stem()
        .and_then(|stem| stem.to_str())
        .map(strip_library_prefix)
        .unwrap_or_else(|| panic!("failed to resolve GHC RTS library name from {}", rts_path.display()));

    let mut search_dirs = BTreeSet::new();
    let mut haskell_libs = Vec::new();
    let mut seen_haskell_libs = HashSet::new();
    let mut native_libs = BTreeSet::new();

    search_dirs.insert(rts_dir);
    seen_haskell_libs.insert(rts_name.clone());
    haskell_libs.push(rts_name);

    for package in ["base", "ghc-prim", "ghc-bignum"] {
        let info = ghc_pkg_describe(package);
        for dir in &info.dynamic_library_dirs {
            search_dirs.insert(dir.clone());
        }
        for library in info.hs_libraries {
            let resolved = resolve_dynamic_hs_library(&library, &info.dynamic_library_dirs);
            if seen_haskell_libs.insert(resolved.clone()) {
                haskell_libs.push(resolved);
            }
        }

        for library in info.extra_libraries {
            native_libs.insert(library);
        }
    }

    let rts_info = ghc_pkg_describe("rts");
    for dir in rts_info.dynamic_library_dirs {
        search_dirs.insert(dir);
    }
    for library in rts_info.extra_libraries {
        native_libs.insert(library);
    }

    for dir in search_dirs {
        println!("cargo:rustc-link-search=native={}", dir.display());
        println!("cargo:rustc-link-arg=-Wl,-rpath,{}", dir.display());
    }

    println!("cargo:rustc-link-arg=-Wl,--no-as-needed");
    for library in haskell_libs {
        println!("cargo:rustc-link-lib=dylib={library}");
    }
    println!("cargo:rustc-link-arg=-Wl,--as-needed");

    for library in native_libs {
        println!("cargo:rustc-link-lib=dylib={library}");
    }
}

#[derive(Default)]
struct PackageInfo {
    hs_libraries: Vec<String>,
    extra_libraries: Vec<String>,
    dynamic_library_dirs: Vec<PathBuf>,
}

fn ghc_print_libdir() -> String {
    let output = Command::new("ghc")
        .arg("--print-libdir")
        .output()
        .unwrap_or_else(|error| panic!("failed to run `ghc --print-libdir`: {error}"));
    if !output.status.success() {
        panic!("`ghc --print-libdir` did not exit successfully");
    }

    String::from_utf8_lossy(&output.stdout).trim().to_string()
}

fn ghc_pkg_describe(package: &str) -> PackageInfo {
    let output = Command::new("ghc-pkg")
        .args(["describe", package])
        .output()
        .unwrap_or_else(|error| panic!("failed to run `ghc-pkg describe {package}`: {error}"));
    if !output.status.success() {
        panic!("`ghc-pkg describe {package}` did not exit successfully");
    }

    let description = String::from_utf8_lossy(&output.stdout);
    parse_package_description(&description)
}

fn parse_package_description(description: &str) -> PackageInfo {
    let mut info = PackageInfo::default();
    let mut current_field = String::new();
    let mut pkgroot = String::new();

    for raw_line in description.lines() {
        let line = raw_line.trim_end();
        if line.is_empty() {
            continue;
        }

        if raw_line.starts_with(' ') || raw_line.starts_with('\t') {
            push_field_values(&current_field, line.trim(), &mut pkgroot, &mut info);
            continue;
        }

        if let Some((field, rest)) = line.split_once(':') {
            current_field = field.trim().to_string();
            push_field_values(&current_field, rest.trim(), &mut pkgroot, &mut info);
        }
    }

    if !pkgroot.is_empty() {
        for dir in &mut info.dynamic_library_dirs {
            let resolved = dir
                .display()
                .to_string()
                .replace("${pkgroot}", &pkgroot);
            *dir = PathBuf::from(resolved);
        }
    }

    info
}

fn push_field_values(
    field: &str,
    values: &str,
    pkgroot: &mut String,
    info: &mut PackageInfo,
) {
    if values.is_empty() {
        return;
    }

    match field {
        "pkgroot" => {
            *pkgroot = values.trim_matches('"').to_string();
        }
        "hs-libraries" => {
            for value in values.split_whitespace() {
                info.hs_libraries
                    .push(strip_library_prefix(value.trim_matches('"')));
            }
        }
        "extra-libraries" => {
            for value in values.split_whitespace() {
                info.extra_libraries.push(value.trim_matches('"').to_string());
            }
        }
        "dynamic-library-dirs" => {
            for value in values.split_whitespace() {
                info.dynamic_library_dirs
                    .push(PathBuf::from(value.trim_matches('"')));
            }
        }
        _ => {}
    }
}

fn resolve_dynamic_hs_library(library: &str, search_dirs: &[PathBuf]) -> String {
    for dir in search_dirs {
        let Ok(entries) = fs::read_dir(dir) else {
            continue;
        };

        for entry in entries.flatten() {
            let path = entry.path();
            let Some(file_name) = path.file_name().and_then(|value| value.to_str()) else {
                continue;
            };

            let exact_name = format!("lib{library}.so");
            let versioned_prefix = format!("lib{library}-");
            if (file_name == exact_name
                || (file_name.starts_with(&versioned_prefix) && file_name.ends_with(".so")))
                && path.is_file()
            {
                if let Some(stem) = path.file_stem().and_then(|value| value.to_str()) {
                    return strip_library_prefix(stem);
                }
            }
        }
    }

    library.to_string()
}

fn find_rts_library(libdir: &Path) -> PathBuf {
    let mut candidates = Vec::new();
    walk_for_rts(libdir, &mut candidates);
    candidates.sort_by_key(|path| rts_priority(path));

    candidates
        .into_iter()
        .next()
        .unwrap_or_else(|| panic!("failed to locate a GHC RTS library under {}", libdir.display()))
}

fn walk_for_rts(root: &Path, candidates: &mut Vec<PathBuf>) {
    let Ok(entries) = fs::read_dir(root) else {
        return;
    };

    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            walk_for_rts(&path, candidates);
            continue;
        }

        if is_threaded_rts_library(&path) {
            candidates.push(path);
        }
    }
}

fn is_threaded_rts_library(path: &Path) -> bool {
    let Some(file_name) = path.file_name().and_then(|value| value.to_str()) else {
        return false;
    };

    path.extension().and_then(|ext| ext.to_str()) == Some("so") && file_name.starts_with("libHSrts-")
}

fn rts_priority(path: &Path) -> (u8, String) {
    let file_name = path
        .file_name()
        .and_then(|value| value.to_str())
        .unwrap_or_default()
        .to_string();

    let rank = if file_name.contains("_debug") {
        3
    } else if file_name.contains("_thr") {
        1
    } else {
        0
    };

    (rank, file_name)
}

fn strip_library_prefix(stem: &str) -> String {
    stem.strip_prefix("lib").unwrap_or(stem).to_string()
}