# BSD 2-CLAUSE LICENSE

# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:

# Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# #ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# Original author: Reza Hosseini

import datetime
import os
import shutil
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import pandas as pd
import plotly.graph_objects as go

from abvelocity.journey.event.gen_event_query import MultiEventTable
from abvelocity.journey.event.gen_event_tables_query import (
    gen_event_tables_query,
    union_tables_query,
)
from abvelocity.journey.seq.conversion_query import CONVERSION_RATE, ConversionQuery
from abvelocity.journey.seq.seq_info import (
    CONSECUTIVE_DEDUPED,
    FULLY_DEDUPED,
    MAP,
    SEQ_INFO_LIST,
    UNDEDUPED,
    SeqInfo,
)
from abvelocity.journey.seq.seq_query import SeqQuery
from abvelocity.journey.viz.journey_barchart import JourneyBarchart
from abvelocity.journey.viz.sankey_plot import SankeyPlot
from abvelocity.journey.viz.sunburst_plot import SunburstPlot
from abvelocity.param.io_param import IOParam
from abvelocity.param.join_query import JoinQuery
from abvelocity.utils.plot_lines_markers import plot_lines_markers, plot_long_df
from abvelocity.utils.transform_time_query import TransformTimeQuery
from abvelocity.utils.write_queries import write_queries

# Default time column
TIME_COL = "time"

# Default event column
EVENT_COL = "event"

# Sequence column
SEQ_COL = "event_seq"

PLOT_VALUE_COLS = ["percent", "count"]

# Name for seq start time column
SEQ_START_TIME = "seq_start_time"

# Name for seq end time column
SEQ_END_TIME = "seq_end_time"

# Name for (total) seq count column (this can be count distinct as well)
SEQ_COUNT = "seq_count"

# Name for sequence length column (at journey level)
SEQ_LENGTH = "seq_length"

# Name for seq avg time taken column
AVG_SEQ_TIME = "avg_seq_time"

# Name for seq avg length column
AVG_SEQ_LENGTH = "avg_seq_length"

# Default name for adoption rate
ADOPTION_RATE = "adoption_rate"

# Agg time column (generated by transforming time)
AGG_TIME = "agg_time"


@dataclass
class CalcRes:
    """This is a dataclass to store results of computations which are used to potentially generate figures or summary data."""

    table_name: Optional[str] = None
    """table_name or an expression of a table view in SQL."""
    query: Optional[str] = None
    """The query to generate the data."""
    group_by_cols: Optional[List[str]] = None
    """The group by columns for the query."""
    group_by_cols_w_time: Optional[List[str]] = None
    """The group by columns including time. The disctinction is useful for timeseries analysis."""
    conditions: Optional[List[str]] = None
    """The conditions applied to generated the data."""
    df: Optional[pd.DataFrame] = None
    """The resulting dataframe from the query."""
    plot_df: Optional[pd.DataFrame] = None
    """A dataframe that is produced during plotting if df needs to be altered."""
    fig: Optional[go.Figure] = None
    """The resulting plotly figure."""


@dataclass
class Seq(ABC):
    """Dataclass to hold journey parameters."""

    event_table_name: str = ""
    """Event table name or a query to extract events.
    This query / table result includes a data schema of the form:
        `UNIT_COL`, `TIME_COL`, `EVENT_COL`, `self.partition_by_cols`

        - `UNIT_COL` (str) is typically the main unit of interest eg member for member data
        - `TIME_COl` (str / int) is a measure of time and if time is to be used to order later on, we will rely on its order
        - `EVENT_COL` (str) includes the events of interest
        - `self.partition_by_cols` This is a collection of columns to partition by in order to define journey boundaries, for example for usage
        data this could be a session id, or some kind of reference id.

    """
    event_queries_dict: Dict[str, str] = field(default_factory=dict)
    """Event queries. This is useful if multiple separate queries should be executed to construct a final event table.
    The choice to allow for multiple queries is as follows:
    These queries are to be executed by a platform and sometimes users need to give us a few queries to be run to make their
    set up simpler or due to computational issues (can materialize several tables in the middle).
    """
    create_table_prefix: str = ""
    """A prefix table name to be used in creating / storing tables."""
    partition_by_cols: List[str] = field(default_factory=list)
    """Columns to partition by and determine journey / seq boundaries."""
    start_date: str = ""
    """Start date used in event queries"""
    end_date: str = ""
    """End date used in event queries."""
    conditions: List[str] = field(default_factory=list)
    """An optional list of conditions to be applied to event queries."""
    seq_info_list: List[SeqInfo] = field(default_factory=lambda: SEQ_INFO_LIST)
    """The sequence types we want to generate with their output table name."""
    seq_queries_dict: Dict[str, str] = field(default_factory=dict)
    """Queries to generate seq data."""
    join_queries_dict: Dict[str, str] = field(default_factory=dict)
    """Queries to perform joins after initial seq tables are generated."""
    order_list: Optional[List[str]] = None
    """A custom order to use instead of time, time is still used when there are ties or label does not appear in `order_list`."""
    time_col_format: Optional[str] = "unix_ms"
    """The time format for the events data which also appears in seq_start_time and seq_end_time.
    Specifies the format of `seq_start_time` and `seq_end_time`.
        Valid options:
            - 'unix_ms': Unix timestamp in milliseconds (e.g., 1678886400000).
              Will use FROM_UNIXTIME(timestamp_col / 1000.0).
            - 'unix_s': Unix timestamp in seconds (e.g., 1678886400).
              Will use FROM_UNIXTIME().
            - 'string': String formatted timestamp (e.g., '2023-03-15 10:00:00').
              Will use CAST(... AS TIMESTAMP).
            - 'timestamp': Already a TIMESTAMP type. No casting applied.
            Defaults to 'unix_ms'.
    Note: This is not needed for seq table generation. It is only used when calling methods such as `get_seq_summary`.
    """
    time_unit: Optional[str] = "minutes"
    """ The time units for the aggregated times computed or just for conversions of times (to less granular ones).
    This unit will be used for the `DATE_TRUNC` function if `add_time_transform_col` is True.
    """
    delta_time_unit: Optional[str] = "minutes"
    """The time units for the delta's computed eg in `get_seq_summary` method.
    Can be 'seconds', 'minutes', 'hours', or 'days'. Defaults to 'minutes'.
    """
    post_joins: Optional[List[JoinQuery]] = None
    """A list of joins given in the `Join` dataclass format.
    These joins are meant to update the tables whose names are stored in the
    `output_table_name` attribute of the `SeqInfo` objects within `self.seq_info_list`.

    Here is how the joins work

        - The joins occur after the sequence tables are generated.
        - The joins includ the `right_table` only and as the `left_table`, we use each of the generated seq tables.
        - The joins are incremental. This means we do one join at a time and we use the resulting table as the `left_table`
          for the next join.
        - The intermediate joined tables have the same name as original table plus f"_join{i}" depending on the order.
        - The final joined table name is the same as the original table plus "_joined"
        - We update the name in the given attributes with the "_joined" version.
    """
    event_color_dict: Optional[Dict[str, str]] = None
    """Dictionary of events with their colors."""
    io_param: Optional[IOParam] = None
    """IO parameters to read / write tables and results."""
    # Note: `event_table_name` is already defined as a dataclass field above.
    # The line `self.event_table_name Optional[str] = None` was incorrect syntax
    # for a dataclass attribute and has been removed.
    multi_event_table: Optional[MultiEventTable] = None
    """This is the set of events we want to use to build the seq data.
    When this is passed, `gen_event_queries_via_multi_event_tables` can be used
    to implement `gen_event_queries` as here we are assuming that the `multi_event_table`
    has the information needed to build the event union table.
    """

    @abstractmethod
    def gen_event_queries(self, **kwargs) -> Dict[str, str]:
        """
        Abstract method to generate event-related queries and stores them in
        `self.event_queries_dict`.
        Note that this may not be needed if `self.event_table_name` is a prepared table.
        This will update the final event_table_name as well if needed.
        """
        pass

    def gen_event_queries_via_multi_event_tables(
        self, multi_event_tables: MultiEventTable
    ) -> Dict[str, str]:
        """
        Generates individual and union event queries from a MultiEventTable instance.
        This method populates `self.event_queries_dict` and `self.event_table_name`.
        It also populates `self.event_color_dict` if it's currently None.

        Args:
            multi_event_tables (MultiEventTable): An instance of MultiEventTable
                containing event definitions.

        Returns:
            Dict[str, str]: A dictionary where keys are event labels and 'union_query',
                and values are their respective SQL queries.
        """
        # Generate queries for individual event tables based on the MultiEventTable definition
        query_list = gen_event_tables_query(
            multi_event_table=multi_event_tables,
            start_date=self.start_date,
            end_date=self.end_date,
            create_table_prefix=self.create_table_prefix,
        )

        # Extract event labels in the correct order from the MultiEventTable's event_tables
        labels = [
            e.event_label for e in multi_event_tables.event_tables if e.event_label is not None
        ]

        # Map labels to their generated queries
        event_queries_dict = {label: query for label, query in zip(labels, query_list)}

        # Generate the union query for all events
        union_query = union_tables_query(
            multi_event_table=multi_event_tables, create_table_prefix=self.create_table_prefix
        )
        event_queries_dict["union_query"] = union_query

        # Update the Seq instance's event_table_name to reflect the union table
        self.event_table_name = f"{self.create_table_prefix}_union"

        # Logic: Update event_color_dict if it's currently None
        if self.event_color_dict is None:
            derived_event_color_dict = {
                event.event_label: event.event_color
                for event in multi_event_tables.event_tables
                if event.event_label is not None and event.event_color is not None
            }
            # Only assign if the derived dictionary is not empty (i.e., if colors were found)
            if derived_event_color_dict:
                self.event_color_dict = derived_event_color_dict

        # Store the generated queries in the instance attribute
        self.event_queries_dict = event_queries_dict

        return event_queries_dict

    def gen_seq_queries(self) -> Dict[str, str]:
        """
        Generates sequence-related queries and stores them in `self.seq_queries_dict`.
        There are two queries stored in the dict:

            - "consecutive_deduped_seq": This version will convert a journey of the form a, a, b, b, c, a -> [a, b, c, a]
                Which means it will eliminate consecutive occurrences of events.
            - "fully_deduped_seq": This version will convert the above journey to: [a, b, c]
                Which means it will only keep the first occurrence of each event.

        The expected output schema for both of these tables is as follows:

            `UNIT_COL`, "event_seq", `self.partition_by_cols`

        where:

            - `UNIT_COL` (str) is typically the main unit of interest eg member for member data
            -  "event_seq" (str) includes the journey / seq in array format eg [a, b, c]
            - `self.partition_by_cols` This is a collection of columns to partition by in order to define journey boundaries, for example for usage
            data this could be a session id, or some kind of reference id.

        """

        for seq_info in self.seq_info_list:
            if not seq_info.output_table_name:
                seq_info.output_table_name = (
                    f"{self.create_table_prefix}_{seq_info.deduping_method}_seq"
                )

        for seq_info in self.seq_info_list:
            deduping_method = seq_info.deduping_method
            self.seq_queries_dict[deduping_method] = SeqQuery(
                event_table_name=self.event_table_name,
                time_col=TIME_COL,
                event_col=EVENT_COL,
                output_table_name=seq_info.output_table_name,
                partition_by_cols=self.partition_by_cols,
                deduping_method=deduping_method,
                max_seq_index=seq_info.max_seq_index,
                order_list=self.order_list,
            ).gen()
            print(f"\n\n\n*** {seq_info.deduping_method}: {self.seq_queries_dict[deduping_method]}")

        return self.seq_queries_dict

    def gen_join_queries(self):
        """Generates join queries."""
        if not self.post_joins:
            return

        # Iterate through the collected table names and produce the join queries
        for seq_info in self.seq_info_list:
            table_name = seq_info.output_table_name
            # At first we use the original table_name as the `left_table`
            table_name_for_join = seq_info.output_table_name
            for i, join in enumerate(self.post_joins):
                # We assign the joined table so far to the `left_table`
                # Note that the `right_table` is already populated
                join.left_table = table_name_for_join
                # Determine the output table name for the current join
                # For the last table we use a simple name without index
                output_table_name = (
                    f"{table_name}_join{i}"
                    if i < len(self.post_joins) - 1
                    else f"{table_name}_joined"
                )
                join.output_table_name = output_table_name
                # Generate and store the SQL query
                query = join.gen()
                self.join_queries_dict[join.output_table_name] = query
                print(f"\n\n\n*** query: {table_name}:\n {query}")
                # The next join will use the newly joined table as `left_table`
                # This is because the join is incremental
                table_name_for_join = output_table_name

        # Reset (append suffix to) the original output table names
        for seq_info in self.seq_info_list:
            current_value = seq_info.output_table_name
            if current_value:
                seq_info.output_table_name = f"{current_value}_joined"

    def gen_sunburst_plots(
        self,
        conditions: List[str] = ["TRUE = TRUE"],
        count_distinct_col: Optional[str] = None,
        value_cols: List[str] = PLOT_VALUE_COLS,
        deduping_methods: List[str] = [UNDEDUPED, CONSECUTIVE_DEDUPED, FULLY_DEDUPED],
        additional_dimensions: Optional[List[str]] = None,
    ):
        """Generates and saves sunburst plots for the sequence tables.

        Args:
            conditions (List[str], optional): List of SQL WHERE clause conditions. Defaults to ["TRUE = TRUE"].
            count_distinct_col (str, optional): Column to use for COUNT(DISTINCT). Defaults to None.
            value_cols (List[str], optional): List of value columns to plot. Defaults to PLOT_VALUE_COLS.
            deduping_methods (List[str], optional): List of deduping methods to use. Defaults to [UNDEDUPED, CONSECUTIVE_DEDUPED, FULLY_DEDUPED].
            additional_dimensions (List[str], optional): Additional columns to group by in the plot. Defaults to None.
        """
        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")

        save_path0 = os.path.join(self.io_param.save_path, "sunburst")
        os.makedirs(save_path0, exist_ok=True)

        sunburst_io_param = IOParam(
            cursor=self.io_param.cursor,
            print_to_html=self.io_param.print_to_html,
            save_path=save_path0,
            file_name_suffix=self.io_param.file_name_suffix,
        )

        sunburst_plot = SunburstPlot(io_param=sunburst_io_param)

        res = {}
        for seq_info in self.seq_info_list:
            deduping_method = seq_info.deduping_method
            table_name = seq_info.output_table_name
            max_seq_index = seq_info.max_seq_index
            if deduping_method in deduping_methods and table_name:
                for value_col in value_cols:
                    print(
                        f"\n***: Generating sunbursts for table: {table_name}: conditions: {conditions}, quantity: {value_col}"
                    )
                    res[(deduping_method, value_col)] = sunburst_plot.gen(
                        table_name=table_name,
                        value_col=value_col,
                        conditions=conditions,
                        max_seq_index=max_seq_index,
                        count_distinct_col=count_distinct_col,
                        color_dict=self.event_color_dict,
                        additional_dimensions=additional_dimensions,
                    )
        return res

    def gen_sankey_plots(
        self,
        conditions: List[str] = ["TRUE = TRUE"],
        count_distinct_col: Optional[str] = None,
        add_end_state: bool = False,
        distinct_nodes_by_stage: bool = False,
        value_cols: List[str] = PLOT_VALUE_COLS,
        deduping_methods: List[str] = [UNDEDUPED, CONSECUTIVE_DEDUPED, FULLY_DEDUPED],
        additional_dimensions: Optional[List[str]] = None,
    ):
        """Generates and saves sankey plots for the sequence tables."""

        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")

        save_path0 = os.path.join(self.io_param.save_path, "sankey")
        os.makedirs(save_path0, exist_ok=True)

        sankey_io_param = IOParam(
            cursor=self.io_param.cursor,
            print_to_html=self.io_param.print_to_html,
            save_path=save_path0,
            file_name_suffix=self.io_param.file_name_suffix,
        )

        sankey_plot = SankeyPlot(io_param=sankey_io_param)

        res = {}
        for seq_info in self.seq_info_list:
            deduping_method = seq_info.deduping_method
            table_name = seq_info.output_table_name
            max_seq_index = seq_info.max_seq_index
            if deduping_method in deduping_methods and table_name:
                for value_col in value_cols:
                    print(
                        f"\n***: Generating sunbursts for table: {table_name}: conditions: {conditions}, quantity: {value_col}"
                    )
                    res[(deduping_method, value_col)] = sankey_plot.gen(
                        table_name=table_name,
                        value_col=value_col,
                        conditions=conditions,
                        max_seq_index=max_seq_index,
                        count_distinct_col=count_distinct_col,
                        color_dict=self.event_color_dict,
                        add_end_state=add_end_state,
                        distinct_nodes_by_stage=distinct_nodes_by_stage,
                        additional_dimensions=additional_dimensions,
                    )
        return res

    def gen_journey_barcharts(
        self,
        conditions: List[str] = ["TRUE = TRUE"],
        count_distinct_col: Optional[str] = None,
        sort_event_array: bool = False,
        deduping_methods: List[str] = [UNDEDUPED, CONSECUTIVE_DEDUPED, FULLY_DEDUPED, MAP],
    ):
        """Generates and saves journey barcharts for the sequence tables."""

        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")

        save_path0 = os.path.join(self.io_param.save_path, "barchart")
        os.makedirs(save_path0, exist_ok=True)

        barchart_io_param = IOParam(
            cursor=self.io_param.cursor,
            print_to_html=self.io_param.print_to_html,
            save_path=save_path0,
            file_name_suffix=self.io_param.file_name_suffix,
        )

        gen_barchart = JourneyBarchart(io_param=barchart_io_param)

        res = {}
        for seq_info in self.seq_info_list:
            deduping_method = seq_info.deduping_method
            table_name = seq_info.output_table_name
            if deduping_method in deduping_methods and table_name:
                print(
                    f"\n***: Generating barcharts for table: {table_name}: conditions: {conditions}"
                )
                is_map = deduping_method == MAP
                res[deduping_method] = gen_barchart.gen(
                    table_name=table_name,
                    conditions=conditions,
                    count_distinct_col=count_distinct_col,
                    sort_event_array=sort_event_array,
                    is_map=is_map,
                )

        return res

    def get_seq_info(self, deduping_method: Optional[str]) -> Optional[SeqInfo]:
        """Given a `deduping_method`, this returns seq_info with that deduping method.
        Args
        deduping_method (str): The `deduping_method` for calculating conversions.
            If not passed, we use `FULLY_DEDUPED`

        Returns:
            seq_info (SeqInfo): The corresponding `SeqInfo` instance.

        """
        if not deduping_method:
            deduping_method = FULLY_DEDUPED

        return next(
            (
                seq_info
                for seq_info in self.seq_info_list
                if seq_info.deduping_method == deduping_method
            ),
            None,  # Default to None if the method is not found
        )

    def aug_agg_time(self, table_name: str, group_by_cols: Optional[list[str]] = None):
        """
        Transforms the `table_name` to include the additional time column
        and adds the new column to `group_by_cols`.
        """
        if not self.time_col_format:
            raise ValueError(
                "Need to set `time_col_format` in `Seq` class to use `gen_seq_summary_query`"
            )

        table_name = TransformTimeQuery(
            table_name=table_name,
            original_time_col=SEQ_START_TIME,
            time_col_format=self.time_col_format,
            new_time_col=AGG_TIME,
            time_unit=self.time_unit,
        ).gen()

        # Create group_by_cols_w_time which includes group_by_cols and the new time col
        if group_by_cols:
            group_by_cols_w_time = group_by_cols + [AGG_TIME]
        else:
            group_by_cols_w_time = [AGG_TIME]

        return table_name, group_by_cols, group_by_cols_w_time

    def get_sql_delta_str(self) -> str:
        """Map common time units to the expected SQL literal for DATE_DIFF."""

        if not self.delta_time_unit:
            raise ValueError(
                "Need to set `delta_time_unit` in `Seq` class to use `gen_seq_summary_query`"
            )

        time_unit_mapping = {
            "second": "SECOND",
            "seconds": "SECOND",
            "minute": "MINUTE",
            "minutes": "MINUTE",
            "hour": "HOUR",
            "hours": "HOUR",
            "day": "DAY",
            "days": "DAY",
        }

        sql_delta_time_unit = time_unit_mapping.get(self.delta_time_unit.lower())
        if not sql_delta_time_unit:
            print(
                f"Warning: Invalid or unhandled delta_time_unit '{self.delta_time_unit}'. Falling back to 'MINUTE'."
            )
            sql_delta_time_unit = "MINUTE"

        # Prepare time columns based on time_col_format
        start_time_col_expr = SEQ_START_TIME
        end_time_col_expr = SEQ_END_TIME

        if self.time_col_format == "unix_ms":
            start_time_col_expr = f"FROM_UNIXTIME({SEQ_START_TIME} / 1000.0)"
            end_time_col_expr = f"FROM_UNIXTIME({SEQ_END_TIME} / 1000.0)"
        elif self.time_col_format == "unix_s":
            start_time_col_expr = f"FROM_UNIXTIME({SEQ_START_TIME})"
            end_time_col_expr = f"FROM_UNIXTIME({SEQ_END_TIME})"
        elif self.time_col_format == "string":
            start_time_col_expr = f"CAST({SEQ_START_TIME} AS TIMESTAMP)"
            end_time_col_expr = f"CAST({SEQ_END_TIME} AS TIMESTAMP)"
        # If `self.time_col_format`` is 'timestamp', no casting is needed, so expressions remain as is.

        sql_delta_str = f"AVG(DATE_DIFF('{sql_delta_time_unit}', {start_time_col_expr}, {end_time_col_expr})) AS {AVG_SEQ_TIME}"

        return sql_delta_str

    def gen_seq_summary_query(
        self,
        deduping_method: Optional[str] = None,
        count_distinct_col: Optional[str] = None,
        group_by_cols: Optional[List[str]] = None,
        group_by_time: Optional[bool] = False,
        conditions: Optional[List[str]] = None,
    ) -> CalcRes:
        """
        Generates a SQL query to summarize event sequences.

        Args:
            deduping_method (str, optional): The deduping_method for calculating conversions.
                If not passed, we use the preferred method via `get_seq_info`.
            count_distinct_col (str, optional): If provided, counts distinct values of this column
                instead of total event_seq count. Defaults to None.
            group_by_cols (list of str, optional): A list of column names to group the results by.
                Defaults to None.
            group_by_time (bool): If we should also group by time, in which case
                we transform the time column as per need before aggregation.
            conditions (List[str], optional): A list of SQL WHERE clause conditions (e.g., ["a = 'b'", "x = 1"]).
                These will be combined with AND. Defaults to None.

        Returns:
            str: The generated SQL query.
        """

        # Validate time_col_format
        allowed_time_formats = {"unix_ms", "unix_s", "string", "timestamp"}
        if self.time_col_format not in allowed_time_formats:
            print(
                f"Warning: Invalid time_col_format '{self.time_col_format}'. Falling back to 'unix_ms'."
            )
            self.time_col_format = "unix_ms"

        # Use get_seq_info to determine the table name based on fallback logic
        seq_info = self.get_seq_info(deduping_method)
        if not seq_info:
            raise ValueError(
                f"Could not find a sequence info for deduping method '{deduping_method}' or its fallbacks."
            )
        table_name = seq_info.output_table_name

        select_parts = []
        group_by_parts = []

        group_by_cols_w_time = group_by_cols
        if group_by_time:
            table_name, group_by_cols, group_by_cols_w_time = self.aug_agg_time(
                table_name, group_by_cols
            )

        # Handle Group By Columns
        if group_by_cols_w_time:
            select_parts.extend(group_by_cols_w_time)
            group_by_parts.extend(group_by_cols_w_time)

        # Count distinct units (eg members, journeys, ...)
        if count_distinct_col:
            select_parts.append(f"COUNT(DISTINCT {count_distinct_col}) AS {SEQ_COUNT}")

        # Calculate average seq_length (total events in a journey)
        select_parts.append(f"AVG({SEQ_LENGTH}) AS {AVG_SEQ_LENGTH}")

        sql_delta_str = self.get_sql_delta_str()

        # Calculate average duration using the prepared time column expressions
        select_parts.append(sql_delta_str)

        query = f"SELECT {', '.join(select_parts)}\nFROM {table_name}"

        # Add WHERE clause if conditions are provided
        if conditions:
            # Combine all conditions with ' AND '
            query += f"\nWHERE {' AND '.join(conditions)}"

        if group_by_parts:
            query += f"\nGROUP BY {', '.join(group_by_parts)}"

        return CalcRes(
            query=query,
            table_name=table_name,
            group_by_cols=group_by_cols,
            group_by_cols_w_time=group_by_cols_w_time,
        )

    def get_seq_summary(
        self,
        deduping_method: Optional[str] = None,
        count_distinct_col: Optional[str] = None,
        group_by_cols: Optional[List[str]] = None,
        group_by_time: Optional[bool] = False,
        conditions: Optional[List[str]] = None,
    ) -> CalcRes:
        """
        Generates a dataframe which summarize event sequences.

        Args:
            deduping_method (str, optional): The deduping_method for calculating conversions.
                If not passed, we use `FULLY_DEDUPED`.
            count_distinct_col (str, optional): If provided, counts distinct values of this column
                instead of total event_seq count. Defaults to None.
            group_by_cols (list of str, optional): A list of column names to group the results by.
                Defaults to None.
            group_by_time (bool): If we should also group by time, in which case
                we transform the time column as per need before aggregation.
            conditions (List[str], optional): A list of SQL WHERE clause conditions (e.g., ["a = 'b'", "x = 1"]).
                These will be combined with AND. Defaults to None.
        Returns:
            pd.DataFrame: The generated dataframe.
        """

        calc_res = self.gen_seq_summary_query(
            deduping_method=deduping_method,
            count_distinct_col=count_distinct_col,
            group_by_cols=group_by_cols,
            group_by_time=group_by_time,
            conditions=conditions,
        )

        query = calc_res.query

        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")

        df = self.io_param.cursor.get_df(query=query).df
        calc_res.df = df
        return calc_res

    def calc_conversion(
        self,
        numerator_list: Optional[List[str]] = None,
        denominator_list: Optional[List[str]] = None,
        require_all_numerator: bool = False,
        require_all_denominator: bool = False,
        count_distinct_col: Optional[str] = None,
        conditions: Optional[List[str]] = field(default_factory=list),
        group_by_cols: Optional[List[str]] = field(default_factory=list),
        group_by_time: Optional[bool] = None,
        deduping_method: Optional[str] = None,
        col_name: Optional[str] = CONVERSION_RATE,
    ):
        """
        Calculates conversion rates based on specified event sequences and returns
        the generated SQL query along with the resulting DataFrame.

        This method leverages the `ConversionQuery` class to construct a SQL query.
        The conversion rate is determined by the ratio of entities (or rows) that
        satisfy the numerator criteria versus those that satisfy the denominator criteria,
        based on the content of an array column (likely event sequences or tags).

        Args:
            numerator_list (Optional[List[str]]): A list of event tags or values that define the
                numerator for the conversion calculation. Entities (rows or distinct items)
                that have any of these tags will be counted. If `None`, all rows are
                considered for the numerator. Defaults to `None`.
            denominator_list (Optional[List[str]]): A list of event tags or values that define the
                denominator for the conversion calculation. Entities (rows or distinct items)
                that have any of these tags will be counted. If `None`, all rows are
                considered for the denominator. Defaults to `None`.
            require_all_numerator (bool): If True, all values in `numerator_list` must be
                present in `array_col` for numerator count. Defaults to False.
            require_all_denominator (bool): If True, all values in `denominator_list` must be
                present in `array_col` for denominator count. Defaults to False.
            count_distinct_col (Optional[str], optional): The column to use for
                `COUNT(DISTINCT)`. If provided (e.g., 'user_id' or 'session_id'), the
                conversion will be based on unique values in this column. If `None`, it
                defaults to counting all relevant rows. Defaults to `None`.
            conditions (Optional[List[str]], optional): A list of additional SQL WHERE
                conditions (e.g., "event_date >= '2023-01-01'"). These conditions will
                be combined with ' AND ' in the generated query. Defaults to an empty list.
            group_by_cols (Optional[List[str]], optional): A list of columns to group the
                conversion results by (e.g., 'country', 'DATE(event_timestamp)'). If
                provided, the query will return conversion rates per group. Defaults to
                an empty list.
            group_by_time (bool): If we should also group by time, in which case
                we transform the time column as per need before aggregation.
            deduping_method: The `deduping_method` for calculating conversions.
                If not passed, we use `FULLY_DEDUPED`
            col_name (Optional[str]): The column name for the conversion rate.


        Returns:
            Dict[str, Union[str, pd.DataFrame]]: A dictionary containing:
                - **'query' (str)**: The generated Presto SQL query string.
                - **'df' (pd.DataFrame)**: A Pandas DataFrame containing the results of
                  executing the query. This DataFrame will be empty if `self.io_param`
                  is `None`.

        Raises:
            ValueError: If `numerator_list` or `denominator_list` are empty (if not None).
                (Validation handled by `ConversionQuery`).
            TypeError: If `count_distinct_col` is not a string, or if `conditions`
                or `group_by_cols` are not lists of strings.
                (Validation handled by `ConversionQuery`).
        """

        table_name = self.get_seq_info(deduping_method).output_table_name

        if group_by_time:
            table_name, group_by_cols = self._aug_agg_time(table_name, group_by_cols)

        query = ConversionQuery(
            table_name=table_name,
            array_col=SEQ_COL,
            numerator_list=numerator_list,
            denominator_list=denominator_list,
            require_all_numerator=require_all_numerator,
            require_all_denominator=require_all_denominator,
            count_distinct_col=count_distinct_col,
            conditions=conditions,
            group_by_cols=group_by_cols,
            col_name=col_name,
        ).gen()

        df = None
        if self.io_param is not None:
            df = self.io_param.cursor.get_df(query).df

        return CalcRes(query=query, df=df)

    def calc_adoption(
        self,
        numerator_list: Optional[List[str]] = None,
        denominator_list: Optional[List[str]] = None,
        # Default changed to True for adoption, but can be overridden
        require_all_numerator: bool = True,
        require_all_denominator: bool = False,
        count_distinct_col: Optional[str] = None,
        conditions: Optional[List[str]] = field(default_factory=list),
        group_by_cols: Optional[List[str]] = field(default_factory=list),
        group_by_time: Optional[bool] = None,
        deduping_method: Optional[str] = None,
        col_name: Optional[str] = ADOPTION_RATE,
    ):
        """
        Calculates adoption rates based on specified event sequences and returns
        the generated SQL query along with the resulting DataFrame.

        This method is a wrapper around `calc_conversion`, with

            - `require_all_numerator` defaulting to `True`.
            - `col_name defaulting` to `ADOPTION_RATE`

        All other arguments are passed directly to `calc_conversion`.
        See that function for Args and Returns.
        """
        # Call calc_conversion directly, overriding the default for require_all_numerator
        return self.calc_conversion(
            numerator_list=numerator_list,
            denominator_list=denominator_list,
            require_all_numerator=require_all_numerator,  # This will use the default `True` or user's override
            require_all_denominator=require_all_denominator,
            count_distinct_col=count_distinct_col,
            conditions=conditions,
            group_by_cols=group_by_cols,
            group_by_time=group_by_time,
            deduping_method=deduping_method,
        )

    # Add `plot_lines_markers` function as a method to become available for usage w/o import
    def plot_lines_markers(self, **kwargs):
        return plot_lines_markers(**kwargs)

    # Add `plot_long_df` function as a method to become available for usage w/o import
    def plot_long_df(self, **kwargs):
        return plot_long_df(**kwargs)

    def plot_timeseries(
        self,
        calc_res: CalcRes,
        y_col: str,
        line_colors: Optional[List[str]] = None,
        title: Optional[str] = None,
    ):
        """Given calc_res it makes a timeseries plot."""
        x_col = AGG_TIME
        group_by_cols = calc_res.group_by_cols

        df = calc_res.df

        for col in [x_col, y_col]:
            if col not in df.columns:
                raise ValueError(f"{col} is not in {df.columns} when `plot_timeseries` is called.")

        if not title:
            title = f"{y_col} for {group_by_cols} across time."

        return plot_long_df(
            df=df,
            x_col=x_col,
            y_col=y_col,
            group_by_cols=group_by_cols,
            line_colors=line_colors,
            title=title,
        )

    def publish_queries(
        self,
        output_base_path: str,
        write_events: bool = True,
        write_sequences: bool = True,
        write_joins: bool = True,
    ):
        """
        Publishes selected generated SQL queries into structured subdirectories within the output path.

        Each type of query (event, sequence, join) gets its own numbered subdirectory
        (e.g., '01_event_queries'). Existing contents of selected subdirectories will be
        removed before new queries are written, ensuring a clean state.

        Args:
            output_base_path: The base directory where the query subdirectories will be created.
            write_events: If True, writes queries from self.event_queries_dict. Defaults to True.
            write_sequences: If True, writes queries from self.seq_queries_dict. Defaults to True.
            write_joins: If True, writes queries from self.join_queries_dict. Defaults to True.
        """
        # A list of tuples, each representing a query set with its write condition
        query_sets_config = [
            (write_events, "event_queries", self.event_queries_dict, "01_event_queries"),
            (write_sequences, "seq_queries", self.seq_queries_dict, "02_seq_queries"),
            (write_joins, "join_queries", self.join_queries_dict, "03_join_queries"),
        ]

        print(f"\n*** Publishing queries to: {output_base_path}")

        queries_written = False

        for should_write, description, queries_dict, subdir_name in query_sets_config:
            if should_write:
                full_output_dir = os.path.join(output_base_path, subdir_name)

                # Clean up the specific subdirectory before writing
                if os.path.exists(full_output_dir):
                    print(f"\n*** Cleaning existing directory: {full_output_dir}")
                    if os.path.isdir(full_output_dir):
                        shutil.rmtree(full_output_dir)
                    elif os.path.isfile(full_output_dir):
                        # This should ideally be prevented, but we handle it defensively
                        os.unlink(full_output_dir)
                        print(f"\n*** Removed conflicting file at: {full_output_dir}")

                # Only write if the dictionary for this query set is not empty
                if queries_dict:
                    print(f"\n*** Processing {description}...")
                    try:
                        # Assumes write_queries is imported and available
                        written_files = write_queries(
                            queries_dict=queries_dict, output_dir=full_output_dir
                        )
                        print(
                            f"\n*** Published {len(written_files)} {description.replace('_', ' ')} files to {full_output_dir}"
                        )
                        queries_written = True
                    except IOError as e:
                        print(
                            f"\n*** ERROR: Failed to publish {description.replace('_', ' ')}: {e}"
                        )
                else:
                    print(f"\n*** Skipping {description}: Dictionary is empty.")
            else:
                print(f"\n*** Skipping {description} as requested.")

        if not queries_written:
            print(
                "\n*** No queries were published. All selected dictionaries were empty or publishing was skipped."
            )
        else:
            print("\n*** All selected queries have been published.")

    def exec_event_queries(self):
        """Executes all event queries."""

        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")
        if self.event_queries_dict:
            for k, query in self.event_queries_dict.items():
                print(f"\n*** event query: {k} will be executed")
                self.io_param.cursor.execute_multi(query=query)

    def exec_seq_queries(self):
        """Executes all sequence queries."""

        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")

        if self.seq_queries_dict:
            for k, query in self.seq_queries_dict.items():
                print(f"\n*** seq query: {k} will be executed")
                self.io_param.cursor.execute_multi(query=query)

    def exec_join_queries(self):
        """Executes all post join queries."""

        if self.io_param is None:
            raise ValueError("io_param must be set to use this method.")

        if self.join_queries_dict:
            for k, query in self.join_queries_dict.items():
                print(f"\n*** seq query: {k} will be executed")
                self.io_param.cursor.execute_multi(query=query)

    def compute_window_end_date(self, window_days: int) -> str:
        """
        Compute the end date by adding a specified number of days to the object's start_date.

        The start_date should be in the format 'YYYY-MM-DD-00' or 'YYYY-MM-DD'.
        The returned end date will always be in the format 'YYYY-MM-DD-00'.

        Args:
            window_days (int): Number of days to add to self.start_date.

        Returns:
            str: The computed end date in the format 'YYYY-MM-DD-00' if it ends with '00' and 'YYYY-MM-DD' otherwise.

        Example:
            If self.start_date = '2024-06-01-00' and window_days = 7,
            the result will be '2024-06-08-00'.
        """
        base_date_str = (
            self.start_date.split("-00")[0] if self.start_date.endswith("-00") else self.start_date
        )

        base_date = datetime.datetime.strptime(base_date_str, "%Y-%m-%d")
        end_date = base_date + datetime.timedelta(days=window_days)
        return (
            end_date.strftime("%Y-%m-%d-00")
            if self.start_date.endswith("-00")
            else end_date.strftime("%Y-%m-%d")
        )
