use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use tracing::debug;
use obipipeline::{
    Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, WorkerPool,
    make_sink, make_source, make_transform,
    throttle,
};

use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, Layer};
use obiskio::{SKError, SKResult};

use crate::common::olm_to_sk;

// ── KmerGraphData ─────────────────────────────────────────────────────────────

enum KmerGraphData {
    File(obipipeline::Throttled<PathBuf>),
    RawBatch(Vec<CanonicalKmer>),
    FilteredBatch(Vec<CanonicalKmer>),
}

// ── build_graph ───────────────────────────────────────────────────────────────

/// Phase 1: pipeline that reads files, filters kmers, pushes into a GraphDeBruijn.
///
/// `flat_fn(path, emit)`: opens path, iterates kmers, calls `emit(batch)` for each batch.
/// `filter(kmer) -> bool`: secondary filter applied in the Transform stage.
pub(crate) fn build_graph<I, F, G>(
    file_source: I,
    flat_fn: F,
    filter: G,
    n_workers: usize,
    max_open: usize,
) -> SKResult<GraphDeBruijn>
where
    I: Iterator<Item = PathBuf> + Send + 'static,
    F: Fn(PathBuf, &mut dyn FnMut(Vec<CanonicalKmer>)) -> SKResult<()> + Send + Sync + 'static,
    G: Fn(CanonicalKmer) -> bool + Send + Sync + 'static,
{
    let capacity = 2;

    let flat_fn = Arc::new(flat_fn);
    let filter = Arc::new(filter);
    let g_shared = Arc::new(Mutex::new(GraphDeBruijn::new()));
    let g_sink = Arc::clone(&g_shared);
    let err_cap: Arc<Mutex<Option<SKError>>> = Arc::new(Mutex::new(None));
    let err_flat = Arc::clone(&err_cap);

    let throttled = throttle(file_source, max_open);

    let pipeline = Pipeline::new(
        make_source!(KmerGraphData, throttled, File),
        vec![
            Stage::Flat(Arc::new(
                move |data: KmerGraphData,
                      push: &PipelineSender<Result<KmerGraphData, PipelineError>>,
                      delta: &PipelineSender<isize>|
                {
                    if let KmerGraphData::File(t) = data {
                        let path = t.item;
                        let _guard = t.guard; // released at end of block
                        let mut count: isize = 0;
                        let push_clone = push.clone();
                        let result = flat_fn(path, &mut |batch: Vec<CanonicalKmer>| {
                            push_clone.send(Ok(KmerGraphData::RawBatch(batch))).ok();
                            count += 1;
                        });
                        match result {
                            Ok(()) => {
                                delta.send(count - 1).ok();
                            }
                            Err(e) => {
                                *err_flat.lock().unwrap() = Some(e);
                                delta.send(-1).ok();
                            }
                        }
                    }
                },
            ) as SharedFlatFn<KmerGraphData>),
            make_transform!(KmerGraphData, {
                let filter = Arc::clone(&filter);
                move |batch: Vec<CanonicalKmer>| -> Vec<CanonicalKmer> {
                    batch.into_iter().filter(|k| filter(*k)).collect()
                }
            }, RawBatch, FilteredBatch),
        ],
        make_sink!(KmerGraphData, {
            move |batch: Vec<CanonicalKmer>| {
                let mut g = g_sink.lock().unwrap();
                for kmer in batch {
                    g.push(kmer);
                }
            }
        }, FilteredBatch),
    );

    WorkerPool::new(pipeline, n_workers, capacity).run();

    if let Some(e) = Arc::try_unwrap(err_cap)
        .unwrap_or_else(|_| panic!("build_graph: err_cap not uniquely owned after pipeline"))
        .into_inner()
        .unwrap_or_else(|e| e.into_inner())
    {
        return Err(e);
    }

    let g = Arc::try_unwrap(g_shared)
        .unwrap_or_else(|_| panic!("build_graph: g_shared not uniquely owned after pipeline"))
        .into_inner()
        .unwrap_or_else(|e| e.into_inner());

    Ok(g)
}

// ── write_graph_as_unitigs ────────────────────────────────────────────────────

/// Phase 2 (write unitigs only): compute degrees, write unitigs to `layer_dir`, drop graph.
///
/// Returns n_kmers. Does NOT build the MPHF — caller does it.
pub(crate) fn write_graph_as_unitigs(g: GraphDeBruijn, layer_dir: &Path) -> SKResult<usize> {
    let n_kmers = g.len();
    g.compute_degrees_and_mark_starts();
    std::fs::create_dir_all(layer_dir)?;
    let mut uw = Layer::<()>::unitig_writer(layer_dir).map_err(|e| olm_to_sk(e, "graph pipeline"))?;
    g.try_for_each_unitig(|unitig| uw.write(unitig))?;
    uw.close()?;
    drop(g);
    Ok(n_kmers)
}

// ── materialize_layer ─────────────────────────────────────────────────────────

/// Phase 2 (full): write_graph_as_unitigs + `Layer::<()>::build`.
///
/// Returns n_kmers.
pub(crate) fn materialize_layer(
    g: GraphDeBruijn,
    layer_dir: &Path,
    block_bits: u8,
    evidence: &IndexMode,
) -> SKResult<usize> {
    let n = write_graph_as_unitigs(g, layer_dir)?;
    debug!("materialize_layer: unitigs written ({n} kmers), building MPHF");
    Layer::<()>::build(layer_dir, block_bits, evidence)
        .map_err(|e| olm_to_sk(e, "graph pipeline"))?;
    debug!("materialize_layer: MPHF build done");
    Ok(n)
}
