(* storage/page.ml *)
let page_size = 4096  (* 4KB 页大小 *)

type page = {
  id: int;
  data: Bytes.t;
  mutable dirty: bool;
}

let create id = {
  id;
  data = Bytes.create page_size;
  dirty = false;
}

(* 页内偏移读写 *)
let read_int page offset =
  Bytes.get_int32_le page.data offset |> Int32.to_int

let write_int page offset value =
  Bytes.set_int32_le page.data offset (Int32.of_int value);
  page.dirty <- true

let read_bytes page offset length =
  Bytes.sub page.data offset length

let write_bytes page offset data =
  Bytes.blit data 0 page.data offset (Bytes.length data);
  page.dirty <- true

let read_string page offset length =
  Bytes.sub_string page.data offset length

let write_string page offset s =
  Bytes.blit_string s 0 page.data offset (String.length s);
  page.dirty <- true

(* 页头结构 *)
type page_header = {
  page_type: int;      (* 0=数据页, 1=索引页, 2=溢出页 *)
  num_cells: int;      (* 单元格数量 *)
  free_offset: int;    (* 空闲区域起始 *)
  cell_offset: int;    (* 单元格区域起始 *)
}

let read_header page =
  {
    page_type = read_int page 0;
    num_cells = read_int page 4;
    free_offset = read_int page 8;
    cell_offset = read_int page 12;
  }

let write_header page header =
  write_int page 0 header.page_type;
  write_int page 4 header.num_cells;
  write_int page 8 header.free_offset;
  write_int page 12 header.cell_offset

(* storage/btree.ml *)
open Page

type btree_node =
  | Leaf of {
      keys: string list;
      values: bytes list;
      next: int;  (* 叶子链表 *)
    }
  | Internal of {
      keys: string list;
      children: int list;
    }

type btree = {
  fd: Unix.file_descr;
  mutable root_page: int;
  mutable num_pages: int;
}

let create filename =
  let fd = Unix.openfile filename
    [Unix.O_RDWR; Unix.O_CREAT; Unix.O_TRUNC] 0o644 in
  let tree = { fd; root_page = 0; num_pages = 1 } in
  (* 初始化根页 *)
  let root = create 0 in
  let header = { page_type = 1; num_cells = 0; free_offset = page_size; cell_offset = 16 } in
  write_header root header;
  (* 写入到文件 *)
  Unix.lseek fd 0 Unix.SEEK_SET |> ignore;
  Unix.write fd root.data 0 page_size |> ignore;
  tree

(* 二分搜索 *)
let binary_search keys target =
  let rec aux lo hi =
    if lo >= hi then lo
    else
      let mid = lo + (hi - lo) / 2 in
      let cmp = String.compare keys.(mid) target in
      if cmp < 0 then aux (mid + 1) hi
      else if cmp > 0 then aux lo mid
      else mid
  in
  aux 0 (Array.length keys)

(* 插入 *)
let rec insert tree key value =
  let root = read_page tree.fd tree.root_page in
  let header = read_header root in
  match header.page_type with
  | 1 -> (* 根是叶子 *)
    insert_into_leaf tree root key value
  | 2 -> (* 根是内部节点 *)
    let new_root = split_if_needed tree root key value in
    if new_root <> tree.root_page then
      tree.root_page <- new_root
  | _ -> failwith "未知页类型"

and insert_into_leaf tree page key value =
  let header = read_header page in
  if header.num_cells < max_leaf_cells then begin
    (* 直接插入 *)
    let idx = find_insert_position page key in
    shift_cells_right page idx;
    write_cell page idx key value;
    write_header page { header with num_cells = header.num_cells + 1 }
  end else begin
    (* 分裂 *)
    split_leaf tree page key value
  end

and split_leaf tree page key value =
  let mid = max_leaf_cells / 2 in
  let new_page_id = allocate_page tree in
  let new_page = create new_page_id in
  (* 将一半数据移动到新页 *)
  (* ... *)
  (* 更新父节点 *)
  update_parent tree page.id new_page_id

(* 查找 *)
let find tree key =
  let page = read_page tree.fd tree.root_page in
  let header = read_header page in
  match header.page_type with
  | 1 -> find_in_leaf page key
  | 2 -> find_in_internal tree page key
  | _ -> None

and find_in_leaf page key =
  let header = read_header page in
  let keys = read_keys page header.num_cells in
  let idx = binary_search (Array.of_list keys) key in
  if idx < Array.length keys && keys.(idx) = key then
    Some (read_value page idx)
  else None

and find_in_internal tree page key =
  let header = read_header page in
  let keys = read_keys page header.num_cells in
  let children = read_children page (header.num_cells + 1) in
  let idx = binary_search (Array.of_list keys) key in
  let child_page_id = List.nth children idx in
  let child_page = read_page tree.fd child_page_id in
  find tree child_page key

(* storage/buffer.ml *)
type buffer_pool = {
  pool_size: int;
  pages: (int, page) Hashtbl.t;
  mutable lru: int list;  (* 最近使用的页 ID *)
  fd: Unix.file_descr;
}

let create_pool pool_size fd = {
  pool_size;
  pages = Hashtbl.create pool_size;
  lru = [];
  fd;
}

(* 获取页（带 LRU 淘汰） *)
let get_page pool page_id =
  match Hashtbl.find_opt pool.pages page_id with
  | Some page ->
    (* 更新 LRU *)
    pool.lru <- page_id :: List.filter ((<>) page_id) pool.lru;
    page
  | None ->
    (* 需要从磁盘读取 *)
    if Hashtbl.length pool.pages >= pool.pool_size then begin
      (* 淘汰最久未使用的页 *)
      let victim_id = List.last pool.lru in
      let victim = Hashtbl.find pool.pages victim_id in
      if victim.dirty then
        flush_page pool.fd victim;
      Hashtbl.remove pool.pages victim_id;
      pool.lru <- List.filter ((<>) victim_id) pool.lru
    end;
    (* 读取新页 *)
    let page = read_from_disk pool.fd page_id in
    Hashtbl.add pool.pages page_id page;
    pool.lru <- page_id :: pool.lru;
    page

(* 刷新脏页 *)
let flush_all pool =
  Hashtbl.iter (fun _ page ->
    if page.dirty then begin
      flush_page pool.fd page;
      page.dirty <- false
    end
  ) pool.pages

(* 预取 *)
let prefetch pool page_ids =
  List.iter (fun id ->
    if not (Hashtbl.mem pool.pages id) then
      ignore (get_page pool id)
  ) page_ids

(* storage/wal.ml *)
type wal_entry = {
  lsn: int;          (* 日志序列号 *)
  txn_id: int;       (* 事务 ID *)
  page_id: int;      (* 页 ID *)
  offset: int;       (* 页内偏移 *)
  before: bytes;     (* 修改前数据 *)
  after: bytes;      (* 修改后数据 *)
}

type wal = {
  fd: Unix.file_descr;
  mutable next_lsn: int;
  mutable buffer: bytes list;
  buffer_size: int;
}

let create_wal filename = {
  fd = Unix.openfile filename
    [Unix.O_RDWR; Unix.O_CREAT; Unix.O_APPEND] 0o644;
  next_lsn = 0;
  buffer = [];
  buffer_size = 1024;
}

(* 写入日志 *)
let write_log wal entry =
  let entry = { entry with lsn = wal.next_lsn } in
  wal.next_lsn <- wal.next_lsn + 1;
  wal.buffer <- entry :: wal.buffer;
  if List.length wal.buffer >= wal.buffer_size then
    flush_wal wal

(* 刷新到磁盘 *)
let flush_wal wal =
  let entries = List.rev wal.buffer in
  wal.buffer <- [];
  List.iter (fun entry ->
    let data = serialize_entry entry in
    Unix.write fd (Bytes.of_string data) 0 (String.length data) |> ignore
  ) entries;
  Unix.fsync wal.fd  (* 确保持久化 *)

(* 恢复 *)
let recover wal =
  let entries = read_all_entries wal.fd in
  (* 重放未提交的事务 *)
  List.iter (fun entry ->
    (* 应用修改 *)
    apply_entry entry
  ) entries

(* 检查点 *)
let checkpoint wal buffer_pool =
  (* 1. 刷新所有脏页 *)
  Buffer.flush_all buffer_pool;
  (* 2. 刷新 WAL *)
  flush_wal wal;
  (* 3. 截断已检查点的日志 *)
  Unix.ftruncate wal.fd 0;
  wal.next_lsn <- 0

(* query/lexer.mll *)
{
  open Parser
}

let digit = ['0'-'9']
let alpha = ['a'-'z' 'A'-'Z']
let ident = alpha (alpha | digit | '_')*
let whitespace = [' ' '\t' '\n' '\r']
let string_literal = '\'' [^ '\'']* '\''

rule token = parse
  | whitespace+ { token lexbuf }
  | "--" [^ '\n']* { token lexbuf }
  | digit+ as num { NUM (float_of_string num) }
  | string_literal as s { STRING (String.sub s 1 (String.length s - 2)) }
  | ident as id {
    match String.uppercase_ascii id with
    | "SELECT" -> SELECT
    | "FROM" -> FROM
    | "WHERE" -> WHERE
    | "AND" -> AND
    | "OR" -> OR
    | "NOT" -> NOT
    | "INSERT" -> INSERT
    | "INTO" -> INTO
    | "VALUES" -> VALUES
    | "UPDATE" -> UPDATE
    | "SET" -> SET
    | "DELETE" -> DELETE
    | "CREATE" -> CREATE
    | "TABLE" -> TABLE
    | "INDEX" -> INDEX
    | "BEGIN" -> BEGIN
    | "COMMIT" -> COMMIT
    | "ROLLBACK" -> ROLLBACK
    | "TRUE" -> BOOL true
    | "FALSE" -> BOOL false
    | "NULL" -> NULL
    | "JOIN" -> JOIN
    | "ON" -> ON
    | "GROUP" -> GROUP
    | "BY" -> BY
    | "ORDER" -> ORDER
    | "ASC" -> ASC
    | "DESC" -> DESC
    | "LIMIT" -> LIMIT
    | "OFFSET" -> OFFSET
    | "AS" -> AS
    | _ -> IDENT id
  }
  | '+' { PLUS }
  | '-' { MINUS }
  | '*' { STAR }
  | '/' { SLASH }
  | '=' { EQ }
  | '<' { LT }
  | '>' { GT }
  | "<=" { LE }
  | ">=" { GE }
  | "!=" | "<>" { NEQ }
  | '(' { LPAREN }
  | ')' { RPAREN }
  | ',' { COMMA }
  | ';' { SEMICOLON }
  | '.' { DOT }
  | eof { EOF }
  | _ as c { failwith (Printf.sprintf "意外字符: %c" c) }

(* query/ast.ml *)

type column = {
  name: string;
  table: string option;
  alias: string option;
}

type literal =
  | LFloat of float
  | LString of string
  | LBool of bool
  | LNull

type binary_op = Add | Sub | Mul | Div | Eq | Neq | Lt | Gt | Le | Ge | And | Or
type unary_op = Not | Neg

type expr =
  | EColumn of column
  | ELiteral of literal
  | EBinaryOp of binary_op * expr * expr
  | EUnaryOp of unary_op * expr
  | EFunc of string * expr list
  | EAgg of string * expr  (* COUNT, SUM, AVG 等 *)

type select_item =
  | Star
  | Expr of expr * string option  (* expression AS alias *)

type join_type = Inner | Left | Right | Full

type table_ref =
  | TableName of string * string option  (* name AS alias *)
  | Join of table_ref * table_ref * join_type * expr option

type order_item = {
  expr: expr;
  direction: [`Asc | `Desc];
}

type statement =
  | Select of {
      columns: select_item list;
      from: table_ref list;
      where: expr option;
      group_by: expr list;
      having: expr option;
      order_by: order_item list;
      limit: int option;
      offset: int option;
    }
  | Insert of {
      table: string;
      columns: string list;
      values: expr list list;
    }
  | Update of {
      table: string;
      set: (string * expr) list;
      where: expr option;
    }
  | Delete of {
      table: string;
      where: expr option;
    }
  | CreateTable of {
      name: string;
      columns: column_def list;
    }
  | CreateIndex of {
      name: string;
      table: string;
      columns: string list;
    }
  | Begin
  | Commit
  | Rollback

and column_def = {
  col_name: string;
  col_type: string;
  not_null: bool;
  primary_key: bool;
}

(* query/parser.mly *)
%{
  open Ast
%}

%token <float> NUM
%token <string> STRING
%token <string> IDENT
%token <bool> BOOL
%token NULL
%token SELECT FROM WHERE AND OR NOT
%token INSERT INTO VALUES UPDATE SET DELETE
%token CREATE TABLE INDEX BEGIN COMMIT ROLLBACK
%token JOIN ON GROUP BY ORDER ASC DESC LIMIT OFFSET AS
%token PLUS MINUS STAR SLASH
%token EQ LT GT LE GE NEQ
%token LPAREN RPAREN COMMA SEMICOLON DOT
%token EOF

%left OR
%left AND
%right NOT
%left EQ NEQ
%left LT GT LE GE
%left PLUS MINUS
%left STAR SLASH

%start <statement> statement

%%

statement:
  | s = select_stmt; SEMICOLON?; EOF { s }
  | s = insert_stmt; SEMICOLON?; EOF { s }
  | s = update_stmt; SEMICOLON?; EOF { s }
  | s = delete_stmt; SEMICOLON?; EOF { s }
  | s = create_stmt; SEMICOLON?; EOF { s }
  | BEGIN; SEMICOLON?; EOF { Begin }
  | COMMIT; SEMICOLON?; EOF { Commit }
  | ROLLBACK; SEMICOLON?; EOF { Rollback }

select_stmt:
  | SELECT; columns = select_list;
    FROM; from = table_list;
    where = where_clause?;
    group = group_clause?;
    having = having_clause?;
    order = order_clause?;
    limit = limit_clause?;
    offset = offset_clause?
    { Select { columns; from; where; group_by = Option.value group ~default:[];
               having; order_by = Option.value order ~default:[];
               limit; offset } }

select_list:
  | STAR { [Star] }
  | items = separated_list(COMMA, select_item) { items }

select_item:
  | e = expr; AS; alias = IDENT { Expr (e, Some alias) }
  | e = expr { Expr (e, None) }

table_list:
  | tables = separated_list(COMMA, table_ref) { tables }

table_ref:
  | name = IDENT; AS; alias = IDENT { TableName (name, Some alias) }
  | name = IDENT { TableName (name, None) }
  | t1 = table_ref; JOIN; t2 = table_ref; ON; cond = expr
    { Join (t1, t2, Inner, Some cond) }

where_clause:
  | WHERE; e = expr { e }

group_clause:
  | GROUP; BY; exprs = separated_list(COMMA, expr) { exprs }

having_clause:
  | HAVING; e = expr { e }

order_clause:
  | ORDER; BY; items = separated_list(COMMA, order_item) { items }

order_item:
  | e = expr; ASC { { expr = e; direction = `Asc } }
  | e = expr; DESC { { expr = e; direction = `Desc } }
  | e = expr { { expr = e; direction = `Asc } }

limit_clause:
  | LIMIT; n = NUM { int_of_float n }

offset_clause:
  | OFFSET; n = NUM { int_of_float n }

expr:
  | e = simple_expr { e }
  | left = expr; op = binop; right = expr { EBinaryOp (op, left, right) }
  | NOT; e = expr { EUnaryOp (Not, e) }
  | MINUS; e = expr %prec NOT { EUnaryOp (Neg, e) }

simple_expr:
  | n = NUM { ELiteral (LFloat n) }
  | s = STRING { ELiteral (LString s) }
  | b = BOOL { ELiteral (LBool b) }
  | NULL { ELiteral LNull }
  | id = IDENT { EColumn { name = id; table = None; alias = None } }
  | table = DOT; col = IDENT { EColumn { name = col; table = Some table; alias = None } }
  | func = IDENT; LPAREN; args = separated_list(COMMA, expr); RPAREN
    { EFunc (func, args) }
  | LPAREN; e = expr; RPAREN { e }

%inline binop:
  | EQ { Eq } | NEQ { Neq }
  | LT { Lt } | GT { Gt }
  | LE { Le } | GE { Ge }
  | PLUS { Add } | MINUS { Sub }
  | STAR { Mul } | SLASH { Div }
  | AND { And } | OR { Or }

insert_stmt:
  | INSERT; INTO; table = IDENT;
    LPAREN; columns = separated_list(COMMA, IDENT); RPAREN;
    VALUES; values = separated_list(COMMA, value_list)
    { Insert { table; columns; values } }

value_list:
  | LPAREN; values = separated_list(COMMA, expr); RPAREN { values }

update_stmt:
  | UPDATE; table = IDENT;
    SET; assignments = separated_list(COMMA, assignment);
    where = where_clause?
    { Update { table; set = assignments; where } }

assignment:
  | col = IDENT; EQ; e = expr { (col, e) }

delete_stmt:
  | DELETE; FROM; table = IDENT;
    where = where_clause?
    { Delete { table; where } }

create_stmt:
  | CREATE; TABLE; name = IDENT;
    LPAREN; columns = separated_list(COMMA, column_def); RPAREN
    { CreateTable { name; columns } }
  | CREATE; INDEX; name = IDENT;
    ON; table = IDENT;
    LPAREN; columns = separated_list(COMMA, IDENT); RPAREN
    { CreateIndex { name; table; columns } }

column_def:
  | name = IDENT; typ = IDENT; constraints = list(constraint_)
    { { col_name = name; col_type = typ;
        not_null = List.mem `NotNull constraints;
        primary_key = List.mem `PrimaryKey constraints } }

constraint_:
  | NOT; NULL { `NotNull }
  | PRIMARY; KEY { `PrimaryKey }

(* query/optimizer.ml *)
open Ast

(* 谓词下推 *)
let push_down_predicates stmt =
  match stmt with
  | Select { columns; from; where; group_by; having; order_by; limit; offset } ->
    let new_where = match where with
      | Some pred ->
        (* 将可以下推的谓词移动到对应的表 *)
        let pushed, remaining = partition_predicates pred from in
        (* 将 pushed 应用到 from 中的表 *)
        apply_pushed_predicates from pushed;
        remaining
      | None -> None
    in
    Select { columns; from; where = new_where; group_by; having; order_by; limit; offset }
  | _ -> stmt

(* 常量折叠 *)
let rec fold_constants = function
  | EBinaryOp (Add, ELiteral (LFloat a), ELiteral (LFloat b)) ->
    ELiteral (LFloat (a +. b))
  | EBinaryOp (Sub, ELiteral (LFloat a), ELiteral (LFloat b)) ->
    ELiteral (LFloat (a -. b))
  | EBinaryOp (Mul, ELiteral (LFloat a), ELiteral (LFloat b)) ->
    ELiteral (LFloat (a *. b))
  | EBinaryOp (Div, ELiteral (LFloat a), ELiteral (LFloat b)) when b <> 0.0 ->
    ELiteral (LFloat (a /. b))
  | EBinaryOp (Eq, ELiteral a, ELiteral b) ->
    ELiteral (LBool (a = b))
  | EBinaryOp (And, ELiteral (LBool a), ELiteral (LBool b)) ->
    ELiteral (LBool (a && b))
  | EBinaryOp (Or, ELiteral (LBool a), ELiteral (LBool b)) ->
    ELiteral (LBool (a || b))
  | EBinaryOp (op, left, right) ->
    EBinaryOp (op, fold_constants left, fold_constants right)
  | EUnaryOp (Not, ELiteral (LBool b)) ->
    ELiteral (LBool (not b))
  | EUnaryOp (Neg, ELiteral (LFloat n)) ->
    ELiteral (LFloat (-.n))
  | e -> e

(* 选择最优执行计划 *)
type execution_plan =
  | FullScan of string
  | IndexScan of string * string * expr
  | NestedLoopJoin of execution_plan * execution_plan * expr
  | HashJoin of execution_plan * execution_plan * string * string

let estimate_cost = function
  | FullScan table ->
    let stats = get_table_stats table in
    stats.num_rows  (* 全表扫描代价 = 行数 *)
  | IndexScan (table, index, _) ->
    let stats = get_table_stats table in
    stats.num_rows / 10  (* 索引扫描代价 ≈ 行数/10 *)
  | NestedLoopJoin (left, right, _) ->
    estimate_cost left * estimate_cost right
  | HashJoin (left, right, _, _) ->
    estimate_cost left + estimate_cost right

let choose_best_plan table where =
  let plans = [
    FullScan table;
    (* 检查是否有可用索引 *)
    (match find_usable_index table where with
     | Some (index, pred) -> IndexScan (table, index, pred)
     | None -> FullScan table);
  ] in
  List.fold_left (fun best plan ->
    if estimate_cost plan < estimate_cost best then plan else best
  ) (List.hd plans) (List.tl plans)