import random
from typing import Callable

import mutator_corpus_splicing
import mutator_equivalent_rewrites
import mutator_extra_statements
import mutator_operators
import mutator_values

_MUTATORS_WEIGHTED: list[tuple[Callable[[str], str], float]] = [
  # # ##############################################################################################
  (mutator_values.mut_replace_num_with_edge_int, 3),
  (mutator_values.mut_replace_num_with_random_int, 2),
  (mutator_values.mut_replace_num_with_edge_float, 2),
  (mutator_values.mut_replace_num_with_edge_string, 2),
  (mutator_values.mut_replace_num_with_random_hex, 2),
  (mutator_values.mut_replace_num_with_null, 2),
  (mutator_values.mut_wrap_num_in_cast, 3),
  (mutator_values.mut_wrap_num_in_datetime_func, 2),
  (mutator_values.mut_wrap_num_in_core_func, 10),
  (mutator_values.mut_wrap_num_in_math_func, 5),
  (mutator_values.mut_wrap_num_in_json_func, 5),
  # # ##############################################################################################
  (mutator_operators.mut_swap_comparison, 2),
  (mutator_operators.mut_swap_boolean, 2),
  (mutator_operators.mut_negate_where, 2),
  (mutator_operators.mut_change_type, 2),
  (mutator_operators.mut_swap_join_type, 2),
  # # ##############################################################################################
  (mutator_extra_statements.mut_append_select_star, 6),
  (mutator_extra_statements.mut_append_count, 6),
  (mutator_extra_statements.mut_append_join, 7),
  (mutator_extra_statements.mut_append_subquery, 9),
  (mutator_extra_statements.mut_append_cte, 9),
  (mutator_extra_statements.mut_append_window, 7),
  (mutator_extra_statements.mut_append_filter, 6),
  (mutator_extra_statements.mut_append_aggregate, 7),
  (mutator_extra_statements.mut_append_update, 6),
  (mutator_extra_statements.mut_append_delete, 6),
  (mutator_extra_statements.mut_append_index, 6),
  (mutator_extra_statements.mut_append_view, 6),
  (mutator_extra_statements.mut_append_trigger, 6),
  (mutator_extra_statements.mut_wrap_savepoint, 6),
  (mutator_extra_statements.mut_wrap_transaction, 6),
  (mutator_extra_statements.mut_wrap_attach_detach, 5),
  (mutator_extra_statements.mut_append_vacuum, 5),
  (mutator_extra_statements.mut_append_analyze, 6),
  (mutator_extra_statements.mut_append_reindex, 5),
  (mutator_extra_statements.mut_append_null_insert, 6),
  (mutator_extra_statements.mut_append_insert_default, 6),
  (mutator_extra_statements.mut_append_insert_select, 6),
  (mutator_extra_statements.mut_append_insert_or_x, 6),
  (mutator_extra_statements.mut_append_upsert, 6),
  (mutator_extra_statements.mut_append_alter_rename_table, 5),
  (mutator_extra_statements.mut_append_alter_rename_column, 5),
  (mutator_extra_statements.mut_append_alter_add_column, 6),
  (mutator_extra_statements.mut_append_alter_drop_column, 5),
  (mutator_extra_statements.mut_append_alter_drop_table_index_view_trigger, 5),
  (mutator_extra_statements.mut_repeat, 6),
  (mutator_extra_statements.mut_prepend_random_pragma, 9),
  (mutator_extra_statements.mut_append_random_pragma, 6),
  (mutator_extra_statements.mut_prepend_random_meta_command, 9),
  # # ##############################################################################################
  (mutator_corpus_splicing.mut_splice_append, 4),
  (mutator_corpus_splicing.mut_splice_prepend, 4),
  (mutator_corpus_splicing.mut_splice_intersperse, 4),
  (mutator_corpus_splicing.mut_splice_replace, 1),
  # # ##############################################################################################
  (mutator_equivalent_rewrites.mut_between_to_and, 3),
  (mutator_equivalent_rewrites.mut_and_to_between, 3),
  (mutator_equivalent_rewrites.mut_in_to_or, 3),
  (mutator_equivalent_rewrites.mut_arithmetic_identity, 3),
  (mutator_equivalent_rewrites.mut_double_negate, 3),
  (mutator_equivalent_rewrites.mut_string_concat_identity, 3),
  (mutator_equivalent_rewrites.mut_double_not, 3),
  (mutator_equivalent_rewrites.mut_de_morgan, 3),
  (mutator_equivalent_rewrites.mut_tautology_and, 3),
  (mutator_equivalent_rewrites.mut_tautology_or, 3),
  (mutator_equivalent_rewrites.mut_is_null_to_isnull, 2),
  (mutator_equivalent_rewrites.mut_eq_to_is, 2),
  (mutator_equivalent_rewrites.mut_inject_null_check, 2),
  (mutator_equivalent_rewrites.mut_subquery_in_from, 3),
  (mutator_equivalent_rewrites.mut_add_indexed_by, 2),
  (mutator_equivalent_rewrites.mut_qualify_with_main, 2),
  (mutator_equivalent_rewrites.mut_predicate_pushdown_blocker, 2),
  (mutator_equivalent_rewrites.mut_force_materialization, 2),
  (mutator_equivalent_rewrites.mut_comma_join_to_explicit, 2),
  (mutator_equivalent_rewrites.mut_swap_on_to_using, 2),
  (mutator_equivalent_rewrites.mut_inject_empty_block_comment, 2),
  (mutator_equivalent_rewrites.mut_inject_non_empty_block_comment, 2),
  (mutator_equivalent_rewrites.mut_inject_line_comment, 2),
  (mutator_equivalent_rewrites.mut_redundant_parens, 2),
  (mutator_equivalent_rewrites.mut_quote_identifier, 2),
  (mutator_equivalent_rewrites.mut_memeify, 2),
  (mutator_equivalent_rewrites.mut_swap_whitespace, 2),
  (mutator_equivalent_rewrites.mut_to_case_expr, 2),
  (mutator_equivalent_rewrites.mut_iif_to_case, 2),
  (mutator_equivalent_rewrites.mut_case_to_iif, 2),
]

MUTATIONS: list[Callable[[str], str]] = [m for m, _ in _MUTATORS_WEIGHTED]
_MUT_FUNCS = MUTATIONS
_MUT_WEIGHTS = [w for _, w in _MUTATORS_WEIGHTED]


def mutate(sql: str) -> str:
  n = random.choices([1, 2, 3, 10, 20], weights=[3, 7, 13, 27, 50], k=1)[0]
  out = sql
  _MUT_SOFT_CAP = 100000
  for _ in range(n):
    if len(out) > _MUT_SOFT_CAP:
      break
    m = random.choices(_MUT_FUNCS, weights=_MUT_WEIGHTS, k=1)[0]
    try:
      new = m(out)
    except Exception:
      continue
    if isinstance(new, str):
      out = new
  return out