use crate::sync::Arc;
use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};

use crate::schema::{columns_affected_by_update, ROWID_SENTINEL};
use crate::translate::emitter::Resolver;
use crate::translate::expr::{bind_and_rewrite_expr, BindingBehavior};
use crate::translate::expression_index::expression_index_column_usage;
use crate::translate::plan::{Operation, Scan};
use crate::translate::planner::{parse_limit, ROWID_STRS};
use crate::{
    bail_parse_error,
    schema::{Schema, Table},
    util::normalize_ident,
    vdbe::builder::{ProgramBuilder, ProgramBuilderOpts},
    CaptureDataChangesExt, Connection,
};
use turso_parser::ast::{self, Expr, SortOrder};

use super::emitter::emit_program;
use super::expr::process_returning_clause;
use super::optimizer::optimize_plan;
use super::plan::{
    ColumnUsedMask, DmlSafety, IterationDirection, JoinedTable, Plan, TableReferences, UpdatePlan,
};
use super::planner::{parse_where, plan_ctes_as_outer_refs};
use super::subquery::{
    plan_subqueries_from_returning, plan_subqueries_from_select_plan,
    plan_subqueries_from_set_clauses, plan_subqueries_from_where_clause,
};
/*
* Update is simple. By default we scan the table, and for each row, we check the WHERE
* clause. If it evaluates to true, we build the new record with the updated value and insert.
*
* EXAMPLE:
*
sqlite> explain update t set a = 100 where b = 5;
addr  opcode         p1    p2    p3    p4             p5  comment
----  -------------  ----  ----  ----  -------------  --  -------------
0     Init           0     16    0                    0   Start at 16
1     Null           0     1     2                    0   r[1..2]=NULL
2     Noop           1     0     1                    0
3     OpenWrite      0     2     0     3              0   root=2 iDb=0; t
4     Rewind         0     15    0                    0
5       Column         0     1     6                    0   r[6]= cursor 0 column 1
6       Ne             7     14    6     BINARY-8       81  if r[6]!=r[7] goto 14
7       Rowid          0     2     0                    0   r[2]= rowid of 0
8       IsNull         2     15    0                    0   if r[2]==NULL goto 15
9       Integer        100   3     0                    0   r[3]=100
10      Column         0     1     4                    0   r[4]= cursor 0 column 1
11      Column         0     2     5                    0   r[5]= cursor 0 column 2
12      MakeRecord     3     3     1                    0   r[1]=mkrec(r[3..5])
13      Insert         0     1     2     t              7   intkey=r[2] data=r[1]
14    Next           0     5     0                    1
15    Halt           0     0     0                    0
16    Transaction    0     1     1     0              1   usesStmtJournal=0
17    Integer        5     7     0                    0   r[7]=5
18    Goto           0     1     0                    0
*/
pub fn translate_update(
    body: ast::Update,
    resolver: &Resolver,
    program: &mut ProgramBuilder,
    connection: &Arc<crate::Connection>,
) -> crate::Result<()> {
    let mut plan = prepare_update_plan(program, resolver, body, connection, false)?;

    // Plan subqueries in the WHERE clause and SET clause
    if let Plan::Update(ref mut update_plan) = plan {
        if let Some(ref mut ephemeral_plan) = update_plan.ephemeral_plan {
            // When using ephemeral plan (key columns are being updated), subqueries are in the ephemeral_plan's WHERE
            plan_subqueries_from_select_plan(program, ephemeral_plan, resolver, connection)?;
        } else {
            // Normal path: subqueries are in the UPDATE plan's WHERE
            plan_subqueries_from_where_clause(
                program,
                &mut update_plan.non_from_clause_subqueries,
                &mut update_plan.table_references,
                &mut update_plan.where_clause,
                resolver,
                connection,
            )?;
        }
        // Plan subqueries in the SET clause (e.g. UPDATE t SET col = (SELECT ...))
        plan_subqueries_from_set_clauses(
            program,
            &mut update_plan.non_from_clause_subqueries,
            &mut update_plan.table_references,
            &mut update_plan.set_clauses,
            resolver,
            connection,
        )?;
    }

    optimize_plan(program, &mut plan, resolver)?;

    if let Plan::Update(ref update_plan) = plan {
        super::stmt_journal::set_update_stmt_journal_flags(
            program,
            update_plan,
            resolver,
            connection,
        )?;
    }

    let opts = ProgramBuilderOpts {
        num_cursors: 1,
        approx_num_insns: 20,
        approx_num_labels: 4,
    };
    program.extend(&opts);
    emit_program(connection, resolver, program, plan, |_| {})?;
    Ok(())
}

pub fn translate_update_for_schema_change(
    body: ast::Update,
    resolver: &Resolver,
    program: &mut ProgramBuilder,
    connection: &Arc<crate::Connection>,
    ddl_query: &str,
    after: impl FnOnce(&mut ProgramBuilder),
) -> crate::Result<()> {
    let mut plan = prepare_update_plan(program, resolver, body, connection, true)?;

    if let Plan::Update(update_plan) = &mut plan {
        if program.capture_data_changes_info().has_updates() {
            update_plan.cdc_update_alter_statement = Some(ddl_query.to_string());
        }

        // Plan subqueries in the WHERE clause
        if let Some(ref mut ephemeral_plan) = update_plan.ephemeral_plan {
            plan_subqueries_from_select_plan(program, ephemeral_plan, resolver, connection)?;
        } else {
            plan_subqueries_from_where_clause(
                program,
                &mut update_plan.non_from_clause_subqueries,
                &mut update_plan.table_references,
                &mut update_plan.where_clause,
                resolver,
                connection,
            )?;
        }
        // Plan subqueries in the SET clause (e.g. UPDATE t SET col = (SELECT ...))
        plan_subqueries_from_set_clauses(
            program,
            &mut update_plan.non_from_clause_subqueries,
            &mut update_plan.table_references,
            &mut update_plan.set_clauses,
            resolver,
            connection,
        )?;
    }

    optimize_plan(program, &mut plan, resolver)?;
    let opts = ProgramBuilderOpts {
        num_cursors: 1,
        approx_num_insns: 20,
        approx_num_labels: 4,
    };
    program.extend(&opts);
    emit_program(connection, resolver, program, plan, after)?;
    Ok(())
}

fn validate_update(
    schema: &Schema,
    body: &ast::Update,
    table_name: &str,
    is_internal_schema_change: bool,
    conn: &Arc<Connection>,
) -> crate::Result<()> {
    // Check if this is a system table that should be protected from direct writes
    if !is_internal_schema_change
        && !conn.is_nested_stmt()
        && !conn.is_mvcc_bootstrap_connection()
        && !crate::schema::can_write_to_table(table_name)
    {
        crate::bail_parse_error!("table {} may not be modified", table_name);
    }
    if body.from.is_some() {
        bail_parse_error!("FROM clause is not supported in UPDATE");
    }
    if !body.order_by.is_empty() {
        bail_parse_error!("ORDER BY is not supported in UPDATE");
    }
    // Check if this is a materialized view
    if schema.is_materialized_view(table_name) {
        bail_parse_error!("cannot modify materialized view {}", table_name);
    }

    // Check if this table has any incompatible dependent views
    let incompatible_views = schema.has_incompatible_dependent_views(table_name);
    if !incompatible_views.is_empty() {
        use crate::incremental::compiler::DBSP_CIRCUIT_VERSION;
        bail_parse_error!(
            "Cannot UPDATE table '{}' because it has incompatible dependent materialized view(s): {}. \n\
             These views were created with a different DBSP version than the current version ({}). \n\
             Please DROP and recreate the view(s) before modifying this table.",
            table_name,
            incompatible_views.join(", "),
            DBSP_CIRCUIT_VERSION
        );
    }
    Ok(())
}

pub fn prepare_update_plan(
    program: &mut ProgramBuilder,
    resolver: &Resolver,
    mut body: ast::Update,
    connection: &Arc<crate::Connection>,
    is_internal_schema_change: bool,
) -> crate::Result<Plan> {
    let database_id = resolver.resolve_database_id(&body.tbl_name)?;
    let schema = resolver.schema();
    let table_name = &body.tbl_name.name;
    let table = match resolver.with_schema(database_id, |s| s.get_table(table_name.as_str())) {
        Some(table) => table,
        None => bail_parse_error!("Parse error: no such table: {}", table_name),
    };
    if program.trigger.is_some() && table.virtual_table().is_some() {
        bail_parse_error!(
            "unsafe use of virtual table \"{}\"",
            body.tbl_name.name.as_str()
        );
    }
    if crate::is_attached_db(database_id) {
        let schema_cookie = resolver.with_schema(database_id, |s| s.schema_version);
        program.begin_write_on_database(database_id, schema_cookie);
    }
    validate_update(
        schema,
        &body,
        table_name.as_str(),
        is_internal_schema_change,
        connection,
    )?;

    // Extract WITH, OR conflict clause, and INDEXED BY before borrowing body mutably
    let with = body.with.take();
    let or_conflict = body.or_conflict.take();
    let indexed = body.indexed.take();

    let table_name = table.get_name();
    let iter_dir = body
        .order_by
        .first()
        .and_then(|ob| {
            ob.order.map(|o| match o {
                SortOrder::Asc => IterationDirection::Forwards,
                SortOrder::Desc => IterationDirection::Backwards,
            })
        })
        .unwrap_or(IterationDirection::Forwards);

    let joined_tables = vec![JoinedTable {
        table: match table.as_ref() {
            Table::Virtual(vtab) => Table::Virtual(vtab.clone()),
            Table::BTree(btree_table) => Table::BTree(btree_table.clone()),
            _ => unreachable!(),
        },
        identifier: body.tbl_name.alias.as_ref().map_or_else(
            || table_name.to_string(),
            |alias| alias.as_str().to_string(),
        ),
        internal_id: program.table_reference_counter.next(),
        op: build_scan_op(&table, iter_dir),
        join_info: None,
        col_used_mask: ColumnUsedMask::default(),
        column_use_counts: Vec::new(),
        expression_index_usages: Vec::new(),
        database_id,
        indexed,
    }];
    let mut table_references = TableReferences::new(joined_tables, vec![]);

    // Plan CTEs and add them as outer query references for subquery resolution
    plan_ctes_as_outer_refs(with, resolver, program, &mut table_references, connection)?;

    let column_lookup: HashMap<String, usize> = table
        .columns()
        .iter()
        .enumerate()
        .filter_map(|(i, col)| col.name.as_ref().map(|name| (name.to_lowercase(), i)))
        .collect();

    let mut set_clauses: Vec<(usize, Box<Expr>)> = Vec::with_capacity(body.sets.len());

    // Process each SET assignment and map column names to expressions
    // e.g the statement `SET x = 1, y = 2, z = 3` has 3 set assigments
    for set in &mut body.sets {
        bind_and_rewrite_expr(
            &mut set.expr,
            Some(&mut table_references),
            None,
            resolver,
            BindingBehavior::ResultColumnsNotAllowed,
        )?;

        let values = match set.expr.as_ref() {
            Expr::Parenthesized(vals) => vals.clone(),
            expr => vec![expr.clone().into()],
        };

        if set.col_names.len() != values.len() {
            bail_parse_error!(
                "{} columns assigned {} values",
                set.col_names.len(),
                values.len()
            );
        }

        for (col_name, expr) in set.col_names.iter().zip(values.iter()) {
            let ident = normalize_ident(col_name.as_str());

            let col_index = match column_lookup.get(&ident) {
                Some(idx) => {
                    // cannot update generated columns directly
                    table.columns()[*idx].ensure_not_generated("UPDATE", col_name.as_str())?;
                    *idx
                }
                None => {
                    // Check if this is the 'rowid' keyword
                    if ROWID_STRS.iter().any(|s| s.eq_ignore_ascii_case(&ident)) {
                        // Find the rowid alias column if it exists
                        if let Some((idx, _col)) = table
                            .columns()
                            .iter()
                            .enumerate()
                            .find(|(_i, c)| c.is_rowid_alias())
                        {
                            // Use the rowid alias column index
                            match set_clauses.iter_mut().find(|(i, _)| i == &idx) {
                                Some((_, existing_expr)) => existing_expr.clone_from(expr),
                                None => set_clauses.push((idx, expr.clone())),
                            }
                            idx
                        } else {
                            // No rowid alias, use sentinel value for actual rowid
                            match set_clauses.iter_mut().find(|(i, _)| *i == ROWID_SENTINEL) {
                                Some((_, existing_expr)) => existing_expr.clone_from(expr),
                                None => set_clauses.push((ROWID_SENTINEL, expr.clone())),
                            }
                            ROWID_SENTINEL
                        }
                    } else {
                        crate::bail_parse_error!("no such column: {}.{}", table_name, col_name);
                    }
                }
            };
            match set_clauses.iter_mut().find(|(idx, _)| *idx == col_index) {
                Some((_, existing_expr)) => {
                    // When multiple SET col[n] = val for the same column are desugared,
                    // compose them: replace the column reference in the new expression
                    // with the existing expression, so
                    //   col = array_set_element(col, 0, 'X')  then  col = array_set_element(col, 2, 'Z')
                    // becomes col = array_set_element(array_set_element(col, 0, 'X'), 2, 'Z')
                    if let Expr::FunctionCall {
                        name,
                        args: new_args,
                        ..
                    } = expr.as_ref()
                    {
                        if name.as_str().eq_ignore_ascii_case("array_set_element")
                            && new_args.len() == 3
                        {
                            let mut composed_args = new_args.clone();
                            composed_args[0].clone_from(existing_expr);
                            *existing_expr = Box::new(Expr::FunctionCall {
                                name: name.clone(),
                                distinctness: None,
                                args: composed_args,
                                order_by: vec![],
                                filter_over: turso_parser::ast::FunctionTail {
                                    filter_clause: None,
                                    over_clause: None,
                                },
                            });
                        } else {
                            existing_expr.clone_from(expr);
                        }
                    } else {
                        existing_expr.clone_from(expr);
                    }
                }
                None => set_clauses.push((col_index, expr.clone())),
            }
        }
    }

    // Plan subqueries in RETURNING expressions before processing
    // (so SubqueryResult nodes are cloned into result_columns)
    let mut non_from_clause_subqueries = vec![];
    plan_subqueries_from_returning(
        program,
        &mut non_from_clause_subqueries,
        &mut table_references,
        &mut body.returning,
        resolver,
        connection,
    )?;

    let result_columns =
        process_returning_clause(&mut body.returning, &mut table_references, resolver)?;

    let order_by = body
        .order_by
        .iter_mut()
        .map(|o| {
            let _ = bind_and_rewrite_expr(
                &mut o.expr,
                Some(&mut table_references),
                Some(&result_columns),
                resolver,
                BindingBehavior::ResultColumnsNotAllowed,
            );
            (o.expr.clone(), o.order.unwrap_or(SortOrder::Asc))
        })
        .collect();

    // Sqlite determines we should create an ephemeral table if we do not have a FROM clause
    // Difficult to say what items from the plan can be checked for this so currently just checking if a RowId Alias is referenced
    // https://github.com/sqlite/sqlite/blob/master/src/update.c#L395
    // https://github.com/sqlite/sqlite/blob/master/src/update.c#L670
    let columns = table.columns();
    let mut where_clause = vec![];
    // Parse the WHERE clause
    parse_where(
        body.where_clause.as_deref(),
        &mut table_references,
        Some(&result_columns),
        &mut where_clause,
        resolver,
    )?;

    // Parse the LIMIT/OFFSET clause
    let (limit, offset) = body
        .limit
        .map_or(Ok((None, None)), |l| parse_limit(l, resolver))?;

    // Determine which indexes need updating
    let indexes: Vec<_> = resolver.with_schema(database_id, |s| {
        s.get_indices(table_name).cloned().collect()
    });
    let rowid_alias_used = set_clauses
        .iter()
        .any(|(idx, _)| *idx == ROWID_SENTINEL || columns[*idx].is_rowid_alias());
    let indexes_to_update = if rowid_alias_used {
        // If the rowid alias is used in the SET clause, we need to update all indexes
        indexes
    } else {
        let updated_cols: HashSet<usize> = set_clauses.iter().map(|(i, _)| *i).collect();
        let affected_cols = columns_affected_by_update(columns, &updated_cols);
        let target_table_ref = table_references
            .joined_tables()
            .first()
            .expect("UPDATE must have a target table reference");
        let mut indexes_to_update = Vec::new();

        // else, we update indexes that match certain conditions
        for idx in indexes {
            let mut needs = false;
            for col in idx.columns.iter() {
                if let Some(expr) = col.expr.as_ref() {
                    let cols_used =
                        expression_index_column_usage(expr.as_ref(), target_table_ref, resolver)?;

                    if cols_used.iter().any(|cidx| affected_cols.contains(&cidx)) {
                        needs = true;
                        break;
                    }
                } else if affected_cols.contains(&col.pos_in_table) {
                    needs = true;
                    break;
                }
            }

            if !needs {
                if let Some(where_expr) = &idx.where_clause {
                    let cols_used = expression_index_column_usage(
                        where_expr.as_ref(),
                        target_table_ref,
                        resolver,
                    )?;
                    // If any column used in the partial index WHERE clause is affected,
                    // this index must be updated as well.
                    needs = cols_used.iter().any(|cidx| affected_cols.contains(&cidx));
                }
            }

            if needs {
                indexes_to_update.push(idx);
            }
        }
        indexes_to_update
    };

    Ok(Plan::Update(UpdatePlan {
        table_references,
        or_conflict,
        set_clauses,
        where_clause,
        returning: if result_columns.is_empty() {
            None
        } else {
            Some(result_columns)
        },
        order_by,
        limit,
        offset,
        contains_constant_false_condition: false,
        indexes_to_update,
        ephemeral_plan: None,
        cdc_update_alter_statement: None,
        non_from_clause_subqueries,
        safety: DmlSafety::default(),
    }))
}

fn build_scan_op(table: &Table, iter_dir: IterationDirection) -> Operation {
    match table {
        Table::BTree(_) => Operation::Scan(Scan::BTreeTable {
            iter_dir,
            index: None,
        }),
        Table::Virtual(_) => Operation::default_scan_for(table),
        _ => unreachable!(),
    }
}
