// 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::cell::{Cell, RefCell};
use std::collections::BTreeSet;
use std::thread::LocalKey;

use pgrx::pg_sys;

use crate::api::HashMap;
use crate::postgres::customscan::aggregatescan::AggregateScan;
use crate::postgres::customscan::joinscan::JoinScan;
use crate::postgres::customscan::CustomScan;

#[derive(Default)]
pub struct WarningData {
    pub contexts: BTreeSet<String>,
    pub details: BTreeSet<String>,
    pub has_contexts: bool,
}

#[derive(Default)]
struct PlannerWarningState {
    /// Maps Category -> Generic warning message -> WarningData
    warnings: HashMap<String, HashMap<String, WarningData>>,
    /// Contexts (e.g., table aliases) that have successfully been planned in the current query planning phase.
    /// If a context is in this set, we suppress any future warnings for it across all categories.
    successful_contexts: BTreeSet<String>,
}

thread_local! {
    /// Global state for planner warnings and successes to deduplicate them across the query planning phase.
    static PLANNER_STATE: RefCell<PlannerWarningState> = RefCell::new(PlannerWarningState::default());
}

/// Trait to convert various types into a list of warning contexts (strings).
pub trait ToWarningContexts {
    fn to_warning_contexts(self) -> Vec<String>;
}

impl ToWarningContexts for () {
    fn to_warning_contexts(self) -> Vec<String> {
        Vec::new()
    }
}

impl ToWarningContexts for &Vec<String> {
    fn to_warning_contexts(self) -> Vec<String> {
        self.to_vec()
    }
}

impl ToWarningContexts for &[String] {
    fn to_warning_contexts(self) -> Vec<String> {
        self.to_vec()
    }
}

impl ToWarningContexts for &str {
    fn to_warning_contexts(self) -> Vec<String> {
        vec![self.to_string()]
    }
}

impl ToWarningContexts for String {
    fn to_warning_contexts(self) -> Vec<String> {
        vec![self]
    }
}

impl ToWarningContexts for Vec<String> {
    fn to_warning_contexts(self) -> Vec<String> {
        self
    }
}

/// Add a warning to be emitted at the end of the planning phase.
///
/// # Arguments
/// * `category` - The category of the warning (e.g., Join, Top K)
/// * `message` - The generic warning message (e.g., "JoinScan not used: query must have a LIMIT clause")
/// * `contexts` - Contexts (e.g., table aliases) to associate with this warning.
///   Supported types:
///   - `()`: No context (global warning)
///   - `&str`: Single context
///   - `&Vec<String>`: List of contexts
pub fn add_planner_warning<S1: Into<String>, S2: Into<String>, C: ToWarningContexts>(
    category: S1,
    message: S2,
    contexts: C,
) {
    add_detailed_planner_warning(category, message, contexts, Vec::<String>::new())
}

/// Add a detailed warning to be emitted at the end of the planning phase.
///
/// # Arguments
/// * `category` - The category of the warning (e.g., Join, Top K)
/// * `message` - The generic warning message (e.g., "JoinScan not used: query must have a LIMIT clause")
/// * `contexts` - Contexts (e.g., table aliases) to associate with this warning.
/// * `details` - Additional details (e.g. types) to union for this warning.
pub fn add_detailed_planner_warning<
    S1: Into<String>,
    S2: Into<String>,
    C: ToWarningContexts,
    D: IntoIterator<Item = String>,
>(
    category: S1,
    message: S2,
    contexts: C,
    details: D,
) {
    let ctxs = contexts.to_warning_contexts();
    let category_str = category.into();
    let message_str = message.into();
    let had_contexts = !ctxs.is_empty();

    let details_vec: Vec<String> = details.into_iter().collect();

    PLANNER_STATE.with(|state| {
        let mut state = state.borrow_mut();

        // Filter out contexts that have already been successfully planned
        let filtered_ctxs: Vec<String> = ctxs
            .into_iter()
            .filter(|c| !state.successful_contexts.contains(c))
            .collect();

        // If we originally had contexts, but all of them were filtered out because they were successful,
        // we should completely suppress this warning.
        if had_contexts && filtered_ctxs.is_empty() {
            return;
        }

        let category_entry = state.warnings.entry(category_str).or_default();
        let entry = category_entry.entry(message_str).or_default();
        if !filtered_ctxs.is_empty() {
            entry.has_contexts = true;
        }
        entry.contexts.extend(filtered_ctxs);
        entry.details.extend(details_vec);
    });
}

/// Mark contexts as successful and clear any previous warnings for them across all categories.
/// This is used when a successful plan is found for a set of tables, invalidating previous
/// warnings about why a plan couldn't be generated.
pub fn mark_contexts_successful<C: ToWarningContexts>(contexts: C) {
    let ctxs_to_remove = contexts.to_warning_contexts();
    if ctxs_to_remove.is_empty() {
        return;
    }

    PLANNER_STATE.with(|state| {
        let mut state = state.borrow_mut();

        for category_warnings in state.warnings.values_mut() {
            for warning_data in category_warnings.values_mut() {
                for ctx in &ctxs_to_remove {
                    warning_data.contexts.remove(ctx);
                }
            }
            // Optional: Clean up empty entries
            category_warnings.retain(|_, warning_data| {
                !warning_data.has_contexts || !warning_data.contexts.is_empty()
            });
        }

        state.successful_contexts.extend(ctxs_to_remove);
    });
}

pub fn clear_planner_warnings() {
    PLANNER_STATE.with(|state| {
        let mut state = state.borrow_mut();
        state.warnings.clear();
        state.successful_contexts.clear();
    });
}

thread_local! {
    /// Tracks whether the current execution is within an `EXPLAIN` statement.
    ///
    /// When `paradedb.planner_warnings` is set to `error`, queries that cannot use
    /// accelerated ParadeDB custom scans fail with an `ERROR`. However, running `EXPLAIN`
    /// on those queries should still succeed so users can inspect the fallback execution plan
    /// and diagnose optimization issues. When `IN_EXPLAIN` is `true`, `emit_planner_warnings()`
    /// downgrades errors to `WARNING`s.
    static IN_EXPLAIN: Cell<bool> = const { Cell::new(false) };
}

/// An RAII guard that sets `IN_EXPLAIN` to `true` if `is_explain` is true, and restores
/// the previous state when dropped (supporting nested utility statements and clean unwinds).
pub struct ExplainGuard {
    prev: bool,
}

impl ExplainGuard {
    /// Creates a new `ExplainGuard`. If `is_explain` is `true`, marks the thread-local
    /// state as being inside an `EXPLAIN` query.
    pub fn new(is_explain: bool) -> Self {
        let prev = IN_EXPLAIN.with(|cell| {
            let prev = cell.get();
            if is_explain {
                cell.set(true);
            }
            prev
        });
        Self { prev }
    }
}

impl Drop for ExplainGuard {
    fn drop(&mut self) {
        IN_EXPLAIN.with(|cell| cell.set(self.prev));
    }
}

/// Returns whether the current thread is currently planning or executing an `EXPLAIN` query.
pub fn is_in_explain() -> bool {
    IN_EXPLAIN.with(|cell| cell.get())
}

static mut PREV_EXPLAIN_ONE_QUERY_HOOK: pg_sys::ExplainOneQuery_hook_type = None;

#[allow(clippy::too_many_arguments)]
#[rustfmt::skip]
#[pgrx::pg_guard]
unsafe extern "C-unwind" fn paradedb_explain_one_query_hook(
    query: *mut pg_sys::Query,
    cursor_options: ::core::ffi::c_int,
    into: *mut pg_sys::IntoClause,
    es: *mut pg_sys::ExplainState,
    query_string: *const ::core::ffi::c_char,
    params: pg_sys::ParamListInfo,
    query_env: *mut pg_sys::QueryEnvironment,
) {
    let _explain_guard = ExplainGuard::new(true);

    if let Some(prev_hook) = PREV_EXPLAIN_ONE_QUERY_HOOK {
        prev_hook(query, cursor_options, into, es, query_string, params, query_env);
    } else {
        #[cfg(any(feature = "pg15", feature = "pg16"))]
        {
            let plan = pg_sys::pg_plan_query(query, query_string, cursor_options, params);
            pg_sys::ExplainOnePlan(
                plan,
                into,
                es,
                query_string,
                params,
                query_env,
                core::ptr::null(),
                core::ptr::null(),
            );
        }
        #[cfg(not(any(feature = "pg15", feature = "pg16")))]
        {
            pg_sys::standard_ExplainOneQuery(
                query,
                cursor_options,
                into,
                es,
                query_string,
                params,
                query_env,
            );
        }
    }
}

/// Register an `ExplainOneQuery_hook` to track when queries are being planned/executed under `EXPLAIN`.
///
/// This allows `emit_planner_warnings()` to downgrade `paradedb.planner_warnings = 'error'`
/// to warnings during `EXPLAIN`, allowing users to inspect the query plan.
pub unsafe fn register_explain_hook() {
    PREV_EXPLAIN_ONE_QUERY_HOOK = pg_sys::ExplainOneQuery_hook;
    pg_sys::ExplainOneQuery_hook = Some(paradedb_explain_one_query_hook);
}

pub fn emit_planner_warnings() {
    let mode = crate::gucs::planner_warnings();
    if mode == crate::gucs::PlannerWarnings::Off {
        return;
    }

    let in_explain = is_in_explain();

    PLANNER_STATE.with(|state| {
        let state = state.borrow();
        let mut messages = Vec::new();

        // Flatten the map to iterate over all messages regardless of category
        for category_warnings in state.warnings.values() {
            for (message, warning_data) in category_warnings.iter() {
                let mut output = message.clone();

                if !warning_data.details.is_empty() {
                    let details_label = if warning_data.details.len() > 1 {
                        "types"
                    } else {
                        "type"
                    };
                    let details_str = warning_data
                        .details
                        .iter()
                        .cloned()
                        .collect::<Vec<_>>()
                        .join(", ");
                    output = format!("{output} ({details_label}: {details_str})");
                }

                if warning_data.contexts.is_empty() {
                    messages.push(output);
                } else {
                    let label = if warning_data.contexts.len() > 1 {
                        "tables"
                    } else {
                        "table"
                    };

                    let context_str = warning_data
                        .contexts
                        .iter()
                        .cloned()
                        .collect::<Vec<_>>()
                        .join(", ");
                    messages.push(format!("{output} ({label}: {context_str})"));
                }
            }
        }

        if messages.is_empty() {
            return;
        }

        if mode == crate::gucs::PlannerWarnings::Error && !in_explain {
            if messages.len() == 1 {
                pgrx::error!("{}", messages[0]);
            } else {
                let combined = messages.join("\n- ");
                pgrx::error!("ParadeDB planner scan errors:\n- {combined}");
            }
        } else {
            for msg in messages {
                pgrx::warning!("{msg}");
            }
        }
    });
}

//
// Per-statement execution-time warnings.
//
// Unlike the planner warnings above, which are batched and flushed at the end of planning, these
// fire during execution and dedup against the currently-executing statement so a warning raised
// from many search-operator predicate nodes surfaces just once.
//

type StatementId = pg_sys::TimestampTz;

const NEVER: StatementId = pg_sys::TimestampTz::MIN;

thread_local! {
    static WARNED_SEQ_SCAN_AT: Cell<StatementId> = const { Cell::new(NEVER) };
    static WARNED_SPILLED_AT: Cell<StatementId> = const { Cell::new(NEVER) };
}

fn current_statement_id() -> StatementId {
    unsafe { pg_sys::GetCurrentStatementStartTimestamp() }
}

fn warn_once_per_statement(warned_at: &'static LocalKey<Cell<StatementId>>, message: &str) {
    let stmt_id = current_statement_id();
    if warned_at.with(|last| last.replace(stmt_id)) != stmt_id {
        pgrx::warning!("{message}");
    }
}

/// Whether the planner already warned (this statement) that a join or aggregate scan wasn't used.
/// Those warnings name a concrete alternative, so the sequential-scan/spill warnings below would
/// just be noise next to them. The planner state is populated during planning and only cleared at
/// the start of the next planning cycle, so it is still readable while the plan executes.
fn superseded_by_scan_warning() -> bool {
    let categories = [
        JoinScan::NAME.to_str().expect("scan NAME is valid UTF-8"),
        AggregateScan::NAME
            .to_str()
            .expect("scan NAME is valid UTF-8"),
    ];
    PLANNER_STATE.with(|state| {
        let state = state.borrow();
        categories.iter().any(|cat| {
            state
                .warnings
                .get(*cat)
                .is_some_and(|msgs| !msgs.is_empty())
        })
    })
}

/// Warn (once per statement) that a search-operator predicate is being applied as a per-row filter
/// over a sequential scan rather than via the index.
pub fn warn_sequential_scan() {
    if superseded_by_scan_warning() {
        return;
    }
    warn_once_per_statement(
        &WARNED_SEQ_SCAN_AT,
        "the table is being sequentially scanned for this query, so performance may be slow\n\
         if you are not sure why, please file an issue: https://github.com/paradedb/paradedb/issues/new/choose",
    );
}

/// Warn (once per statement) that the materialized filter match set exceeded `work_mem` and
/// spilled to a temporary file.
pub fn warn_filter_spilled() {
    if superseded_by_scan_warning() {
        return;
    }
    warn_once_per_statement(
        &WARNED_SPILLED_AT,
        "the query's filter match set exceeded work_mem and spilled to a temporary file; query performance may be degraded",
    );
}
