import logging
import boto3
import awswrangler as wr
import pandas as pd
import json
import re
import ast
from typing import Optional, Dict, Any, Union
from dataclasses import dataclass
from datetime import datetime, date
import uuid 
import pytz
import base64
import numpy as np
from decimal import Decimal
from gen3_validator.dict import DataDictionary
from acdc_aws_etl_pipeline.validate.validate import load_schema_from_s3_uri

logger = logging.getLogger(__name__)

@dataclass
class AthenaConfig:
    """
    Configuration class for dbt release info writing.

    Args:
        aws_region (str): AWS region for Athena and S3 operations.
        aws_profile (str): AWS profile to use for authentication.
        athena_s3_output (str): S3 location for Athena query output.
    """
    aws_region: str
    aws_profile: str
    athena_s3_output: str

    def as_dict(self) -> Dict[str, Any]:
        return {
            "aws_region": self.aws_region,
            "aws_profile": self.aws_profile,
            "athena_s3_output": self.athena_s3_output,
        }

# ----------------- Athena Helpers -----------------

class AthenaQuery:
    def __init__(self, athena_config):
        self.config = athena_config

    def _get_boto_session(self):
        """Creates a boto3 session with the configured region and optional profile."""
        region = self.config.aws_region
        profile = getattr(self.config, "aws_profile", None)
        logger.debug(
            f"Creating boto3 session with region: {region}, profile: {profile}"
        )
        if not region:
            logger.error("config region must be set for boto3 session.")
            raise RuntimeError("config region must be set")
        # Only require region; profile is optional
        if profile:
            session = boto3.Session(region_name=region, profile_name=profile)
        else:
            session = boto3.Session(region_name=region)
        logger.debug("boto3 session successfully created.")
        return session
    
    def list_tables(self, database: str) -> list:
        """
        List all table names in the specified Athena database using awswrangler.

        Args:
            database (str): The Athena database to list tables from.

        Returns:
            list: A list of table names (str).
        """
        boto3_session = self._get_boto_session()
        try:
            tables = wr.catalog.get_tables(database=database, boto3_session=boto3_session)
            table_names = [tbl['Name'] for tbl in tables]
            logger.info(f"Found {len(table_names)} tables in database '{database}'.")
            return table_names
        except Exception as e:
            logger.error(f"Error listing tables in Athena database '{database}': {e}", exc_info=True)
            raise

    def query_athena(self, sql: str, athena_database: str, ctas_approach: bool = True) -> pd.DataFrame:
        """Runs an Athena query and returns the results as a pandas DataFrame."""
        logger.info(f"Running Athena query: {sql}")
        boto3_session = self._get_boto_session()
        try:
            df = wr.athena.read_sql_query(
                sql=sql,
                boto3_session=boto3_session,
                database=athena_database,
                ctas_approach=ctas_approach,
                s3_output=self.config.athena_s3_output
            )
            logger.info(
                f"Athena query completed successfully. Returned {len(df)} rows."
            )
        except Exception as e:
            logger.error(f"Error running Athena query: {e}", exc_info=True)
            raise
        return df
    
    def create_release_table(self) -> None:
        """
        Create the release tracking table in the specified database if it does not already exist.

        The table is created as an Iceberg table with Parquet format and Snappy compression.
        The schema includes release_tag, model_name, db_name, snapshot_id, committed_at, inserted_at, and github_sha.

        Raises:
            Exception: If table creation fails.
        """
        release_db = 'acdc_dataops_metadata_db'
        release_table = 'releases'
        release_s3_location = 's3://acdc-dataops-metadata/'

        create_sql = f"""
            CREATE TABLE IF NOT EXISTS {release_db}.{release_table} (
                release_tag        STRING,
                model_name         STRING,
                db_name            STRING,
                snapshot_id        BIGINT,
                committed_at       TIMESTAMP,
                inserted_at        TIMESTAMP,
                github_sha         STRING
            )
            PARTITIONED BY (release_tag)
            LOCATION '{release_s3_location}'
            TBLPROPERTIES (
                'table_type'='ICEBERG',
                'format'='parquet',
                'write_compression'='snappy'
            )
        """
        logger.info(f"Ensuring release table exists: {release_db}.{release_table}")
        try:
            self.query_athena(create_sql, release_db, False)
            logger.info(f"Release table {release_db}.{release_table} created or already exists.")
        except Exception as e:
            logger.error(f"Failed to create release table {release_db}.{release_table}: {e}", exc_info=True)
            raise

    def insert_to_iceberg_table(self, df, table_name, athena_database):
        """Inserts row into table."""
        logger.info(
            f"Inserting DataFrame with {len(df)} rows into iceberg table '{table_name}' "
            f"in database '{athena_database}'."
        )
        boto3_session = self._get_boto_session()
        try:
            logger.info(f"Inserting to icerberg table with s3 output: {self.config.athena_s3_output}")
            temp_s3_path = f"{self.config.athena_s3_output}/temp/{uuid.uuid4()}/"
            wr.athena.to_iceberg(
                df=df,
                database=athena_database,
                table=table_name,
                boto3_session=boto3_session,
                workgroup='primary',
                temp_path=temp_s3_path
            )
            logger.info(
                f"Insert to iceberg table '{table_name}' successful."
            )
        except Exception as e:
            logger.error(
                f"Error inserting to iceberg table '{table_name}': {e}",
                exc_info=True
            )
            raise

    def find_db_for_model(self, model_name: str) -> Optional[str]:
        """
        Search all Athena databases for a table with the given name and return the database name if found.

        This function iterates through all available Athena databases and checks if a table with the specified
        model_name exists in any of them. If found, it returns the name of the database; otherwise, it returns None.

        Args:
            model_name (str): The name of the table/model to search for.
        Returns:
            Optional[str]: The name of the database containing the table, or None if not found.

        Raises:
            Exception: If there is an error fetching databases or tables from Athena.

        Example:
            >>> find_db_for_model("my_table", config)
            'my_database'
            >>> find_db_for_model("nonexistent_table", config)
            None

        Notes:
            - Requires AWS credentials and permissions to list Athena databases and tables.
            - Uses the awswrangler (wr) library.
            - The AWS region is determined from config.aws_region.

        """
        try:
            # Create a session and pass it to wrangler
            boto3_session = self._get_boto_session()
            databases = wr.catalog.databases(boto3_session=boto3_session)
            db_list = databases.get('Database', [])
            for db in db_list:
                try:
                    # Pass the session to other wrangler calls too
                    tables = wr.catalog.get_tables(database=db, boto3_session=boto3_session)
                    table_names = [t.get('Name') for t in tables]
                    if model_name in table_names:
                        return db
                except Exception as table_exc:
                    logger.warning(f"Could not fetch tables for database '{db}': {table_exc}")
                    continue
            return None
        except Exception as e:
            logger.error(f"Error searching for model '{model_name}' in Athena databases: {e}")
            return None


def convert_dataframe_types_for_json(df: pd.DataFrame) -> pd.DataFrame:
    """
    Converts DataFrame column types to JSON-serialisable formats.
    """
    df_copy = df.copy()
    for col in df_copy.columns:
        # Handle Decimal objects first, converting them to float
        if any(isinstance(x, Decimal) for x in df_copy[col].dropna()):
            df_copy[col] = df_copy[col].apply(
                lambda x: float(x) if isinstance(x, Decimal) else x
            )

        # Handle datetime-like types, converting to ISO strings
        if pd.api.types.is_datetime64_any_dtype(df_copy[col].dtype):
            df_copy[col] = df_copy[col].apply(
                lambda x: x.isoformat() if pd.notna(x) else None
            )
            continue 

        # For float columns (either original or from Decimal conversion)
        if pd.api.types.is_float_dtype(df_copy[col].dtype):
            # Only convert to object type if NaN is present
            if df_copy[col].isna().any():
                df_copy[col] = df_copy[col].astype(object).where(df_copy[col].notna(), None)
            else:
                # Ensures columns of Decimals without Nones become float
                df_copy[col] = df_copy[col].astype(float)

        # Convert numpy integers to standard Python integers
        elif pd.api.types.is_integer_dtype(df_copy[col].dtype):
            df_copy[col] = df_copy[col].astype('Int64')

        # Handle object columns that might contain Timestamps or other special types
        elif pd.api.types.is_object_dtype(df_copy[col].dtype):
            df_copy[col] = df_copy[col].apply(
                lambda x: x.isoformat() if isinstance(x, (pd.Timestamp, datetime))
                else None if pd.isna(x)
                else x
            )

    return df_copy

def replace_nan_with_none(obj):
    """
    Recursively traverses a dictionary or list and replaces float NaN values with None.
    """
    if isinstance(obj, dict):
        return {k: replace_nan_with_none(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [replace_nan_with_none(i) for i in obj]
    # Check if it's a float and is NaN
    elif isinstance(obj, float) and np.isnan(obj):
        return None
    else:
        return obj


def json_serialiser(obj):
    """
    Custom JSON serialiser for handling types that aren't natively JSON-serialisable.

    This improved version handles:
    - Null/NA values (None, np.nan, pd.NaT), which are converted to None (JSON null).
    - Date and time objects (datetime.date, datetime.datetime, pd.Timestamp),
      which are converted to ISO 8601 strings.
    - Decimal objects, which are converted to floats.
    - NumPy numeric types (integer and floating), which are converted to
      standard Python int and float types.
    - Bytes objects, which are Base64-encoded into a string for safe JSON transport.
    - Set and frozenset objects, which are converted to lists.
    """
    # This check is critical and must be first. It correctly handles
    # None, np.nan, and pd.NaT, converting them to None for JSON null.
    if pd.isna(obj):
        return None

    # Handle specific, non-native JSON types
    if isinstance(obj, (datetime, pd.Timestamp, date)):
        return obj.isoformat()
    if isinstance(obj, Decimal):
        return float(obj)
    if isinstance(obj, np.integer):
        return int(obj)
    if isinstance(obj, np.floating):
        # This handles numpy float types, including np.inf.
        # np.nan is already caught by the pd.isna() check above.
        return float(obj)
    if isinstance(obj, bytes):
        # JSON does not support bytes. A common practice is to encode them
        # as a Base64 string, which is safe for JSON.
        return base64.b64encode(obj).decode('utf-8')
    if isinstance(obj, (set, frozenset)):
        # Sets are not JSON serialisable; convert them to a list.
        return list(obj)

    # For any unhandled types, this will raise an error. This is the correct
    # protocol for a default function used with `json.dumps()`.
    raise TypeError(f"Type {type(obj)} not JSON serialisable")




class AthenaValidationWriter:
    """
    AthenaWriter encapsulates common operations for extracting tabular data from an Amazon Athena table,
    transforming it as needed, and preparing a JSON serialization suitable for downstream processing or archiving to S3.

    Responsibilities:
        - Retrieve the most recent study_id and snapshot_id from the specified Athena table.
        - Extract the entire contents of a table.
        - Clean and format specific columns, such as decoding any stringified dictionary values for submitter_id links.
        - Construct clean, consistent JSON output.

    Args:
        athena_config (AthenaConfig): Configuration for Athena session, including region, workgroup, and output S3 location.
        db_name (str): Name of the Athena database.
        table_name (str): Name of the Athena table.
    """
    def __init__(self, athena_config, db_name, table_name):
        """
        Initializes an AthenaWriter instance.

        Args:
            athena_config (AthenaConfig): AthenaConfig instance with AWS and Athena parameters.
            db_name (str): Database name in Athena.
            table_name (str): Table name in Athena.
        """
        self.athena_config = athena_config
        self.db_name = db_name
        self.table_name = table_name
        self.study_id = None
        self.snapshot_id = None
        logger.info(f"AthenaWriter initialized with db: {db_name}, table: {table_name}")


    def _get_latest_snapshot_id(self, return_commit_datetime: bool = False) -> str:
        """
        Retrieves the latest snapshot_id from the Iceberg/Athena table's $snapshots metadata.

        Returns:
            str: The most recent snapshot ID based on the newest committed_at timestamp.

        Notes:
            The method stores the result as self.snapshot_id for convenient reuse.
            If no results are present, this will raise an IndexError.
        """
        # Query using CAST to avoid bringing in unsupported timestamp(3) with time zone types,
        # instead cast everything to string (safe for Athena <-> pandas & avoids Hive errors)
        query = f"""
            SELECT
                CAST(snapshot_id AS VARCHAR) AS snapshot_id,
                CAST(committed_at AS VARCHAR) AS committed_at
            FROM "{self.db_name}"."{self.table_name}$snapshots"
            ORDER BY committed_at DESC
            LIMIT 1
        """
        athena_query = AthenaQuery(self.athena_config)
        result = athena_query.query_athena(sql=query, athena_database=self.db_name)

        if result.empty:
            logger.warning(
                f"No snapshot rows found for {self.db_name}.{self.table_name}"
            )
            self.snapshot_id = None
            if return_commit_datetime:
                return None, None
            return None

        snapshot_id = result['snapshot_id'].iloc[0]
        self.snapshot_id = snapshot_id
        logger.info(f"Retrieved snapshot_id: {snapshot_id}")

        if return_commit_datetime:
            commit_datetime = result['committed_at'].iloc[0]
            return snapshot_id, commit_datetime if commit_datetime is not None else None
        return snapshot_id

    def _get_full_table(self):
        """
        Retrieves the complete contents of the Athena table, optionally at a specific snapshot.

        If self.snapshot_id is set, queries the table at that snapshot id.
        Otherwise, retrieves the latest/current version.

        Returns:
            list[dict]: Each row as a dictionary where keys are column names.

        Notes:
            This fetches all rows, which may have memory implications for large tables.
        """
        table_ref = f'"{self.db_name}"."{self.table_name}"'
        if self.snapshot_id is not None:
            table_ref = f'{table_ref} FOR VERSION AS OF {self.snapshot_id}'
            logger.info(f"Querying table at snapshot_id: {self.snapshot_id}")
        else:
            logger.info("Querying table at latest/current version (no snapshot_id).")

        query = f"""
            SELECT *
            FROM {table_ref}
        """
        athena_query = AthenaQuery(self.athena_config)
        result = athena_query.query_athena(sql=query, athena_database=self.db_name, ctas_approach=False)
        logger.info(f"Retrieved {len(result)} rows from the table.")
        return result

    def _format_submitter_id_value(self, submitter_id_value: str) -> Union[dict, list, str]:
        """
        Attempts to parse a string as a dictionary or list of dictionaries,
        handling both JSON format and Python dict string representations.
        
        This handles cases where Athena returns:
        - Single dict as JSON: '{"submitter_id": "value"}'
        - Single dict as Python: "{'submitter_id': 'value'}"
        - List of dicts as JSON: '[{"submitter_id": "val1"}, {"submitter_id": "val2"}]'
        
        Args:
            submitter_id_value: The value from the raw Athena row.
        
        Returns:
            dict, list, or original value: Parsed dict/list if successfully parsed,
            else the input value unchanged.
        """
        if not isinstance(submitter_id_value, str):
            return submitter_id_value
        
        # Try JSON parsing first (handles both dicts and lists with escaped quotes)
        try:
            parsed = json.loads(submitter_id_value)
            # Validate that result is dict or list of dicts
            if isinstance(parsed, dict):
                return parsed
            elif isinstance(parsed, list):
                return parsed
        except (json.JSONDecodeError, ValueError):
            pass
        
        # Fallback to ast.literal_eval for Python-style dicts/lists
        dict_or_list_pattern = re.compile(r'^\s*[\[{]\s*["\']?submitter_id["\']?')
        if dict_or_list_pattern.match(submitter_id_value):
            try:
                parsed = ast.literal_eval(submitter_id_value)
                if isinstance(parsed, (dict, list)):
                    return parsed
            except (ValueError, SyntaxError):
                logger.warning(f"Could not parse string-like dictionary/list: {submitter_id_value}")
                pass
        
        return submitter_id_value



    def construct_json(self):
        """
        Constructs a JSON string from the Athena table data, adding and normalizing
        certain metadata fields (including consistent snapshot_id and study_id), and
        parsing stringified dict values for proper downstream JSON compatibility.

        Returns:
            str: Formatted JSON string with all records, ready for saving to file or S3.

        Steps:
            1. Fetch latest study_id and snapshot_id for the table.
            2. Retrieve all table rows.
            3. For each row, update/add 'snapshot_id' and 'study_id', and parse any
               stringified dict fields (e.g., submitter_id link columns).
            4. Serialize rows as a pretty JSON string.
        """
        logger.info("Getting latest snapshot_id")
        self._get_latest_snapshot_id()
        
        logger.info("Constructing JSON data from Athena table...")
        full_table = self._get_full_table()

        if hasattr(full_table, "to_dict"):
            full_table_json = full_table.to_dict(orient="records")
        else:
            logger.error("Expected a pandas DataFrame from _get_full_table(), but got type: %s", type(full_table))
            raise ValueError("Unable to convert table to JSON, not a DataFrame.")
        
        # Convert data types to JSON-serialisable formats
        logger.info("Converting DataFrame types for JSON serialisation...")
        full_table = convert_dataframe_types_for_json(full_table)

        full_table_json = full_table.to_dict(orient="records")

        for obj in full_table_json:
            try:
                # Convert string submitter_id link dicts to valid dict
                for key, value in list(obj.items()):
                    obj[key] = self._format_submitter_id_value(value)
            except Exception as e:
                logger.warning(f"Error processing row for JSON output: {e}")

        try:
            full_table_json = replace_nan_with_none(full_table_json)
            json_data = json.dumps(full_table_json, indent=4, default=json_serialiser)
            logger.info("JSON data construction complete.")
        except (TypeError, ValueError) as e:
            logger.error(f"Failed to serialise JSON: {e}")
            raise
        
        return json_data


class AthenaGoldWriter(AthenaValidationWriter):
    def __init__(self, athena_config, db_name, table_name):
        super().__init__(athena_config, db_name, table_name)
        self.study_id = None
        self.snapshot_id = None
        self.json_data = None
        logger.info(f"AthenaGoldWriter initialized with db: {db_name}, table: {table_name}")
    
    def construct_json(self) -> str:
        """
        This 'construct_json' is specialized for the AthenaGoldWriter.

        For Gold tables, the expected behavior is:
          - Retrieve the full gold table (via _get_full_table())
          - Attempt to parse stringified dict fields just as in the parent, but may need table-specific logic
          - Serialize as a pretty JSON string
        """
        logger.info("Getting latest snapshot_id")
        self._get_latest_snapshot_id()
        
        logger.info("Constructing JSON data from Athena GOLD table...")
        full_table = self._get_full_table()

        if hasattr(full_table, "to_dict"):
            full_table_json = full_table.to_dict(orient="records")
        else:
            logger.error("Expected a pandas DataFrame from _get_full_table(), but got type: %s", type(full_table))
            raise ValueError("Unable to convert table to JSON, not a DataFrame.")

        # Convert data types to JSON-serialisable formats
        logger.info("Converting DataFrame types for JSON serialisation (Gold)...")
        full_table = convert_dataframe_types_for_json(full_table)
        full_table_json = full_table.to_dict(orient="records")

        for obj in full_table_json:
            try:
                # Convert string submitter_id link dicts to valid dict, if present
                for key, value in list(obj.items()):
                    obj[key] = self._format_submitter_id_value(value)
            except Exception as e:
                logger.warning(f"Error processing GOLD row for JSON output: {e}")

        try:
            full_table_json = replace_nan_with_none(full_table_json)
            json_data = json.dumps(full_table_json, indent=4, default=json_serialiser)
            logger.info("GOLD JSON data construction complete.")
            self.json_data = json_data
        except (TypeError, ValueError) as e:
            logger.error(f"Failed to serialise GOLD JSON: {e}")
            raise

        return json_data


def generate_validation_id():
    """
    Generates a unique validation identifier string, based on the current date and time in Australian Eastern Time.

    Returns:
        str: A string formatted as "YYYYMMDDHHMMSS", e.g., "20240531173010".
    """
    # Get Australian Eastern timezone (automatically handles AEST/AEDT)
    australian_tz = pytz.timezone('Australia/Melbourne')
    
    # Get current time in Australian Eastern timezone
    current_date_time = datetime.now(australian_tz).strftime("%Y%m%d%H%M%S")
    validation_id = f"{current_date_time}"
    logger.info(f"Generated validation_id (Australian Eastern Time): {validation_id}")
    return validation_id

def write_validation_json_to_s3(s3_bucket,
                study_id,
                validation_id,
                table_name,
                snapshot_id,
                json_data):
    """
    Writes the supplied JSON data to an S3 bucket using the specified logical path and metadata.

    Args:
        s3_bucket (str): Name of the S3 bucket.
        study_id (str): Study identifier to include in the S3 key/path.
        validation_id (str): Unique validation run identifier string.
        table_name (str): Source Athena table name.
        snapshot_id (str): Table snapshot ID (for versioning).
        json_data (str): The JSON-serialized string to be uploaded.

    The S3 object key is formatted as:
        validation/study_id=<study_id>/validation_id=<validation_id>/table_name=<table_name>/snapshot_id=<snapshot_id>/<table_name>.json

    Side Effects:
        Uploads (puts) the JSON data to the given S3 bucket.

    Raises:
        Any exception raised by boto3.client('s3').put_object.

    Logging:
        Logs both intent and completion of S3 upload.
    """
    s3 = boto3.client('s3')
    s3_object_key = f"validation/study_id={study_id}/validation_id={validation_id}/table_name={table_name}/snapshot_id={snapshot_id}/{table_name}.json"
    logger.info(f"Writing JSON data to S3 bucket: {s3_bucket}, object key: {s3_object_key}")
    s3.put_object(Body=json_data, Bucket=s3_bucket, Key=s3_object_key)
    logger.info(f"Object created at s3://{s3_bucket}/{s3_object_key}")



def write_gold_json_to_s3(
    s3_bucket,
    study_id,
    table_name,
    snapshot_id,
    json_data,
):
    """
    Write JSON data for a "Gold" Athena table to S3 in a unique and logical path.

    This function uploads the supplied JSON string to the given S3 bucket, placing it under a "validation"
    directory structure specific to "gold" Athena tables. The S3 object key removes the "gold_" prefix from
    the table name for the actual file name and omits any validation ID in the path.

    Args:
        s3_bucket (str): S3 bucket name to upload to.
        study_id (str): The unique identifier for the study; appears in the S3 path.
        table_name (str): Athena gold table name (e.g., "gold_diagnosis"); used in path and for filename.
        snapshot_id (str): Identifier for the data snapshot/version used in the path.
        json_data (str): JSON string (already serialized), to upload.

    S3 Path Pattern:
        validation/study_id=<study_id>/table_name=<table_name>/snapshot_id=<snapshot_id>/<stripped_table_name>.json

        - <stripped_table_name> is the table_name with the "gold_" prefix removed.

    Example:
        gold_table_name = "gold_diagnosis"
        table_name = gold_table_name
        Actual file will be e.g.:
            validation/study_id=foo/table_name=gold_diagnosis/snapshot_id=bar/diagnosis.json

    Side Effects:
        Uploads JSON to S3 using boto3.

    Raises:
        Any exception from boto3's client('s3').put_object.

    Logging:
        Logs before and after the S3 put operation, including path details for traceability.
    """
    s3 = boto3.client('s3')

    filename = table_name
    if filename.startswith("gold_"):
        filename = filename.replace("gold_", "")

    # Remove study_id from filename if present
    if study_id in filename:
        filename = filename.replace(f"{study_id}_", "")
    else:
        logger.warn(f"Filename {filename} does not contain study_id {study_id}, writing filename as {filename}")

    s3_object_key = f"gold_jsons/study_id={study_id}/table_name={table_name}/snapshot_id={snapshot_id}/{filename}.json"
    logger.info(f"Writing JSON data to S3 bucket: {s3_bucket}, object key: {s3_object_key}")
    s3.put_object(Body=json_data, Bucket=s3_bucket, Key=s3_object_key)
    logger.info(f"Object created at s3://{s3_bucket}/{s3_object_key}")


def construct_data_import_order(s3_uri) -> list:
    schema_dict = load_schema_from_s3_uri(s3_uri)
    dd = DataDictionary(schema_dict)
    dd.schema = schema_dict
    dd.calculate_node_order()
    return dd.node_order

def write_release_jsons_to_s3(s3_bucket, release_id, study_id, table_name, json_data):
    """
    Write a JSON string to a specific S3 location for a given release and study.

    Args:
        s3_bucket (str): The S3 bucket where the file will be uploaded.
        release_id (str): Release identifier used in the S3 key path.
        study_id (str): Study identifier used in the S3 key path.
        table_name (str): Table name (used for naming the .json file).
        json_data (str): JSON data (as a string) to be uploaded.

    Returns:
        str: The output directory path in S3 where the file was written.

    Raises:
        Exception: Any exception raised by boto3.client('s3').put_object.

    Example:
        >>> output_dir = write_release_jsons_to_s3('my-bucket', 'release123', 'study1', 'gold_foo', '{"x":1}')
    """
    s3 = boto3.client('s3')
    output_dir = f"release_jsons/{release_id}/{study_id}"
    s3_object_key = f"{output_dir}/{table_name}.json"
    logger.info(f"Writing JSON data to S3 bucket: {s3_bucket}, object key: {s3_object_key}")
    s3.put_object(Body=json_data, Bucket=s3_bucket, Key=s3_object_key)
    logger.info(f"Object created at s3://{s3_bucket}/{s3_object_key}")
    return output_dir