use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use crate::schema::ROWID_SENTINEL;
use crate::translate::emitter::Resolver;
use crate::translate::expr::{bind_and_rewrite_expr, walk_expr, BindingBehavior, WalkControl};
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},
    Connection,
};
use turso_parser::ast::{self, Expr, Indexed, SortOrder};

use super::emitter::emit_program;
use super::expr::process_returning_clause;
use super::optimizer::optimize_plan;
use super::plan::{
    ColumnUsedMask, IterationDirection, JoinedTable, Plan, TableReferences, UpdatePlan,
};
use super::planner::parse_where;
/*
* 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,
    mut program: ProgramBuilder,
    connection: &Arc<crate::Connection>,
) -> crate::Result<ProgramBuilder> {
    let mut plan = prepare_update_plan(&mut program, resolver.schema, body, connection, false)?;
    optimize_plan(&mut program, &mut plan, resolver.schema)?;
    let opts = ProgramBuilderOpts {
        num_cursors: 1,
        approx_num_insns: 20,
        approx_num_labels: 4,
    };
    program.extend(&opts);
    emit_program(connection, resolver, &mut program, plan, |_| {})?;
    Ok(program)
}

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

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

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

pub fn prepare_update_plan(
    program: &mut ProgramBuilder,
    schema: &Schema,
    mut body: ast::Update,
    connection: &Arc<crate::Connection>,
    is_internal_schema_change: bool,
) -> crate::Result<Plan> {
    if body.with.is_some() {
        bail_parse_error!("WITH clause is not supported in UPDATE");
    }
    if body.or_conflict.is_some() {
        bail_parse_error!("ON CONFLICT clause is not supported in UPDATE");
    }
    if body.from.is_some() {
        bail_parse_error!("FROM clause is not supported in UPDATE");
    }
    if body
        .indexed
        .as_ref()
        .is_some_and(|i| matches!(i, Indexed::IndexedBy(_)))
    {
        bail_parse_error!("INDEXED BY clause is not supported in UPDATE");
    }

    if !body.order_by.is_empty() {
        bail_parse_error!("ORDER BY is not supported in UPDATE");
    }

    let table_name = &body.tbl_name.name;

    // Check if this is a system table that should be protected from direct writes
    // Skip this check for internal schema change operations (like ALTER TABLE)
    if !is_internal_schema_change && crate::schema::is_system_table(table_name.as_str()) {
        bail_parse_error!("table {} may not be modified", table_name);
    }

    if schema.table_has_indexes(&table_name.to_string()) && !schema.indexes_enabled() {
        // Let's disable altering a table with indices altogether instead of checking column by
        // column to be extra safe.
        bail_parse_error!(
            "UPDATE table disabled for table with indexes is disabled. Omit the `--experimental-indexes=false` flag to enable this feature."
        );
    }
    let table = match schema.get_table(table_name.as_str()) {
        Some(table) => table,
        None => bail_parse_error!("Parse error: no such table: {}", table_name),
    };

    // Check if this is a materialized view
    if schema.is_materialized_view(table_name.as_str()) {
        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.as_str());
    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
        );
    }

    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: 0,
    }];
    let mut table_references = TableReferences::new(joined_tables, vec![]);

    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::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,
            connection,
            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) => *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 = expr.clone(),
                                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 = expr.clone(),
                                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)) => *existing_expr = expr.clone(),
                None => set_clauses.push((col_index, expr.clone())),
            }
        }
    }

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

    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),
                connection,
                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,
        connection,
    )?;

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

    // Check what indexes will need to be updated by checking set_clauses and see
    // if a column is contained in an index.
    let indexes = schema.get_indices(table_name);
    let updated_cols: HashSet<usize> = set_clauses.iter().map(|(i, _)| *i).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.cloned().collect()
    } else {
        // otherwise we need to update the indexes whose columns are set in the SET clause,
        // or if the colunns used in the partial index WHERE clause are being updated
        indexes
            .filter_map(|idx| {
                let mut needs = idx.columns.iter().any(|c| {
                    c.expr.as_ref().map_or_else(
                        || updated_cols.contains(&c.pos_in_table),
                        |expr| {
                            columns_used_by_index_expr(expr, &table_references, connection)
                                .iter()
                                .any(|cidx| updated_cols.contains(cidx))
                        },
                    )
                });

                if !needs {
                    if let Some(w) = &idx.where_clause {
                        let mut where_copy = w.as_ref().clone();
                        let mut tr =
                            TableReferences::new(table_references.joined_tables().to_vec(), vec![]);
                        bind_and_rewrite_expr(
                            &mut where_copy,
                            Some(&mut tr),
                            None,
                            connection,
                            BindingBehavior::ResultColumnsNotAllowed,
                        )
                        .ok()?;
                        let cols_used = collect_cols_used_in_expr(&where_copy);
                        // if any of the columns used in the partial index WHERE clause is being
                        // updated, we need to update this index
                        needs = cols_used.iter().any(|c| updated_cols.contains(c));
                    }
                }
                if needs {
                    Some(idx.clone())
                } else {
                    None
                }
            })
            .collect()
    };

    Ok(Plan::Update(UpdatePlan {
        table_references,
        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,
    }))
}

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!(),
    }
}

/// Returns a set of column indices used in the expression.
/// *Must* be used on an Expr already processed by `bind_and_rewrite_expr`
fn collect_cols_used_in_expr(expr: &Expr) -> HashSet<usize> {
    let mut acc = HashSet::new();
    let _ = walk_expr(expr, &mut |expr| match expr {
        Expr::Column { column, .. } => {
            acc.insert(*column);
            Ok(WalkControl::Continue)
        }
        _ => Ok(WalkControl::Continue),
    });
    acc
}

/// Returns a set of column indices used by an index expression.
fn columns_used_by_index_expr(
    expr: &Expr,
    table_references: &TableReferences,
    connection: &Arc<Connection>,
) -> HashSet<usize> {
    let mut expr_copy = expr.clone();
    let mut tr = TableReferences::new(table_references.joined_tables().to_vec(), vec![]);
    if bind_and_rewrite_expr(
        &mut expr_copy,
        Some(&mut tr),
        None,
        connection,
        BindingBehavior::ResultColumnsNotAllowed,
    )
    .is_err()
    {
        return HashSet::new();
    }
    collect_cols_used_in_expr(&expr_copy)
}
