use std::path::PathBuf;

use eyre::{Context, Result};
use itertools::Itertools;
use kernels_data::config::TorchNoarch;
use kernels_data::config::{Backend, Build, General};
use minijinja::{context, Environment};

use crate::pyproject::common::write_compat_py;
use crate::pyproject::common::write_metadata;
use crate::pyproject::fileset::FileSet;
use crate::pyproject::ops_identifier::KernelIdentifier;

static SETUP_PY: &str = include_str!("../templates/torch/noarch/setup.py");

pub fn write_torch_ext_noarch(
    env: &Environment,
    build: &Build,
    kernel_id: &KernelIdentifier,
) -> Result<FileSet> {
    let mut file_set = FileSet::default();

    write_compat_py(&mut file_set)?;
    write_ops_py(
        env,
        &build.general.name.python_name(),
        kernel_id,
        &mut file_set,
    )?;
    write_pyproject_toml(
        env,
        build.framework.torch_noarch(),
        &build.general,
        &mut file_set,
    )?;
    write_setup_py(&mut file_set)?;
    write_metadata(build, kernel_id, &mut file_set)?;

    Ok(file_set)
}

fn write_ops_py(
    env: &Environment,
    name: &str,
    kernel_id: &KernelIdentifier,
    file_set: &mut FileSet,
) -> Result<()> {
    let mut path = PathBuf::new();
    path.push("torch-ext");
    path.push(name);
    path.push("_ops.py");
    let writer = file_set.entry(path);

    env.get_template("torch/noarch/_ops.py")
        .wrap_err("Cannot get noarch _ops.py template")?
        .render_captured_to(
            context! {
                kernel_name => kernel_id.name(),
                kernel_unique_id => kernel_id.unique_id(),
            },
            writer,
        )
        .wrap_err("Cannot render kernel template")?;

    Ok(())
}

fn write_pyproject_toml(
    env: &Environment,
    torch_noarch: Option<&TorchNoarch>,
    general: &General,
    file_set: &mut FileSet,
) -> Result<()> {
    let writer = file_set.entry("pyproject.toml");

    let name = &general.name;
    let data_globs = torch_noarch.and_then(|torch| {
        torch
            .data_extensions()
            .map(|exts| exts.iter().map(|ext| format!("\"**/*.{ext}\"")).join(", "))
    });

    // Common python dependencies (no backend-specific ones)
    let python_dependencies =
        itertools::process_results(general.general_python_depends(), |iter| {
            iter.flat_map(|(_, deps)| deps.python.iter().map(|d| format!("\"{}\"", d.pkg)))
                .join(", ")
        })?;

    // Collect backend-specific dependencies for all backends
    let mut backend_dependencies = Vec::new();
    for backend in &Backend::all() {
        let deps = itertools::process_results(general.backend_python_depends(*backend), |iter| {
            iter.flat_map(|(_, deps)| deps.python.iter().map(|d| format!("\"{}\"", d.pkg)))
                .join(", ")
        })?;

        if !deps.is_empty() {
            backend_dependencies.push((backend.to_string(), deps));
        }
    }

    env.get_template("torch/noarch/pyproject.toml")
        .wrap_err("Cannot get noarch pyproject.toml template")?
        .render_captured_to(
            context! {
                data_globs => data_globs,
                python_dependencies => python_dependencies,
                backend_dependencies => backend_dependencies,
                name => name,
                python_name => general.name.python_name(),
            },
            writer,
        )
        .wrap_err("Cannot render kernel template")?;

    Ok(())
}

fn write_setup_py(file_set: &mut FileSet) -> Result<()> {
    file_set
        .entry("setup.py")
        .extend_from_slice(SETUP_PY.as_bytes());
    Ok(())
}
