/*
    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;
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,
) -> Result<Box<dyn Iterator<Item = Result<(RowNum, Vec<u8>)>> + 'a>> {
    Ok(match sort {
        Sort::None => Box::new(
            PrefixIterator {
                inner: txn.iter(&prefix, false)?,
                prefix: prefix.into(),
            }
            .map(|k| {
                k.map(|k| {
                    (
                        RowNum::from_be_bytes(k[k.len() - 8..].try_into().unwrap()),
                        k[4..].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");
            Box::new(
                PrefixIterator {
                    inner: match value_sort.order {
                        SortOrder::Ascending => txn.iter(&prefix, false)?,
                        SortOrder::Descending => txn.iter(&prefix_end(prefix.clone()), true)?,
                    },
                    prefix: prefix.into(),
                }
                .map(|k| {
                    k.map(|k| {
                        (
                            RowNum::from_be_bytes(k[k.len() - 8..].try_into().unwrap()),
                            k[4..].to_vec(),
                        )
                    })
                }),
            )
        }
        Sort::TextSearch(_, _) => todo!(),
    })
}

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

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