from __future__ import annotations

import copy
from abc import ABC, abstractmethod
from datetime import datetime, timedelta
from typing import Any, Generic, TypeVar, cast

from py4j.java_gateway import JVMView, JavaObject
from pyspark.sql import DataFrame, SparkSession

from ai.chronon.api.ttypes import (
    GroupBy, Join, JoinPart, JoinSource, Query, Source, StagingQuery
)
from ai.chronon.pyspark.constants import PARTITION_COLUMN_FORMAT
from ai.chronon.repo.serializer import thrift_simple_json
from ai.chronon.utils import (
    __set_name as set_name, get_max_window_for_gb_in_days, output_table_name
)

# Define type variable for our executables
T = TypeVar('T', GroupBy, Join, StagingQuery)


class PySparkExecutable(Generic[T], ABC):
    """
    Abstract base class defining common functionality for executing features via PySpark.
    This class provides shared initialization and utility methods but does not
    define abstract execution methods, which are left to specialized subclasses.
    """

    def __init__(self, obj: T, spark_session: SparkSession):
        """
        Initialize the executable with an object and SparkSession.

        Args:
            obj: The object (GroupBy, Join, StagingQuery) to execute
            spark_session: The SparkSession to use for execution
        """
        self.obj: T = obj
        self.spark: SparkSession = spark_session
        self.jvm: JVMView = self.spark._jvm
        self.platform: PlatformInterface = self.get_platform()
        self.java_spark_session: JavaObject = self.spark._jsparkSession
        self.default_start_date: str = (
            datetime.now() - timedelta(days=8)
        ).strftime(PARTITION_COLUMN_FORMAT)
        self.default_end_date: str = (
            datetime.now() - timedelta(days=1)
        ).strftime(PARTITION_COLUMN_FORMAT)

    @abstractmethod
    def get_platform(self) -> PlatformInterface:
        """
        Get the platform interface for platform-specific operations.

        Returns:
            A PlatformInterface instance
        """
        pass

    def _update_query_dates(
        self, query: Query, start_date: str, end_date: str
    ) -> Query:
        """
        Update start and end dates of a query.

        Args:
            query: The query to update
            start_date: The new start date
            end_date: The new end date

        Returns:
            The updated query
        """
        query_copy = copy.deepcopy(query)
        query_copy.startPartition = start_date
        query_copy.endPartition = end_date
        return query_copy

    def _update_source_dates(
        self, source: Source, start_date: str, end_date: str
    ) -> Source:
        """
        Update start and end dates of a source.

        Args:
            source: The source to update
            start_date: The new start date
            end_date: The new end date

        Returns:
            The updated source
        """
        source_copy = copy.deepcopy(source)
        if source_copy.events and source_copy.events.query:
            source_copy.events.query = self._update_query_dates(
                cast(Query, source_copy.events.query), start_date, end_date)
        elif source_copy.entities and source_copy.entities.query:
            source_copy.entities.query = self._update_query_dates(
                cast(Query, source_copy.entities.query), start_date, end_date)
        return source_copy

    def _execute_underlying_join_sources(
        self, group_bys: list[GroupBy], start_date: str, end_date: str, step_days: int
    ) -> None:
        """
        Execute underlying join sources.

        Args:
            group_bys: List of GroupBy objects
            start_date: Start date for execution
            end_date: End date for execution
            step_days: Number of days to process in each step
        """
        joins_to_execute: list[Join] = []
        join_sources_to_execute_start_dates: dict[str, str] = {}

        for group_by in group_bys:
            group_by_join_sources: list[JoinSource] = [
                s.joinSource for s in cast(list[Source], group_by.sources)
                if s.joinSource and not s.joinSource.outputTableNameOverride
            ]

            if not group_by_join_sources:
                continue

            # Recall that records generated by the inner join are input events
            # for the outer join. Therefore in order to correctly aggregate the
            # outer join, your inner join needs to be run from
            # start_date - max_window_for_gb_in_days
            max_window_for_gb_in_days: int = get_max_window_for_gb_in_days(
                group_by)

            shifted_start_date = (
                datetime.strptime(start_date, PARTITION_COLUMN_FORMAT) -
                timedelta(days=max_window_for_gb_in_days)
            ).strftime(PARTITION_COLUMN_FORMAT)

            for js in group_by_join_sources:
                js_name: str | None = js.join.metaData.name

                if js_name is None:
                    raise ValueError(
                        f"Join source {js} does not have a name. "
                        "Was set_metadata called?"
                    )

                if js_name not in join_sources_to_execute_start_dates:
                    join_sources_to_execute_start_dates[js_name] = shifted_start_date
                    joins_to_execute.append(js.join)
                else:
                    existing_start_date: str = join_sources_to_execute_start_dates[
                        js_name
                    ]
                    join_sources_to_execute_start_dates[js_name] = min(
                        shifted_start_date, existing_start_date
                    )

        if not joins_to_execute:
            return

        self.platform.log_operation(
            f"Executing {len(joins_to_execute)} Join Sources"
        )

        for join in joins_to_execute:
            j_start_date: str = join_sources_to_execute_start_dates[
                join.metaData.name
            ]

            executable_join_cls: type[JoinExecutable] = (
                self.platform.get_executable_join_cls()
            )

            executable_join = executable_join_cls(join, self.spark)

            self.platform.log_operation(
                f"Executing Join Source {join.metaData.name} "
                f"from {j_start_date} to {end_date}"
            )
            _ = executable_join.run(
                start_date=j_start_date, end_date=end_date, step_days=step_days
            )

            output_table_name_for_js: str = output_table_name(
                join, full_name=True
            )
            self.platform.log_operation(
                f"Join Source {join.metaData.name} will be read from "
                f"{output_table_name_for_js}"
            )

        self.platform.log_operation("Finished executing Join Sources")

    def group_by_to_java(self, group_by: GroupBy) -> JavaObject:
        """
        Convert GroupBy object to Java representation.

        Args:
            group_by: The GroupBy object to convert

        Returns:
            Java representation of the GroupBy
        """
        json_representation: str = thrift_simple_json(group_by)
        java_group_by: JavaObject = self.jvm.ai.chronon.spark.PySparkUtils.parseGroupBy(
            json_representation
        )
        return java_group_by

    def join_to_java(self, join: Join) -> JavaObject:
        """
        Convert Join object to Java representation.

        Args:
            join: The Join object to convert

        Returns:
            Java representation of the Join
        """
        json_representation: str = thrift_simple_json(join)
        java_join: JavaObject = self.jvm.ai.chronon.spark.PySparkUtils.parseJoin(
            json_representation
        )
        return java_join

    def staging_query_to_java(self, staging_query: StagingQuery) -> JavaObject:
        """
        Convert StagingQuery object to Java representation.

        Args:
            staging_query: The StagingQuery object to convert

        Returns:
            Java representation of the StagingQuery
        """
        json_representation: str = thrift_simple_json(staging_query)
        java_staging_query: JavaObject = (
            self.jvm.ai.chronon.spark.PySparkUtils.parseStagingQuery(
                json_representation
            )
        )
        return java_staging_query


class GroupByExecutable(PySparkExecutable[GroupBy], ABC):
    """Interface for executing GroupBy objects"""

    def _update_source_dates_for_group_by(
        self, group_by: GroupBy, start_date: str, end_date: str
    ) -> GroupBy:
        """
        Update start and end dates of sources in GroupBy.

        The source start_date is shifted back by the maximum window size to ensure
        enough historical data is available for windowed aggregations. For example,
        if computing features for 20251201-20251202 with a 3-day window, the source
        needs data starting from 20251128 to correctly aggregate the 3-day window.

        Args:
            group_by: The GroupBy object to update
            start_date: The new start date for feature computation
            end_date: The new end date

        Returns:
            The updated GroupBy object
        """
        if not group_by.sources:
            return group_by

        # Shift start_date back by max window to ensure enough data for aggregations
        max_window_days: int = get_max_window_for_gb_in_days(group_by)
        shifted_start_date: str = (
            datetime.strptime(start_date, PARTITION_COLUMN_FORMAT) -
            timedelta(days=max_window_days)
        ).strftime(PARTITION_COLUMN_FORMAT)

        for i, source in enumerate(group_by.sources):
            group_by.sources[i] = self._update_source_dates(
                source, shifted_start_date, end_date
            )
        return group_by

    def run(
        self,
        start_date: str | None = None,
        end_date: str | None = None,
        step_days: int = 30,
        skip_execution_of_underlying_join: bool = False
    ) -> DataFrame:
        """
        Execute the GroupBy object.

        Args:
            start_date: Start date for the execution (format: YYYYMMDD)
            end_date: End date for the execution (format: YYYYMMDD)
            step_days: Number of days to process in each step
            skip_execution_of_underlying_join: Whether to skip execution of
                underlying joins

        Returns:
            DataFrame with the execution results
        """
        effective_start_date = start_date or self.default_start_date
        effective_end_date = end_date or self.default_end_date

        self.platform.log_operation(
            f"Executing GroupBy {self.obj.metaData.name} from "
            f"{effective_start_date} to {effective_end_date} with step_days {step_days}"
        )
        self.platform.log_operation(
            f"Skip Execution of Underlying Join Sources: "
            f"{skip_execution_of_underlying_join}"
        )

        if not skip_execution_of_underlying_join:
            self._execute_underlying_join_sources(
                group_bys=[self.obj],
                start_date=effective_start_date,
                end_date=effective_end_date,
                step_days=step_days
            )

        # Prepare GroupBy for execution
        group_by_to_execute: GroupBy = copy.deepcopy(self.obj)
        group_by_to_execute.backfillStartDate = effective_start_date

        # Update sources with correct dates
        group_by_to_execute = self._update_source_dates_for_group_by(
            group_by_to_execute, effective_start_date, effective_end_date
        )

        # Get output table name
        group_by_output_table: str = output_table_name(
            group_by_to_execute, full_name=True
        )

        # GroupBy backfills don't store the semantic hash as a property in
        # the table the same way joins do.
        # Therefore we drop the backfill table to avoid data quality issues.
        self.platform.drop_table_if_exists(table_name=group_by_output_table)

        # Find starting point for log capture just before executing JVM calls
        log_token = self.platform.start_log_capture(
            f"Run GroupBy: {self.obj.metaData.name}"
        )

        try:
            # Convert to Java GroupBy
            java_group_by = self.group_by_to_java(group_by_to_execute)
            # Execute GroupBy
            result_df_scala: JavaObject = (
                self.jvm.ai.chronon.spark.PySparkUtils.runGroupBy(
                    java_group_by,
                    effective_end_date,
                    self.jvm.ai.chronon.spark.PySparkUtils.getIntOptional(
                        str(step_days)),
                    self.platform.get_table_utils()
                )
            )

            result_df = DataFrame(result_df_scala, self.spark)
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"GroupBy {self.obj.metaData.name} executed successfully "
                f"and was written to {group_by_output_table}"
            )
            return result_df
        except Exception as e:
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Execution failed for GroupBy {self.obj.metaData.name}: {str(e)}"
            )
            raise e

    def analyze(
        self,
        start_date: str | None = None,
        end_date: str | None = None,
        enable_hitter_analysis: bool = False
    ) -> None:
        """
        Analyze the GroupBy object.

        Args:
            start_date: Start date for analysis (format: YYYYMMDD)
            end_date: End date for analysis (format: YYYYMMDD)
            enable_hitter_analysis: Whether to enable hitter analysis
        """
        start_date = start_date or self.default_start_date
        end_date = end_date or self.default_end_date

        self.platform.log_operation(
            f"Analyzing GroupBy {self.obj.metaData.name} from "
            f"{start_date} to {end_date}"
        )
        self.platform.log_operation(
            f"Enable Hitter Analysis: {enable_hitter_analysis}"
        )

        # Prepare GroupBy for analysis
        group_by_to_analyze: GroupBy = copy.deepcopy(self.obj)

        # Update sources with correct dates
        group_by_to_analyze = self._update_source_dates_for_group_by(
            group_by_to_analyze, start_date, end_date
        )

        # Start log capture just before executing JVM calls
        log_token = self.platform.start_log_capture(
            f"Analyze GroupBy: {self.obj.metaData.name}"
        )

        try:
            # Convert to Java GroupBy
            java_group_by = self.group_by_to_java(group_by_to_analyze)
            # Analyze GroupBy
            self.jvm.ai.chronon.spark.PySparkUtils.analyzeGroupBy(
                java_group_by,
                start_date,
                end_date,
                enable_hitter_analysis,
                self.platform.get_table_utils()
            )
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"GroupBy {self.obj.metaData.name} analyzed successfully"
            )
        except Exception as e:
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Analysis failed for GroupBy {self.obj.metaData.name}: {str(e)}"
            )
            raise e


class JoinExecutable(PySparkExecutable[Join], ABC):
    """Interface for executing Join objects"""

    def _update_source_dates_for_join_parts(
        self, join_parts: list[JoinPart], start_date: str, end_date: str
    ) -> list[JoinPart]:
        """
        Update start and end dates of sources in JoinParts.

        Each JoinPart contains a GroupBy with windowed aggregations. The source
        start_date for each GroupBy is shifted back by its maximum window size
        to ensure enough historical data is available for aggregations.

        Args:
            join_parts: List of JoinPart objects
            start_date: The new start date for feature computation
            end_date: The new end date

        Returns:
            The updated list of JoinPart objects
        """
        if not join_parts:
            return []

        for jp in join_parts:
            # Shift start_date back by max window for this GroupBy
            max_window_days: int = get_max_window_for_gb_in_days(jp.groupBy)
            shifted_start_date: str = (
                datetime.strptime(start_date, PARTITION_COLUMN_FORMAT) -
                timedelta(days=max_window_days)
            ).strftime(PARTITION_COLUMN_FORMAT)

            for i, source in enumerate(jp.groupBy.sources):
                jp.groupBy.sources[i] = self._update_source_dates(
                    source, shifted_start_date, end_date
                )
        return join_parts

    def run(
        self,
        start_date: str | None = None,
        end_date: str | None = None,
        step_days: int = 30,
        skip_first_hole: bool = False,
        skip_execution_of_underlying_join: bool = False
    ) -> DataFrame:
        """
        Execute the Join object with Join-specific parameters.

        Args:
            start_date: Start date for the execution (format: YYYYMMDD)
            end_date: End date for the execution (format: YYYYMMDD)
            step_days: Number of days to process in each step
            skip_first_hole: Whether to skip the first hole in the join
            skip_execution_of_underlying_join: Whether to skip execution of
                underlying joins

        Returns:
            DataFrame with the execution results
        """
        start_date = start_date or self.default_start_date
        end_date = end_date or self.default_end_date

        self.platform.log_operation(
            f"Executing Join {self.obj.metaData.name} from "
            f"{start_date} to {end_date} with step_days {step_days}"
        )
        self.platform.log_operation(f"Skip First Hole: {skip_first_hole}")
        self.platform.log_operation(
            f"Skip Execution of Underlying Join: {skip_execution_of_underlying_join}"
        )

        # Prepare Join for execution
        join_to_execute = copy.deepcopy(self.obj)
        join_to_execute.left = self._update_source_dates(
            join_to_execute.left, start_date, end_date
        )

        if not skip_execution_of_underlying_join and self.obj.joinParts:
            self._execute_underlying_join_sources(
                [jp.groupBy for jp in join_to_execute.joinParts],
                start_date, end_date, step_days
            )

        # Update join parts sources
        join_to_execute.joinParts = self._update_source_dates_for_join_parts(
            join_to_execute.joinParts, start_date, end_date
        )

        # Get output table name
        join_output_table = output_table_name(join_to_execute, full_name=True)

        # Start log capture just before executing JVM calls
        log_token = self.platform.start_log_capture(
            f"Run Join: {self.obj.metaData.name}"
        )

        try:
            # Convert to Java Join
            java_join: JavaObject = self.join_to_java(join_to_execute)
            # Execute Join
            result_df_scala: JavaObject = self.jvm.ai.chronon.spark.PySparkUtils.runJoin(
                java_join,
                end_date,
                self.jvm.ai.chronon.spark.PySparkUtils.getIntOptional(
                    str(step_days)
                ),
                skip_first_hole,
                self.platform.get_table_utils()
            )

            result_df = DataFrame(result_df_scala, self.spark)
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Join {self.obj.metaData.name} executed successfully and "
                f"was written to {join_output_table}"
            )
            return result_df
        except Exception as e:
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Execution failed for Join {self.obj.metaData.name}: {str(e)}"
            )
            raise e

    def analyze(
        self,
        start_date: str | None = None,
        end_date: str | None = None,
        enable_hitter_analysis: bool = False
    ) -> None:
        """
        Analyze the Join object.

        Args:
            start_date: Start date for analysis (format: YYYYMMDD)
            end_date: End date for analysis (format: YYYYMMDD)
            enable_hitter_analysis: Whether to enable hitter analysis
        """
        effective_start_date = start_date or self.default_start_date
        effective_end_date = end_date or self.default_end_date

        self.platform.log_operation(
            f"Analyzing Join {self.obj.metaData.name} from "
            f"{effective_start_date} to {effective_end_date}"
        )
        self.platform.log_operation(
            f"Enable Hitter Analysis: {enable_hitter_analysis}"
        )

        # Prepare Join for analysis
        join_to_analyze: Join = copy.deepcopy(self.obj)
        join_to_analyze.left = self._update_source_dates(
            join_to_analyze.left, effective_start_date, effective_end_date
        )

        # Update join parts sources
        join_to_analyze.joinParts = self._update_source_dates_for_join_parts(
            join_to_analyze.joinParts, effective_start_date, effective_end_date
        )

        # Start log capture just before executing JVM calls
        log_token = self.platform.start_log_capture(
            f"Analyze Join: {self.obj.metaData.name}"
        )

        try:
            # Convert to Java Join
            java_join: JavaObject = self.join_to_java(join_to_analyze)
            # Analyze Join
            self.jvm.ai.chronon.spark.PySparkUtils.analyzeJoin(
                java_join,
                effective_start_date,
                effective_end_date,
                enable_hitter_analysis,
                self.platform.get_table_utils()
            )
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Join {self.obj.metaData.name} analyzed successfully"
            )

        except Exception as e:
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Analysis failed for Join {self.obj.metaData.name}: {str(e)}"
            )
            raise e


class StagingQueryExecutable(PySparkExecutable[StagingQuery], ABC):
    """
    Interface for executing StagingQuery objects.

    StagingQuery is used to materialize intermediate data transformations with
    SQL-based queries. Unlike GroupBy and Join which aggregate and join features,
    StagingQuery runs arbitrary SQL transformations with date template substitution.

    The query supports these templates:
    - {{ start_date }}: Start partition for the current computation range
    - {{ end_date }}: End partition for the current computation range
    - {{ latest_date }}: The end partition independent of computation range
    - {{ max_date(table=namespace.my_table) }}: Max partition of a given table
    """

    def run(
        self,
        end_date: str | None = None,
        step_days: int | None = None,
        enable_auto_expand: bool = True,
        override_start_partition: str | None = None,
        skip_first_hole: bool = True
    ) -> DataFrame:
        """
        Execute the StagingQuery object.

        Args:
            end_date: End date for the execution (format: YYYYMMDD).
                     Defaults to yesterday.
            step_days: Number of days to process in each step. None for single batch.
            enable_auto_expand: Whether to auto-expand the output table schema
                               if new columns are added (default: True)
            override_start_partition: Optional override for the start partition.
                                     If not provided, uses the StagingQuery's
                                     configured startPartition.
            skip_first_hole: Whether to skip the first hole in partition range
                            (default: True)

        Returns:
            DataFrame with the staging query results
        """
        end_date = end_date or self.default_end_date

        self.platform.log_operation(
            f"Executing StagingQuery {self.obj.metaData.name} "
            f"ending at {end_date}"
        )
        if step_days:
            self.platform.log_operation(f"Step Days: {step_days}")
        self.platform.log_operation(f"Enable Auto Expand: {enable_auto_expand}")
        if override_start_partition:
            self.platform.log_operation(
                f"Override Start Partition: {override_start_partition}"
            )
        self.platform.log_operation(f"Skip First Hole: {skip_first_hole}")

        # Prepare StagingQuery for execution
        staging_query_to_execute: StagingQuery = copy.deepcopy(self.obj)

        # Get output table name
        staging_query_output_table: str = staging_query_to_execute.metaData.outputTable

        # Start log capture just before executing JVM calls
        log_token = self.platform.start_log_capture(
            f"Run StagingQuery: {self.obj.metaData.name}"
        )

        try:
            # Convert to Java StagingQuery
            java_staging_query = self.staging_query_to_java(staging_query_to_execute)

            # Execute StagingQuery
            result_df_scala: JavaObject = (
                self.jvm.ai.chronon.spark.PySparkUtils.runStagingQuery(
                    java_staging_query,
                    end_date,
                    self.jvm.ai.chronon.spark.PySparkUtils.getIntOptional(
                        str(step_days) if step_days else None
                    ),
                    self.jvm.ai.chronon.spark.PySparkUtils.getBooleanOptional(
                        enable_auto_expand
                    ),
                    self.jvm.ai.chronon.spark.PySparkUtils.getStringOptional(
                        override_start_partition
                    ),
                    skip_first_hole,
                    self.platform.get_table_utils()
                )
            )

            result_df = DataFrame(result_df_scala, self.spark)
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"StagingQuery {self.obj.metaData.name} executed successfully "
                f"and was written to {staging_query_output_table}"
            )
            return result_df
        except Exception as e:
            self.platform.end_log_capture(log_token)
            self.platform.log_operation(
                f"Execution failed for StagingQuery "
                f"{self.obj.metaData.name}: {str(e)}"
            )
            raise e


class PlatformInterface(ABC):
    """
    Interface for platform-specific operations.

    This class defines operations that vary by platform (Databricks, Jupyter, etc.)
    and should be implemented by platform-specific classes.
    """

    def __init__(self, spark: SparkSession) -> None:
        """
        Initialize with a SparkSession.

        Args:
            spark: The SparkSession to use
        """
        self.spark = spark
        self.jvm = spark._jvm
        self.java_spark_session = spark._jsparkSession
        self.register_udfs()

    @abstractmethod
    def get_executable_join_cls(self) -> type[JoinExecutable]:
        """
        Get the class for executing joins.

        Returns:
            The class for executing joins
        """
        pass

    @abstractmethod
    def start_log_capture(self, job_name: str) -> Any:
        """
        Start capturing logs for a job.

        Args:
            job_name: The name of the job for log headers

        Returns:
            A token representing the capture state (platform-specific)
        """
        pass

    @abstractmethod
    def end_log_capture(self, capture_token: Any) -> None:
        """
        End log capturing and print the logs.

        Args:
            capture_token: The token returned from start_log_capture
        """
        pass

    def get_table_utils(self) -> JavaObject:
        """
        Get the table utils class that will be used for read/write operations
        on the JVM side. Can be overridden by subclasses to provide
        platform-specific implementations.

        Returns:
            A JavaObject representing the table utilities
        """
        return self.jvm.ai.chronon.spark.TableUtils(self.java_spark_session)

    def register_udfs(self) -> None:
        """
        Register UDFs for the self.platform.

        This method is intentionally left empty but not abstract, as some
        platforms may not need to register UDFs.

        Subclasses can override this method to provide platform-specific UDF
        registration.

        Pro tip: Both the JVM Spark Session and Python Spark Session use the
        same spark-sql engine. You can register Python UDFS and use them in
        the JVM, as well as vice-versa.
        At Stripe we currently only use Scala UDFs, so we include a JAR of our
        UDFs in the cluster and register them via:

        self.jvm.com.stripe.shepherd.udfs.ShepherdUdfRegistry.registerAllUdfs(
            self.java_spark_session
        )
        """
        pass

    def log_operation(self, message: str) -> None:
        """
        Log an operation.

        Args:
            message: The message to log
        """
        current_utc_time = datetime.utcnow()
        time_str = current_utc_time.strftime('[%Y-%m-%d %H:%M:%S UTC]')
        print(f'{time_str} {message}')

    def drop_table_if_exists(self, table_name: str) -> None:
        """
        Drop a table if it exists.

        Args:
            table_name: The name of the table to drop
        """
        _ = self.spark.sql(f"DROP TABLE IF EXISTS {table_name}")

    def set_metadata(
        self,
        obj: GroupBy | Join | StagingQuery,
        mod_prefix: str,
        name_prefix: str | None = None,
        output_namespace: str | None = None
    ) -> T:
        """
        Set the metadata for the object.

        Args:
            obj: The object to set metadata on
            mod_prefix: The root directory of where your features exist
            name_prefix: Optional prefix to add to the name
            output_namespace: Optional namespace to override the set output namespace

        Returns:
            The updated object
        """
        obj_type: type[GroupBy] | type[Join] | type[StagingQuery] = type(obj)

        # Handle object naming
        if not obj.metaData.name:
            try:
                set_name(obj=obj, cls=obj_type, mod_prefix=mod_prefix)
            except AttributeError:
                raise AttributeError(
                    "Please provide a name when defining "
                    + "group_bys/joins/staging_queries adhoc."
                )
        # We do this to avoid adding the name prefix multiple times
        elif obj.metaData.name and name_prefix:
            obj.metaData.name = obj.metaData.name.replace(f"{name_prefix}_", "")

        # Handle nested objects for GroupBy
        if obj_type == GroupBy:
            for s in obj.sources:
                if s.joinSource and not s.joinSource.join.metaData.name:
                    set_name(obj=s.joinSource.join, cls=Join, mod_prefix=mod_prefix)
        # Handle nested objects for Join
        elif obj_type == Join:
            for jp in obj.joinParts:
                for s in jp.groupBy.sources:
                    if s.joinSource and not s.joinSource.join.metaData.name:
                        set_name(
                            obj=s.joinSource.join, cls=Join, mod_prefix=mod_prefix
                        )

        # Set output namespace
        if output_namespace:
            obj.metaData.outputNamespace = output_namespace

            if obj_type == Join:
                for jp in obj.joinParts:
                    jp.groupBy.metaData.outputNamespace = output_namespace

        # Add user prefix to name
        if name_prefix:
            obj.metaData.name = f"{name_prefix}_{obj.metaData.name}"

        return obj
