use crate::backend::{Backend, python::PythonBackend};
use crate::opt::const_fold::const_fold;
use crate::{
    Diag, DiagConfig, lower_with_config_collect, optimize_with_info_and_config,
    parse_dsl_with_errors, render_simplified_expr,
};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ApiError {
    #[error("parse failed: {0}")]
    Parse(String),
    #[error("lowering failed: {0}")]
    Lower(String),
    #[error("optimization failed: {0}")]
    Optimize(String),
    #[error("expected a single output expression, got {0}")]
    UnsupportedOutputs(usize),
}

fn format_diags(diags: &[Diag]) -> String {
    diags
        .iter()
        .map(|diag| diag.message.clone())
        .collect::<Vec<_>>()
        .join("; ")
}

fn lowered_dag(source: &str, optimize: bool) -> Result<crate::ir::TypedDag, ApiError> {
    let func =
        parse_dsl_with_errors(source).map_err(|diags| ApiError::Parse(format_diags(&diags)))?;
    let lower_result = lower_with_config_collect(&func, DiagConfig::default());
    if !lower_result.errors.is_empty() {
        return Err(ApiError::Lower(format_diags(&lower_result.errors)));
    }

    let mut dag = lower_result.dag;
    if optimize {
        dag = optimize_with_info_and_config(&dag, &Default::default())
            .map_err(|err| ApiError::Optimize(err.to_string()))?
            .dag;
    }
    Ok(dag)
}

pub fn simplify_expr(source: &str, optimize: bool) -> Result<String, ApiError> {
    let dag = const_fold(&lowered_dag(source, optimize)?).dag;
    if dag.outputs.len() != 1 {
        return Err(ApiError::UnsupportedOutputs(dag.outputs.len()));
    }
    Ok(render_simplified_expr(&dag, dag.outputs[0], true))
}

pub fn emit_python(source: &str, function_name: &str, optimize: bool) -> Result<String, ApiError> {
    let dag = lowered_dag(source, optimize)?;
    let backend = PythonBackend;
    Ok(backend.emit_kernel(&dag, function_name))
}
