// 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/>.

//! Physical optimizer rule that wires FFHelper instances from PgSearchScanPlan
//! into VisibilityFilterExec for ctid resolution.
//!
//! VisibilityFilterExec needs real ctids to check visibility, but when deferred
//! visibility is enabled, ctid columns hold packed DocAddresses. This rule
//! finds the PgSearchScanPlan that owns each ctid column and wires its FFHelper
//! into the VisibilityFilterExec so it can resolve the packed addresses itself.
//!
//! This is interior mutation only (Mutex-based wiring) — no structural plan changes.

use std::sync::Arc;

use datafusion::common::config::ConfigOptions;
use datafusion::common::{DataFusionError, Result};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;

use crate::index::fast_fields_helper::FFHelper;
use crate::postgres::customscan::joinscan::visibility_filter::VisibilityFilterExec;
use crate::scan::execution_plan::PgSearchScanPlan;

#[derive(Debug)]
pub struct VisibilityCtidResolverRule;

impl PhysicalOptimizerRule for VisibilityCtidResolverRule {
    fn name(&self) -> &str {
        "VisibilityCtidResolver"
    }

    fn schema_check(&self) -> bool {
        // Interior mutation only — no schema changes.
        true
    }

    fn optimize(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        _config: &ConfigOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        walk_plan(&plan)?;
        Ok(plan)
    }
}

/// Walk the plan tree. When we find a VisibilityFilterExec, wire FFHelpers
/// from matching PgSearchScanPlans in its subtree.
fn walk_plan(plan: &Arc<dyn ExecutionPlan>) -> Result<()> {
    if let Some(vis_exec) = plan.as_any().downcast_ref::<VisibilityFilterExec>() {
        for &(plan_pos, _) in vis_exec.plan_pos_oids() {
            let ffhelper =
                find_ffhelper_for_plan_position(plan.as_ref(), plan_pos).ok_or_else(|| {
                    DataFusionError::Internal(format!(
                        "VisibilityCtidResolverRule: no PgSearchScanPlan found \
                     for deferred ctid plan_position {plan_pos}"
                    ))
                })?;
            vis_exec.set_ctid_resolver(plan_pos, ffhelper);
        }
    }

    for child in plan.children() {
        walk_plan(child)?;
    }
    Ok(())
}

/// Search the subtree for a PgSearchScanPlan whose deferred ctid metadata matches
/// the given plan position. Returns its FFHelper if found.
fn find_ffhelper_for_plan_position(
    plan: &dyn ExecutionPlan,
    plan_position: usize,
) -> Option<Arc<FFHelper>> {
    if let Some(scan) = plan.as_any().downcast_ref::<PgSearchScanPlan>() {
        if scan.deferred_ctid_plan_position() == Some(plan_position) {
            return scan.ffhelper();
        }
    }

    for child in plan.children() {
        if let Some(helper) = find_ffhelper_for_plan_position(child.as_ref(), plan_position) {
            return Some(helper);
        }
    }

    None
}

#[cfg(any(test, feature = "pg_test"))]
#[pgrx::pg_schema]
mod tests {
    use super::find_ffhelper_for_plan_position;
    use std::sync::Arc;

    use arrow_schema::{Schema, SchemaRef};
    use pgrx::prelude::*;

    use crate::index::fast_fields_helper::FFHelper;
    use crate::query::SearchQueryInput;
    use crate::scan::execution_plan::PgSearchScanPlan;

    fn empty_schema() -> SchemaRef {
        Arc::new(Schema::empty())
    }

    #[pg_test]
    fn matches_scan_by_deferred_ctid_plan_position() {
        let ffhelper = Arc::new(FFHelper::empty());
        let scan = PgSearchScanPlan::new(
            vec![],
            empty_schema(),
            SearchQueryInput::All,
            None,
            Vec::new(),
            Some(ffhelper.clone()),
            0,
            Some(7),
        );

        let found = find_ffhelper_for_plan_position(&scan, 7)
            .expect("matching plan_position should find ffhelper");
        assert!(Arc::ptr_eq(&found, &ffhelper));
        assert!(find_ffhelper_for_plan_position(&scan, 6).is_none());
    }
}
