// 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 crate::api::{Cardinality, FieldName, HashSet, OrderByFeature, OrderByInfo, SortDirection};
use crate::index::fast_fields_helper::WhichFastField;
use crate::postgres::customscan::basescan::projections::window_agg::WindowAggregateInfo;
use crate::postgres::customscan::limit_offset::LimitOffset;
use crate::postgres::customscan::CustomScan;
use pgrx::{pg_sys, PgList};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone)]
pub enum OrderByStyle {
    Score {
        pathkey: *mut pg_sys::PathKey,
        rti: u32,
    },
    Field {
        pathkey: *mut pg_sys::PathKey,
        name: FieldName,
        rti: u32,
    },
    VectorDistance {
        pathkey: *mut pg_sys::PathKey,
        name: FieldName,
        rti: u32,
        /// See `OrderByFeature::VectorDistance::query_vector`.
        query_vector: crate::api::QueryVector,
        /// See `OrderByFeature::VectorDistance::metric`.
        metric: crate::vector::metric::VectorMetric,
    },
}

impl OrderByStyle {
    pub fn pathkey(&self) -> *mut pg_sys::PathKey {
        match self {
            OrderByStyle::Score { pathkey, .. } => *pathkey,
            OrderByStyle::Field { pathkey, .. } => *pathkey,
            OrderByStyle::VectorDistance { pathkey, .. } => *pathkey,
        }
    }

    pub fn direction(&self) -> SortDirection {
        unsafe {
            let pathkey = self.pathkey();
            assert!(!pathkey.is_null());
            let nulls_first = (*pathkey).pk_nulls_first;

            #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))]
            let is_asc = match (*pathkey).pk_strategy as u32 {
                pg_sys::BTLessStrategyNumber => true,
                pg_sys::BTGreaterStrategyNumber => false,
                value => panic!("unrecognized sort strategy number: {value}"),
            };
            #[cfg(feature = "pg18")]
            let is_asc = match (*pathkey).pk_cmptype {
                pg_sys::CompareType::COMPARE_LT => true,
                pg_sys::CompareType::COMPARE_GT => false,
                value => panic!("unrecognized compare type: {value}"),
            };

            match (is_asc, nulls_first) {
                (true, true) => SortDirection::AscNullsFirst,
                (true, false) => SortDirection::AscNullsLast,
                (false, true) => SortDirection::DescNullsFirst,
                (false, false) => SortDirection::DescNullsLast,
            }
        }
    }

    /// Extract ORDER BY information from query pathkeys
    /// In this case, we convert OrderByStyle to OrderByInfo for serialization.
    pub fn extract_orderby_info(order_pathkeys: Option<&Vec<OrderByStyle>>) -> Vec<OrderByInfo> {
        order_pathkeys
            .unwrap_or(&vec![])
            .iter()
            .map(|style| style.into())
            .collect()
    }
}

impl From<&OrderByStyle> for OrderByInfo {
    fn from(value: &OrderByStyle) -> Self {
        let feature = match value {
            OrderByStyle::Field { name, rti, .. } => OrderByFeature::Field {
                name: name.to_owned(),
                rti: *rti,
            },
            OrderByStyle::Score { rti, .. } => OrderByFeature::Score { rti: *rti },
            OrderByStyle::VectorDistance {
                name,
                rti,
                query_vector,
                metric,
                ..
            } => OrderByFeature::VectorDistance {
                name: name.to_owned(),
                rti: *rti,
                query_vector: query_vector.clone(),
                metric: *metric,
            },
        };
        OrderByInfo {
            feature,
            direction: value.direction(),
        }
    }
}

///
/// The type of ExecMethod that was chosen at planning time. We fully select an ExecMethodType at
/// planning time in order to be able to make claims about the sortedness and estimates for our
/// execution.
///
/// `which_fast_fields` lists in this enum are _all_ of the fast fields which were identified at
/// planning time: based on the join order that the planner ends up choosing, only a subset of
/// these might be used at execution time (in an order specified by the execution time target
/// list), but never a superset.
///
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub enum ExecMethodType {
    #[default]
    Normal,
    TopK {
        heaprelid: pg_sys::Oid,
        limit_offset: LimitOffset,
        orderby_info: Option<Vec<OrderByInfo>>,
        window_aggregates: Vec<WindowAggregateInfo>,
    },
    Columnar {
        which_fast_fields: HashSet<WhichFastField>,
        limit_offset: Option<LimitOffset>,
    },
}

impl ExecMethodType {
    /// Returns true if this execution method declares sorted output.
    /// Only Top K with an `ORDER BY` produces sorted output.
    pub fn declares_sorted_output(&self) -> bool {
        matches!(
            self,
            ExecMethodType::TopK {
                orderby_info: Some(..),
                ..
            }
        )
    }
}

#[derive(Debug, Eq, PartialEq, Hash, Copy, Clone)]
#[repr(u32)]
#[allow(dead_code)]
pub enum Flags {
    /// #define CUSTOMPATH_SUPPORT_BACKWARD_SCAN	0x0001
    BackwardScan = 0x0001,

    /// #define CUSTOMPATH_SUPPORT_MARK_RESTORE		0x0002
    MarkRestore = 0x0002,

    /// #define CUSTOMPATH_SUPPORT_PROJECTION		0x0004
    Projection = 0x0004,

    /// ParadeDB custom flag for indicating we want to force the plan to be used
    Force = 0x0008,

    /// ParadeDB custom flag set on a parallel-aware partial path that PostgreSQL is allowed to
    /// reject in favour of a serial sibling. Unlike a plain parallel-aware path (which clears the
    /// complete pathlist and installs a prohibitively-costed serial stub so the Gather must win),
    /// an `OfferParallel` partial path is merely added via `add_partial_path`: the complete pathlist
    /// is left intact and no stub is injected, so PostgreSQL costs the Gather against the real serial
    /// sibling and picks the cheaper. See `hook::add_path`.
    OfferParallel = 0x0010,
}

pub struct CustomPathBuilder<CS: CustomScan> {
    args: CS::Args,
    flags: HashSet<Flags>,

    custom_path_node: pg_sys::CustomPath,
}

#[derive(Copy, Clone, Debug)]
pub enum RestrictInfoType {
    BaseRelation,
    Join,
    None,
}

pub fn restrict_info(rel: &pg_sys::RelOptInfo) -> (PgList<pg_sys::RestrictInfo>, RestrictInfoType) {
    unsafe {
        let baseri = PgList::from_pg(rel.baserestrictinfo);
        let joinri = PgList::from_pg(rel.joininfo);

        if baseri.is_empty() && joinri.is_empty() {
            // both lists are empty, so return an empty list
            (PgList::new(), RestrictInfoType::None)
        } else if !baseri.is_empty() {
            // the baserestrictinfo has entries, so we prefer that first
            (baseri, RestrictInfoType::BaseRelation)
        } else {
            // only the joininfo has entries, so that's what we'll use
            (joinri, RestrictInfoType::Join)
        }
    }
}

impl<CS: CustomScan> CustomPathBuilder<CS> {
    pub fn new(
        root: *mut pg_sys::PlannerInfo,
        rel: *mut pg_sys::RelOptInfo,
        args: CS::Args,
    ) -> Self {
        unsafe {
            Self {
                args,
                flags: Default::default(),

                custom_path_node: pg_sys::CustomPath {
                    path: pg_sys::Path {
                        type_: pg_sys::NodeTag::T_CustomPath,
                        pathtype: pg_sys::NodeTag::T_CustomScan,
                        parent: rel,
                        pathtarget: (*rel).reltarget,
                        param_info: pg_sys::get_baserel_parampathinfo(
                            root,
                            rel,
                            pg_sys::bms_copy((*rel).lateral_relids),
                        ),
                        ..Default::default()
                    },
                    methods: CS::custom_path_methods(),
                    ..Default::default()
                },
            }
        }
    }

    pub fn args(&self) -> &CS::Args {
        &self.args
    }

    //
    // public settings
    //

    #[allow(dead_code)]
    pub fn clear_flags(mut self) -> Self {
        self.flags.clear();
        self
    }

    pub fn set_flag(mut self, flag: Flags) -> Self {
        self.flags.insert(flag);
        self
    }

    pub fn set_rows(mut self, rows: Cardinality) -> Self {
        self.custom_path_node.path.rows = rows;
        self
    }

    pub fn set_startup_cost(mut self, cost: pg_sys::Cost) -> Self {
        self.custom_path_node.path.startup_cost = cost;
        self
    }

    pub fn set_total_cost(mut self, cost: pg_sys::Cost) -> Self {
        self.custom_path_node.path.total_cost = cost;
        self
    }

    pub fn startup_cost(&self) -> pg_sys::Cost {
        self.custom_path_node.path.startup_cost
    }

    pub fn total_cost(&self) -> pg_sys::Cost {
        self.custom_path_node.path.total_cost
    }

    pub fn add_path_key(mut self, style: &OrderByStyle) -> Self {
        unsafe {
            let mut pklist =
                PgList::<pg_sys::PathKey>::from_pg(self.custom_path_node.path.pathkeys);
            pklist.push(style.pathkey());

            self.custom_path_node.path.pathkeys = pklist.into_pg();
        }
        self
    }

    pub fn set_pathkeys(mut self, pathkeys: *mut pg_sys::List) -> Self {
        self.custom_path_node.path.pathkeys = pathkeys;
        self
    }

    pub fn set_force_path(mut self, force: bool) -> Self {
        if force {
            self.flags.insert(Flags::Force);
        } else {
            self.flags.remove(&Flags::Force);
        }
        self
    }

    /// Configures the path for Parallel-Aware (Partial) execution.
    ///
    /// Setting this to `nworkers > 0` marks the path as:
    /// - `parallel_safe`: It can execute in a worker.
    /// - `parallel_aware`: It will coordinate with other workers to divide the work.
    pub fn set_parallel(mut self, nworkers: usize) -> Self {
        self.custom_path_node.path.parallel_aware = true;
        self.custom_path_node.path.parallel_safe = true;
        self.custom_path_node.path.parallel_workers =
            nworkers.try_into().expect("nworkers should be a valid i32");

        self
    }

    /// Configures whether the path is safe to run in a worker backend.
    ///
    /// If `safe` is true but `parallel_aware` is false (and `workers` is 0), the path
    /// can still run inside a worker (e.g. inner side of a Parallel Hash Join), but
    /// it will be a "replicated" execution where the worker performs the full scan.
    pub fn set_parallel_safe(mut self, safe: bool) -> Self {
        self.custom_path_node.path.parallel_safe = safe;
        self
    }

    /// Attach child paths; core plans them into `CustomScan.custom_plans` via
    /// `create_customscan_plan`.
    pub fn set_custom_paths(mut self, paths: PgList<pg_sys::Path>) -> Self {
        self.custom_path_node.custom_paths = paths.into_pg();
        self
    }

    /// Build a CustomPath using the given private data.
    ///
    /// `custom_private` can be used to store the custom path's private data.
    pub fn build(mut self, custom_private: CS::PrivateData) -> pg_sys::CustomPath {
        self.custom_path_node.custom_private = custom_private.into();
        self.custom_path_node.flags = self
            .flags
            .into_iter()
            .fold(0, |acc, flag| acc | flag as u32);

        self.custom_path_node
    }
}
