use std::path::PathBuf;

use clap::Args;
use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
use obilayeredmap::IndexMode;

fn parse_key_value(s: &str) -> Result<(String, String), String> {
    let pos = s.find('=').ok_or_else(|| format!("invalid key=value: no '=' in '{s}'"))?;
    Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
}
use obisys::Reporter;
use tracing::info;

use crate::cli::{CommonArgs, block_size_to_bits, partitions_to_bits};
use crate::steps::scatter;

#[derive(Args)]
pub struct IndexArgs {
    /// Output index directory
    #[arg(short, long)]
    pub output: PathBuf,

    /// Overwrite output directory if it already exists
    #[arg(long, default_value_t = false)]
    pub force: bool,

    /// Genome label (default: input filename without path/extension)
    #[arg(long)]
    pub label: Option<String>,

    /// Genome categorical metadata as key=value pairs (repeatable)
    #[arg(long = "meta", value_parser = parse_key_value)]
    pub meta: Vec<(String, String)>,

    /// Minimum kmer abundance (inclusive)
    #[arg(long, default_value_t = 1)]
    pub min_abundance: u32,

    /// Maximum kmer abundance (inclusive)
    #[arg(long)]
    pub max_abundance: Option<u32>,

    /// Store kmer counts in the index (default: set membership only)
    #[arg(long, default_value_t = false)]
    pub with_counts: bool,

    /// Keep intermediate build files (dereplicated superkmers, mphf1, counts1)
    #[arg(long, default_value_t = false)]
    pub keep_intermediate: bool,

    /// Use approximate (fingerprint-based) evidence instead of exact evidence.
    /// False-positive rate per z-window: 1/2^(b·z).
    #[arg(long, default_value_t = false)]
    pub approx: bool,

    /// Findere z parameter: number of consecutive k-mers that must all match.
    /// Effective indexed k-mer size is kmer_size - z + 1.
    #[arg(short = 'z', long, default_value = None)]
    pub findere_z: Option<u8>,

    /// Fingerprint bits per slot (b).  FP per z-window = 1/2^(b·z).
    #[arg(long, default_value = None)]
    pub evidence_bits: Option<u8>,

    /// Target false-positive rate per z-window (e.g. 0.01).
    /// Used to derive missing b or z.
    #[arg(long, default_value = None)]
    pub fp: Option<f64>,

    /// Block size for exact evidence `.idx` (number of unitigs per block).
    /// Must be a power of two; rounded up if not. Default 1 = O(1) random access.
    #[arg(long, default_value_t = 1)]
    pub block_size: usize,

    #[command(flatten)]
    pub common: CommonArgs,
}

/// Resolve the (z, b, fp) triplet from the user-supplied subset.
///
/// Model: FP = 1/2^(b·z)  ⟹  b·z = ⌈-log₂(fp)⌉
///
/// Rules when one value is missing (conservative = ceiling):
///   given z, b       → fp = 1/2^(b·z)
///   given z, fp      → b  = ⌈-log₂(fp) / z⌉
///   given b, fp      → z  = ⌈-log₂(fp) / b⌉
///   given z only     → b  = 8  (default),  fp derived
///   given b only     → z  = 1  (default),  fp derived
///   given fp only    → b  = 8  (default),  z  derived
///   none given       → z  = 1, b = 8, fp = 1/256
pub(crate) fn resolve_approx_params(
    z_opt: Option<u8>,
    b_opt: Option<u8>,
    fp_opt: Option<f64>,
) -> (u8, u8, f64) {
    const DEFAULT_B: u8 = 8;
    const DEFAULT_Z: u8 = 1;

    let bits_needed = |fp: f64| -> u8 {
        (-fp.log2()).ceil() as u8
    };

    match (z_opt, b_opt, fp_opt) {
        // All three given: use b and z, recompute fp conservatively.
        (Some(z), Some(b), Some(_fp)) => {
            let fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, fp)
        }
        // Two given, derive third.
        (Some(z), Some(b), None) => {
            let fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, fp)
        }
        (Some(z), None, Some(fp)) => {
            let bz = (-fp.log2()).ceil() as u32;
            let b = ((bz + z as u32 - 1) / z as u32).max(1) as u8;
            let actual_fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, actual_fp)
        }
        (None, Some(b), Some(fp)) => {
            let bz = (-fp.log2()).ceil() as u32;
            let z = ((bz + b as u32 - 1) / b as u32).max(1) as u8;
            let actual_fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, actual_fp)
        }
        // One given, apply defaults for the other.
        (Some(z), None, None) => {
            let b = DEFAULT_B;
            let fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, fp)
        }
        (None, Some(b), None) => {
            let z = DEFAULT_Z;
            let fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, fp)
        }
        (None, None, Some(fp)) => {
            let b = DEFAULT_B;
            let z = ((bits_needed(fp) as u32 + b as u32 - 1) / b as u32).max(1) as u8;
            let actual_fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, actual_fp)
        }
        // None given: defaults.
        (None, None, None) => {
            let b = DEFAULT_B;
            let z = DEFAULT_Z;
            let fp = 1.0_f64 / (1u64 << (b as u32 * z as u32)) as f64;
            (z, b, fp)
        }
    }
}

pub fn run(args: IndexArgs) {
    args.common.validate();

    let output = args.output.clone();
    let mut rep = Reporter::new();

    // ── Resolve evidence kind ────────────────────────────────────────────────
    let (evidence, effective_kmer_size) = if args.approx {
        let (z, b, fp) = resolve_approx_params(args.findere_z, args.evidence_bits, args.fp);
        let k = args.common.kmer_size;
        if z as usize >= k {
            eprintln!(
                "error: Findere z={z} must be < kmer-size={k} \
                 (effective kmer size k−z+1 = {} ≤ 0)",
                k as isize - z as isize + 1
            );
            std::process::exit(1);
        }
        let s = k - z as usize + 1;
        info!("approximate evidence: b={b}, z={z}, fp={fp:.2e}, indexed kmer size={s}");
        (IndexMode::Approx { b, z }, s)
    } else {
        (IndexMode::Exact, args.common.kmer_size)
    };

    // ── Open or create the index ─────────────────────────────────────────────
    let mut idx = if KmerIndex::exists(&output) && !args.force {
        info!("resuming from existing index at {}", output.display());
        KmerIndex::open(&output).unwrap_or_else(|e| {
            eprintln!("error opening index: {e}");
            std::process::exit(1);
        })
    } else {
        if args.force && KmerIndex::exists(&output) {
            info!("--force: removing existing index at {}", output.display());
            std::fs::remove_dir_all(&output).unwrap_or_else(|e| {
                eprintln!("error removing existing index: {e}");
                std::process::exit(1);
            });
        }
        let n_bits = partitions_to_bits(args.common.partitions);
        let effective = 1usize << n_bits;
        if effective != args.common.partitions {
            info!("partitions: {} → {} (next power of 2)", args.common.partitions, effective);
        }
        let block_bits = block_size_to_bits(args.block_size);
        let config = IndexConfig {
            kmer_size:      effective_kmer_size,
            minimizer_size: args.common.minimizer_size,
            n_bits,
            with_counts:    args.with_counts,
            evidence:       evidence.clone(),
            block_bits,
        };
        let genome_info = args.label.as_ref().map(|label| {
            validate_label(label).unwrap_or_else(|e| {
                eprintln!("error: --label: {e}");
                std::process::exit(1);
            });
            let mut info = GenomeInfo::new(label.clone());
            for (k, v) in &args.meta {
                info.meta.insert(k.clone(), v.clone());
            }
            info
        });
        KmerIndex::create(&output, config, genome_info, false).unwrap_or_else(|e| {
            eprintln!("error creating index: {e}");
            std::process::exit(1);
        })
    };


    // ── Stage 1: scatter ─────────────────────────────────────────────────────
    if idx.state() < IndexState::Scattered {
        let k         = idx.kmer_size();
        let level_max = args.common.level_max;
        let theta     = args.common.theta;
        let n_workers = args.common.threads.max(1);

        let max_open = args.common.effective_max_open();
        scatter(idx.partition_mut(), args.common.seqfile_paths(), k, level_max, theta, n_workers, max_open, &mut rep);

        idx.mark_scattered().unwrap_or_else(|e| {
            eprintln!("error marking scatter done: {e}");
            std::process::exit(1);
        });
    } else {
        info!("scatter already done, skipping");
    }

    // ── Stage 2: dereplicate + count ─────────────────────────────────────────
    if idx.state() < IndexState::Counted {
        idx.dereplicate_and_count(args.keep_intermediate, &mut rep).unwrap_or_else(|e| {
            eprintln!("error: {e}");
            std::process::exit(1);
        });
    } else {
        info!("dereplicate+count already done, skipping");
    }

    // ── Stage 3: build layered index ─────────────────────────────────────────
    if idx.state() < IndexState::Indexed {
        idx.build_layers(
            args.min_abundance,
            args.max_abundance,
            args.keep_intermediate,
            &mut rep,
        ).unwrap_or_else(|e| {
            eprintln!("error: {e}");
            std::process::exit(1);
        });
    } else {
        info!("index already built, skipping");
    }

    rep.print();
}
