# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
#
# DISCLAIMER: This code is provided as a sample/educational resource only.
# It is NOT intended for use in production environments without thorough review,
# testing, and adaptation to your specific requirements. AWS makes no warranties,
# express or implied, about the completeness, reliability, or suitability of this
# code for any particular purpose. Use at your own risk.
#
# Intended to run as a Glue 5.1 Spark job, as the Data-Analyst definer IAM role,
# in producer account that owns the definer role.
# Glue database and base tables are present in a central catalog account.
# The Glue catalog view also gets created in the central account.

import logging
import sys
import os
import boto3
import time
from datetime import datetime
from pyspark.sql import SparkSession

# Logging configuration
formatter = logging.Formatter('[%(asctime)s] %(levelname)s @ line %(lineno)d: %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)
handler.setFormatter(formatter)
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(handler)

catalog_name = "spark_catalog"

# Initialize Spark session
logger.info("Initializing Spark Session")
try:
    spark = (SparkSession
                .builder
                .appName("GlueView-BothDialects")
                .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
                .config(f"spark.sql.catalog.{catalog_name}", "org.apache.iceberg.spark.SparkSessionCatalog")
                .config(f"spark.sql.catalog.{catalog_name}.io-impl", "org.apache.iceberg.aws.s3.S3FileIO")
                .config(f"spark.sql.catalog.{catalog_name}.catalog", "org.apache.iceberg.aws.glue.GlueCatalog")
                .config(f"spark.sql.catalog.{catalog_name}.warehouse", "s3://<your-bucket-name>/<prefix>/transaction-table1")
                .config(f"spark.sql.catalog.{catalog_name}.glue.account-id", "<your-producer-account-id>")
                .config(f"spark.sql.catalog.{catalog_name}.client.region", "<your-region>")
                .config("spark.sql.defaultCatalog", catalog_name)
                .config(f"spark.sql.catalog.{catalog_name}.glue.skip-name-validation", "true")
                .getOrCreate())
except Exception as e:
    logger.error(f"Failed to initialize Spark session: {e}")
    raise

spark.sparkContext.setLogLevel("INFO")
logger.info("Spark session initialized successfully - running as Data-Analyst")

print('boto3 version = ', boto3.__version__)

try:
    glueclient = boto3.client('glue')
    stsclient = boto3.client("sts")
    catalog_id = stsclient.get_caller_identity()["Account"]
except Exception as e:
    logger.error(f"Failed to initialize AWS clients or retrieve account ID: {e}")
    raise

database_rl_name  = 'rl_bankdata_icebergdb'
original_db_name  = 'bankdata_icebergdb'
basetable1_name   = 'transaction_table1'
basetable2_name   = 'transaction_table2'
view_name         = 'view_fromglue'

## sql_query should refer to the original database name and table name, not resource links
sql_query = (
    "SELECT a.transaction_id, a.transaction_type, "
    "b.transaction_location "
    "FROM bankdata_icebergdb.transaction_table1 a "
    "RIGHT JOIN bankdata_icebergdb.transaction_table2 b "
    "ON a.transaction_id = b.transaction_id"
)

## you need backticks in SQL query if you have special characters or unusual notations
## in database or table names. Underscore is a valid character; no backticks needed.

try:
    logger.info("Showing available databases")
    spark.sql("show databases").show()

    spark.sql(f"show tables from {database_rl_name}").show()

    logger.info("Select from base table using resourcelink of db")
    spark.sql(f"select * from {database_rl_name}.{basetable1_name} limit 7").show()
    spark.sql(f"select * from {database_rl_name}.{basetable2_name} limit 7").show()
except Exception as e:
    logger.error(f"Failed during Spark SQL exploration queries: {e}")
    raise


## Spark SQL cannot be used to create a Glue multi-dialect/protected view.
## You need to use boto3 or the CLI.

try:
    logger.info(f"Creating Glue view '{view_name}' in database '{database_rl_name}'")
    response_createtable = glueclient.create_table(
        CatalogId=catalog_id,
        DatabaseName=database_rl_name,  ## use resource link name here, not original database name
        TableInput={
            'Name': view_name,
            'Description': 'ViewFromGlueBothDialects',
            'StorageDescriptor': {
                'Columns': [
                    {'Name': 'transaction_id',       'Type': 'string'},
                    {'Name': 'transaction_type',     'Type': 'string'},
                    {'Name': 'transaction_location', 'Type': 'string'}
                ]
            },
            'ViewDefinition': {
                'IsProtected': True,
                'Definer': 'arn:aws:iam::<your-producer-account-id>:role/Data-Analyst',
                'Representations': [
                    {
                        'Dialect': 'SPARK',
                        'DialectVersion': '1.0',
                        'ViewOriginalText': sql_query,
                        'ViewExpandedText': sql_query
                    },
                    {
                        'Dialect': 'ATHENA',
                        'DialectVersion': '3',
                        'ViewOriginalText': sql_query,
                        'ValidationConnection': 'glue-view-validation-connection',
                    }
                ],
                'SubObjects': [
                    'arn:aws:glue:<your-region>:<your-central-account-id>:table/bankdata_icebergdb/transaction_table1',  ## use original database name and table names of source account
                    'arn:aws:glue:<your-region>:<your-central-account-id>:table/bankdata_icebergdb/transaction_table2'
                ]
            }
        }
    )
    logger.info("create_table API call succeeded, polling for view readiness...")
except Exception as e:
    logger.error(f"Failed to create Glue view: {e}")
    raise


def wait_for_view_ready(glue_client, catalog_id, database_name, table_name,
                        poll_interval_seconds=15, max_attempts=8):
    """
    Poll Glue get_table() and check Table.Status.State until it reaches SUCCESS.
    Possible State values observed from the API: SUCCESS | FAILED | IN_PROGRESS.
    Returns True on SUCCESS, raises RuntimeError on FAILED or if max_attempts exceeded.
    """
    for attempt in range(1, max_attempts + 1):
        try:
            response = glue_client.get_table(
                CatalogId=catalog_id,
                DatabaseName=database_name,
                Name=table_name,
                IncludeStatusDetails=True
            )
            table = response.get('Table', {})
            status = table.get('Status', {})
            state = status.get('State')
            action = status.get('Action', 'UNKNOWN')
            logger.info(f"[Attempt {attempt}/{max_attempts}] View status — Action: {action}, State: {state}")

            if state == 'SUCCESS':
                logger.info(f"View '{table_name}' created successfully (Action: {action}, State: {state}).")
                return True

            if state == 'FAILED':
                details = status.get('Details', {})
                raise RuntimeError(
                    f"View '{table_name}' creation failed. Action: {action}, Details: {details}"
                )

            # state is IN_PROGRESS or absent — keep polling
            logger.info(f"View not ready yet (State: {state}), waiting {poll_interval_seconds}s...")
            time.sleep(poll_interval_seconds)

        except glue_client.exceptions.EntityNotFoundException:
            logger.info(f"[Attempt {attempt}/{max_attempts}] View not yet found in Glue catalog, waiting...")
            time.sleep(poll_interval_seconds)

    raise RuntimeError(
        f"View '{table_name}' did not reach SUCCESS after {max_attempts} attempts "
        f"({max_attempts * poll_interval_seconds}s total)."
    )


try:
    wait_for_view_ready(
        glue_client=glueclient,
        catalog_id=catalog_id,
        database_name=database_rl_name,
        table_name=view_name
    )
except RuntimeError as e:
    logger.error(f"View readiness check failed: {e}")
    raise
except Exception as e:
    logger.error(f"Unexpected error while polling for view readiness: {e}")
    raise


try:
    logger.info("-----Select from the created view-------")
    spark.sql(f"select * from {database_rl_name}.{view_name} limit 7").show()
except Exception as e:
    logger.error(f"Failed to query the created view: {e}")
    raise

logger.info("Glue view creation and query completed successfully.")