/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2026 metamuffin <metamuffin.org>
*/

use crate::{
    MultiBehaviour, RowNum, Sort, SortOrder,
    kv::{
        SubtreeNum,
        helpers::{read_counter, write_counter},
        index_key::{IndexKey, SortKey},
        prefix_iterator::PrefixIterator,
    },
};
use anyhow::{Result, bail};
use jellyobject::Object;
use std::iter::empty;

pub fn update_index(
    txn: &mut dyn jellykv::Transaction,
    is: SubtreeNum,
    ik: &IndexKey,
    row: RowNum,
    ob: &Object,
    remove: bool,
) -> Result<()> {
    let mut ks = vec![is.to_be_bytes().to_vec()];
    ik.0.apply(ob, &mut ks);
    match &ik.1 {
        SortKey::None => {
            for mut k in ks {
                k.extend(row.to_be_bytes());
                if remove {
                    txn.del(&k)?;
                } else {
                    txn.set(&k, &[])?;
                }
            }
        }
        SortKey::Random => {
            for mut k in ks {
                k.extend(randomize(row).to_be_bytes());
                if remove {
                    txn.del(&k)?;
                } else {
                    txn.set(&k, &row.to_be_bytes())?;
                }
            }
        }
        SortKey::Count => {
            for k in ks {
                let mut c = read_counter(txn, &k, 0)?;
                if remove {
                    c -= 1;
                } else {
                    c += 1;
                }
                write_counter(txn, &k, c)?;
            }
        }
        SortKey::Value(path, multi_behaviour) => {
            assert_eq!(*multi_behaviour, MultiBehaviour::First, "TODO");
            if let Some(value) = path.get_matching_value(ob) {
                for mut k in ks {
                    k.extend(value);
                    k.extend(row.to_be_bytes());
                    if remove {
                        txn.del(&k)?;
                    } else {
                        txn.set(&k, &[])?;
                    }
                }
            }
        }
        SortKey::Text(_) => todo!(),
    }
    Ok(())
}

pub fn read_count_index(txn: &dyn jellykv::Transaction, prefix: Vec<u8>) -> Result<u64> {
    read_counter(txn, &prefix, 0)
}

pub fn iter_index<'a>(
    txn: &'a dyn jellykv::Transaction,
    prefix: Vec<u8>,
    sort: &Sort,
    resume: Option<Vec<u8>>,
) -> Result<Box<dyn Iterator<Item = Result<(RowNum, Vec<u8>)>> + 'a>> {
    Ok(match sort {
        Sort::None => {
            let mut start = prefix.clone();
            if let Some(resume) = resume {
                if resume.len() != 8 {
                    bail!("invalid resume length")
                }
                start.extend(resume);
                start = inc_key(start);
            }
            Box::new(
                PrefixIterator {
                    inner: txn.iter(&start, false)?,
                    prefix: prefix.into(),
                }
                .map(|k| {
                    k.map(|k| {
                        let rn = RowNum::from_be_bytes(k[k.len() - 8..].try_into().unwrap());
                        (rn, rn.to_be_bytes().to_vec())
                    })
                }),
            )
        }
        Sort::Random(seed) => {
            let mut k = prefix.clone();
            k.extend(randomize(*seed).to_be_bytes());
            let Some(it) = txn.iter(&k, false)?.next() else {
                return Ok(Box::new(empty()));
            };
            let res = it?;
            if !res.starts_with(&prefix) {
                return Ok(Box::new(empty()));
            }
            let r = txn.get(&res)?.unwrap();
            let row = RowNum::from_be_bytes(r.try_into().unwrap());
            Box::new([Ok((row, vec![]))].into_iter())
        }
        Sort::Value(value_sort) => {
            assert!(value_sort.offset.is_none(), "TODO");
            let (mut start, rev) = match value_sort.order {
                SortOrder::Ascending => (prefix.clone(), false),
                SortOrder::Descending => (inc_key(prefix.clone()), true),
            };
            if let Some(resume) = resume {
                start = prefix.iter().copied().chain(resume).collect();
                if rev {
                    start = dec_key(start);
                } else {
                    start = inc_key(start);
                }
            }
            Box::new(
                PrefixIterator {
                    inner: txn.iter(&start, rev)?,
                    prefix: prefix.clone().into(),
                }
                .map(move |k| {
                    k.map(|k| {
                        (
                            RowNum::from_be_bytes(k[k.len() - 8..].try_into().unwrap()),
                            k[prefix.len()..].to_vec(),
                        )
                    })
                }),
            )
        }
        Sort::TextSearch(_, _) => todo!(),
    })
}

fn inc_key(mut k: Vec<u8>) -> Vec<u8> {
    for v in k.iter_mut().rev() {
        let (nv, carry) = v.overflowing_add(1);
        *v = nv;
        if !carry {
            break;
        }
    }
    k
}
fn dec_key(mut k: Vec<u8>) -> Vec<u8> {
    for v in k.iter_mut().rev() {
        let (nv, carry) = v.overflowing_sub(1);
        *v = nv;
        if !carry {
            break;
        }
    }
    k
}

fn randomize(mut x: u64) -> u64 {
    for _ in 0..15 {
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
    }
    x
}