use pyo3::prelude::*;
use pyo3::types::PyList;
use snail_ast::*;
use snail_error::LowerError;

use super::constants::*;
use super::expr::lower_call_arguments;
use super::helpers::{assign_name, name_expr, string_expr};
use super::py_ast::{AstBuilder, py_err_to_lower};

use super::stmt::{TailBehavior, lower_block_with_tail};

/// Lower a `xargs` statement: iterates files from sources or argv.
pub(crate) fn lower_xargs_stmt(
    builder: &AstBuilder<'_>,
    sources: &[Argument],
    body: &[Stmt],
    span: &SourceSpan,
    tail: TailBehavior,
) -> Result<Vec<PyObject>, LowerError> {
    if sources.is_empty() {
        lower_xargs_no_source(builder, body, span, tail)
    } else {
        lower_xargs_with_sources(builder, sources, body, span, tail)
    }
}

/// Shared loop structure for all `xargs` variants.
///
/// Generates:
/// ```python
/// for __snail_src in <iter_expr>:
///     with __SnailLazyFile(__snail_src, 'r') as __snail_fd:
///         __snail_text = __SnailLazyText(__snail_fd)
///         <body>
/// ```
fn lower_xargs_loop(
    builder: &AstBuilder<'_>,
    iter_expr: PyObject,
    body: &[Stmt],
    span: &SourceSpan,
    tail: TailBehavior,
) -> Result<Vec<PyObject>, LowerError> {
    let mut stmts = Vec::new();

    let with_body = build_xargs_with_body(builder, body, span, tail)?;
    let lazy_file_call = build_lazy_file_call(builder, span)?;

    let with_item = builder
        .call_node_no_loc(
            "withitem",
            vec![
                lazy_file_call,
                name_expr(
                    builder,
                    SNAIL_XARGS_FD_PYVAR,
                    span,
                    builder.store_ctx().map_err(py_err_to_lower)?,
                )?,
            ],
        )
        .map_err(py_err_to_lower)?;

    let with_stmt = builder
        .call_node(
            "With",
            vec![
                PyList::new_bound(builder.py(), vec![with_item]).into_py(builder.py()),
                PyList::new_bound(builder.py(), with_body).into_py(builder.py()),
            ],
            span,
        )
        .map_err(py_err_to_lower)?;

    let for_loop = builder
        .call_node(
            "For",
            vec![
                name_expr(
                    builder,
                    SNAIL_XARGS_SRC_PYVAR,
                    span,
                    builder.store_ctx().map_err(py_err_to_lower)?,
                )?,
                iter_expr,
                PyList::new_bound(builder.py(), vec![with_stmt]).into_py(builder.py()),
                PyList::empty_bound(builder.py()).into_py(builder.py()),
            ],
            span,
        )
        .map_err(py_err_to_lower)?;
    stmts.push(for_loop);
    Ok(stmts)
}

fn lower_xargs_with_sources(
    builder: &AstBuilder<'_>,
    sources: &[Argument],
    body: &[Stmt],
    span: &SourceSpan,
    tail: TailBehavior,
) -> Result<Vec<PyObject>, LowerError> {
    // Lower all arguments via lower_call_arguments and generate:
    //   __snail_normalize_sources(arg1, arg2, *args, key=val, ...)
    let (positional, keywords) = lower_call_arguments(builder, sources, None)?;
    let iter_expr = builder
        .call_node(
            "Call",
            vec![
                name_expr(
                    builder,
                    SNAIL_NORMALIZE_SOURCES_HELPER,
                    span,
                    builder.load_ctx().map_err(py_err_to_lower)?,
                )?,
                PyList::new_bound(builder.py(), positional).into_py(builder.py()),
                PyList::new_bound(builder.py(), keywords).into_py(builder.py()),
            ],
            span,
        )
        .map_err(py_err_to_lower)?;

    lower_xargs_loop(builder, iter_expr, body, span, tail)
}

fn lower_xargs_no_source(
    builder: &AstBuilder<'_>,
    body: &[Stmt],
    span: &SourceSpan,
    tail: TailBehavior,
) -> Result<Vec<PyObject>, LowerError> {
    // for __snail_src in __snail_stdin_args():
    //     with __SnailLazyFile(__snail_src, 'r') as __snail_fd:
    //         __snail_text = __SnailLazyText(__snail_fd)
    //         <body>
    let iter_expr = builder
        .call_node(
            "Call",
            vec![
                name_expr(
                    builder,
                    SNAIL_STDIN_ARGS_HELPER,
                    span,
                    builder.load_ctx().map_err(py_err_to_lower)?,
                )?,
                PyList::empty_bound(builder.py()).into_py(builder.py()),
                PyList::empty_bound(builder.py()).into_py(builder.py()),
            ],
            span,
        )
        .map_err(py_err_to_lower)?;

    lower_xargs_loop(builder, iter_expr, body, span, tail)
}

fn build_xargs_with_body(
    builder: &AstBuilder<'_>,
    body: &[Stmt],
    span: &SourceSpan,
    tail: TailBehavior,
) -> Result<Vec<PyObject>, LowerError> {
    let mut with_body = Vec::new();

    // __snail_text = __SnailLazyText(__snail_fd)
    let lazy_text_call = builder
        .call_node(
            "Call",
            vec![
                name_expr(
                    builder,
                    SNAIL_LAZY_TEXT_CLASS,
                    span,
                    builder.load_ctx().map_err(py_err_to_lower)?,
                )?,
                PyList::new_bound(
                    builder.py(),
                    vec![name_expr(
                        builder,
                        SNAIL_XARGS_FD_PYVAR,
                        span,
                        builder.load_ctx().map_err(py_err_to_lower)?,
                    )?],
                )
                .into_py(builder.py()),
                PyList::empty_bound(builder.py()).into_py(builder.py()),
            ],
            span,
        )
        .map_err(py_err_to_lower)?;
    with_body.push(assign_name(
        builder,
        SNAIL_XARGS_TEXT_PYVAR,
        lazy_text_call,
        span,
    )?);

    // Lower user code
    let user_code = lower_block_with_tail(builder, body, tail, span)?;
    with_body.extend(user_code);

    Ok(with_body)
}

fn build_lazy_file_call(
    builder: &AstBuilder<'_>,
    span: &SourceSpan,
) -> Result<PyObject, LowerError> {
    builder
        .call_node(
            "Call",
            vec![
                name_expr(
                    builder,
                    SNAIL_LAZY_FILE_CLASS,
                    span,
                    builder.load_ctx().map_err(py_err_to_lower)?,
                )?,
                PyList::new_bound(
                    builder.py(),
                    vec![
                        name_expr(
                            builder,
                            SNAIL_XARGS_SRC_PYVAR,
                            span,
                            builder.load_ctx().map_err(py_err_to_lower)?,
                        )?,
                        string_expr(builder, "r", false, StringDelimiter::Double, span)?,
                    ],
                )
                .into_py(builder.py()),
                PyList::empty_bound(builder.py()).into_py(builder.py()),
            ],
            span,
        )
        .map_err(py_err_to_lower)
}
