// Copyright (c) 2023-2026 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use std::fmt;

use crate::api::{AsCStr, FieldName, HashMap, HashSet, OrderByInfo, Varno};
use crate::index::fast_fields_helper::WhichFastField;
use crate::postgres::customscan::basescan::projections::window_agg::WindowAggregateInfo;
use crate::postgres::customscan::basescan::ExecMethodType;
use crate::postgres::customscan::builders::custom_path::OrderByStyle;
use crate::query::SearchQueryInput;

use pgrx::pg_sys::AsPgCStr;
use pgrx::{pg_sys, FromDatum, PgList};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Limit {
    Static(usize),
    Parameterized { limit_param_id: i32 },
}

impl fmt::Display for Limit {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Limit::Static(n) => write!(f, "{n}"),
            Limit::Parameterized { .. } => write!(f, "<parameterized>"),
        }
    }
}

impl Limit {
    pub unsafe fn from_param(node: *mut pg_sys::Node) -> Option<Self> {
        use crate::nodecast;

        if node.is_null() {
            return None;
        }

        if let Some(param) = nodecast!(Param, T_Param, node) {
            if (*param).paramkind == pg_sys::ParamKind::PARAM_EXTERN {
                return Some(Limit::Parameterized {
                    limit_param_id: (*param).paramid,
                });
            }
            return None;
        }

        if let Some(func_expr) = nodecast!(FuncExpr, T_FuncExpr, node) {
            let args = PgList::<pg_sys::Node>::from_pg((*func_expr).args);
            if args.len() == 1 {
                return Limit::from_param(args.get_ptr(0).unwrap());
            }
        }

        if let Some(relabel) = nodecast!(RelabelType, T_RelabelType, node) {
            return Limit::from_param((*relabel).arg.cast());
        }

        if let Some(coerce) = nodecast!(CoerceViaIO, T_CoerceViaIO, node) {
            return Limit::from_param((*coerce).arg.cast());
        }

        None
    }

    pub fn param_id(&self) -> Option<i32> {
        match self {
            Limit::Parameterized { limit_param_id } => Some(*limit_param_id),
            Limit::Static(_) => None,
        }
    }

    pub fn static_value(&self) -> Option<usize> {
        match self {
            Limit::Static(n) => Some(*n),
            Limit::Parameterized { .. } => None,
        }
    }

    pub unsafe fn resolve(&self, estate: *mut pg_sys::EState) -> usize {
        match self {
            Limit::Static(n) => *n,
            Limit::Parameterized { limit_param_id } => {
                let param_list = (*estate).es_param_list_info;
                assert!(
                    !param_list.is_null(),
                    "es_param_list_info is NULL but we have a parameterized LIMIT"
                );

                let idx = (*limit_param_id - 1) as usize;
                assert!(
                    idx < (*param_list).numParams as usize,
                    "LIMIT param_id {} out of range (numParams={})",
                    limit_param_id,
                    (*param_list).numParams
                );

                let param_data = &(*param_list)
                    .params
                    .as_slice((*param_list).numParams as usize)[idx];
                assert!(
                    !param_data.isnull,
                    "LIMIT parameter ${} is NULL",
                    limit_param_id
                );

                i64::from_datum(param_data.value, param_data.isnull)
                    .expect("LIMIT parameter should be convertible to i64") as usize
            }
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct PrivateData {
    heaprelid: Option<pg_sys::Oid>,
    indexrelid: Option<pg_sys::Oid>,
    range_table_index: Option<pg_sys::Index>,
    query: Option<SearchQueryInput>,
    limit: Option<Limit>,
    // ORDER-BY info that will be used iff the appropriate ExecMethodType is chosen.
    maybe_orderby_info: Option<Vec<OrderByInfo>>,
    #[serde(with = "var_attname_lookup_serializer")]
    var_attname_lookup: Option<HashMap<(Varno, pg_sys::AttrNumber), FieldName>>,
    segment_count: usize,
    // The fast fields which were identified during planning time as potentially being
    // needed at execution time. In order for our planning-time-chosen ExecMethodType to be
    // accurate, this must always be a superset of the fields extracted from the execution
    // time target list.
    planned_which_fast_fields: Option<HashSet<WhichFastField>>,
    target_list_len: Option<usize>,
    referenced_columns_count: usize,
    need_scores: bool,
    exec_method_type: ExecMethodType,
    // Additional search predicates from join filters that are relevant for snippet/score generation
    // Stores the entire simplified Boolean expression to preserve OR structures like (TRUE OR name:"Rowling")
    join_predicates: Option<SearchQueryInput>,
    ambulkdelete_epoch: u32,
    // Window aggregates to compute during Top K execution
    window_aggregates: Vec<WindowAggregateInfo>,
    // Whether this path was chosen as a sorted path (declares pathkeys for index's sort_by field).
    // When true, the execution should use the sorted merge path for segment scanning.
    use_sorted_path: bool,
}

mod var_attname_lookup_serializer {
    use super::*;

    use serde::{de::Error, Deserializer, Serializer};

    fn key_to_string(key: &(Varno, i16)) -> String {
        format!("{},{}", key.0, key.1)
    }

    fn key_from_string(s: &str) -> Result<(Varno, i16), String> {
        let mut parts = s.splitn(2, ',');
        let p1_str = parts
            .next()
            .ok_or_else(|| "Missing first part of key".to_string())?;
        let p2_str = parts
            .next()
            .ok_or_else(|| "Missing second part of key".to_string())?;

        let p1 = p1_str
            .parse::<Varno>()
            .map_err(|e| format!("Failed to parse first key part '{p1_str}': {e}"))?;
        let p2 = p2_str
            .parse::<i16>()
            .map_err(|e| format!("Failed to parse second key part '{p2_str}': {e}"))?;

        Ok((p1, p2))
    }

    pub fn serialize<S>(
        map_option: &Option<HashMap<(Varno, i16), FieldName>>,
        serializer: S,
    ) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let Some(map) = map_option else {
            return serializer.serialize_none();
        };

        // Serialize as Vec<(String, String)>.
        map.iter()
            .map(|(k, v)| (key_to_string(k), v))
            .collect::<Vec<(String, &FieldName)>>()
            .serialize(serializer)
    }

    #[allow(clippy::type_complexity)]
    pub fn deserialize<'de, D>(
        deserializer: D,
    ) -> Result<Option<HashMap<(Varno, i16), FieldName>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        // Deserialize as Vec<(String, String)>.
        let Some(string_map) = Option::<Vec<(&'de str, FieldName)>>::deserialize(deserializer)?
        else {
            return Ok(None);
        };

        let mut map = HashMap::default();
        map.reserve(string_map.len());

        for (k_str, v) in string_map {
            let key_tuple = key_from_string(k_str)
                .map_err(|e| D::Error::custom(format!("Invalid key format '{k_str}': {e}")))?;
            map.insert(key_tuple, v);
        }
        Ok(Some(map))
    }
}

impl From<*mut pg_sys::List> for PrivateData {
    fn from(list: *mut pg_sys::List) -> Self {
        unsafe {
            let list = PgList::<pg_sys::Node>::from_pg(list);
            let node = list.get_ptr(0).unwrap();
            let content = node
                .as_c_str()
                .unwrap()
                .to_str()
                .expect("string node should be valid utf8");
            serde_json::from_str(content).unwrap()
        }
    }
}

impl From<PrivateData> for *mut pg_sys::List {
    fn from(value: PrivateData) -> Self {
        let content = serde_json::to_string(&value).unwrap();
        unsafe {
            let mut ser = PgList::new();
            ser.push(pg_sys::makeString(content.as_pg_cstr()).cast::<pg_sys::Node>());
            ser.into_pg()
        }
    }
}

//
// setter functions
//

impl PrivateData {
    pub fn set_heaprelid(&mut self, oid: pg_sys::Oid) {
        self.heaprelid = Some(oid);
    }

    pub fn set_indexrelid(&mut self, oid: pg_sys::Oid) {
        self.indexrelid = Some(oid);
    }

    pub fn set_range_table_index(&mut self, rti: pg_sys::Index) {
        self.range_table_index = Some(rti);
    }

    pub fn set_query(&mut self, query: SearchQueryInput) {
        self.query = Some(query);
    }

    pub fn set_limit(&mut self, limit: Option<Limit>) {
        self.limit = limit;
    }

    pub fn set_maybe_orderby_info(&mut self, style: Option<&Vec<OrderByStyle>>) {
        self.maybe_orderby_info = Some(OrderByStyle::extract_orderby_info(style));
    }

    pub fn set_var_attname_lookup(
        &mut self,
        var_attname_lookup: HashMap<(Varno, pg_sys::AttrNumber), FieldName>,
    ) {
        self.var_attname_lookup = Some(var_attname_lookup);
    }

    pub fn set_ambulkdelete_epoch(&mut self, epoch: u32) {
        self.ambulkdelete_epoch = epoch;
    }

    pub fn set_segment_count(&mut self, segment_count: usize) {
        self.segment_count = segment_count;
    }

    pub fn set_planned_which_fast_fields(
        &mut self,
        planned_which_fast_fields: HashSet<WhichFastField>,
    ) {
        self.planned_which_fast_fields = Some(planned_which_fast_fields);
    }

    pub fn set_exec_method_type(&mut self, exec_method_type: ExecMethodType) {
        self.exec_method_type = exec_method_type;
    }

    pub fn set_target_list_len(&mut self, len: Option<usize>) {
        self.target_list_len = len;
    }

    pub fn set_referenced_columns_count(&mut self, count: usize) {
        self.referenced_columns_count = count;
    }

    pub fn set_need_scores(&mut self, maybe: bool) {
        self.need_scores = maybe;
    }

    pub fn set_join_predicates(&mut self, predicates: Option<SearchQueryInput>) {
        self.join_predicates = predicates;
    }

    pub fn set_window_aggregates(&mut self, window_aggregates: Vec<WindowAggregateInfo>) {
        self.window_aggregates = window_aggregates;
    }

    pub fn set_use_sorted_path(&mut self, use_sorted: bool) {
        self.use_sorted_path = use_sorted;
    }
}

//
// getter functions
//

impl PrivateData {
    pub fn heaprelid(&self) -> Option<pg_sys::Oid> {
        self.heaprelid
    }

    pub fn indexrelid(&self) -> Option<pg_sys::Oid> {
        self.indexrelid
    }

    pub fn range_table_index(&self) -> Option<pg_sys::Index> {
        self.range_table_index
    }

    pub fn query(&self) -> &Option<SearchQueryInput> {
        &self.query
    }

    pub fn limit(&self) -> &Option<Limit> {
        &self.limit
    }

    pub fn static_limit(&self) -> Option<usize> {
        self.limit.as_ref().and_then(Limit::static_value)
    }

    pub fn maybe_orderby_info(&self) -> &Option<Vec<OrderByInfo>> {
        &self.maybe_orderby_info
    }

    pub fn var_attname_lookup(&self) -> &Option<HashMap<(Varno, pg_sys::AttrNumber), FieldName>> {
        &self.var_attname_lookup
    }

    pub fn segment_count(&self) -> usize {
        self.segment_count
    }

    pub fn target_list_len(&self) -> Option<usize> {
        self.target_list_len
    }

    pub fn planned_which_fast_fields(&self) -> &Option<HashSet<WhichFastField>> {
        &self.planned_which_fast_fields
    }

    pub fn exec_method_type(&self) -> &ExecMethodType {
        &self.exec_method_type
    }

    pub fn referenced_columns_count(&self) -> usize {
        debug_assert!(self.referenced_columns_count >= self.target_list_len.unwrap_or(0));
        self.referenced_columns_count
    }

    pub fn need_scores(&self) -> bool {
        self.need_scores
    }

    pub fn join_predicates(&self) -> &Option<SearchQueryInput> {
        &self.join_predicates
    }

    pub fn ambulkdelete_epoch(&self) -> u32 {
        self.ambulkdelete_epoch
    }

    pub fn window_aggregates(&self) -> &Vec<WindowAggregateInfo> {
        &self.window_aggregates
    }

    pub fn use_sorted_path(&self) -> bool {
        self.use_sorted_path
    }
}
