import json
import logging
from typing import Optional  # Python 3.7–3.9 compatibility
import boto3
from botocore.exceptions import ClientError

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

class DynamoDBDemo:
    def __init__(self, region: str):
        self.region = region
        self.dynamodb_resource = boto3.resource("dynamodb", region_name=self.region)
        self.client = self.dynamodb_resource.meta.client  # low-level client for waiters

    def wait_for_table_exists(self, table_name: str, delay: int = 5, max_attempts: int = 20) -> bool:
        try:
            waiter = self.client.get_waiter("table_exists")
            waiter.wait(TableName=table_name, WaiterConfig={"Delay": delay, "MaxAttempts": max_attempts})
            return True
        except ClientError as e:
            logging.error("Waiter table_exists failed: %s", e)
            return False  # [web:1]

    def create_table(self, table_name, key_schema, attribute_definitions, provisioned_throughput):
        try:
            print(f"\nCreating table {table_name} ...")
            self.dynamodb_resource.create_table(
                TableName=table_name,
                KeySchema=key_schema,
                AttributeDefinitions=attribute_definitions,
                ProvisionedThroughput=provisioned_throughput,
            )
            if self.wait_for_table_exists(table_name):
                print(f"Table {table_name} is ACTIVE.")
                return True
            return False
        except ClientError as e:
            logging.error(e)
            return False  # [web:3]

    def store_an_item(self, table_name, item):
        try:
            print(f"\nStoring item in table {table_name} ...")
            table = self.dynamodb_resource.Table(table_name)
            table.put_item(Item=item)
            print("Item stored.")
            return True
        except ClientError as e:
            logging.error(e)
            return False  # [web:59]

    def get_an_item(self, table_name, key):
        try:
            print(f"\nRetrieving item with key from table {table_name} ...")
            table = self.dynamodb_resource.Table(table_name)
            response = table.get_item(Key=key)
            item = response.get("Item")
            if item:
                print("Item found:")
                print(json.dumps(item, indent=2, default=str))
                return True
            else:
                print("Item not found.")
                return False
        except ClientError as e:
            logging.error(e)
            return False  # [web:3]

    def describe_table(self, table_name: str):
        try:
            resp = self.client.describe_table(TableName=table_name)
            print(json.dumps(resp["Table"], indent=2, default=str))
            return True
        except ClientError as e:
            logging.error(e)
            return False  # [web:3]

    def list_tables(self):
        try:
            tables = []
            paginator = self.client.get_paginator("list_tables")
            for page in paginator.paginate():
                tables.extend(page.get("TableNames", []))
            if tables:
                print("Tables:")
                for t in tables:
                    print(f"- {t}")
            else:
                print("No tables found.")
            return True
        except ClientError as e:
            logging.error(e)
            return False  # [web:3]

    def update_item(
        self,
        table_name: str,
        key: dict,
        update_expression: str,
        expr_attr_values: Optional[dict] = None,
        expr_attr_names: Optional[dict] = None,
        condition_expression: Optional[str] = None,
        return_values: str = "UPDATED_NEW",
    ):
        """
        Update an item using DynamoDB UpdateExpression with optional conditions.
        ReturnValues: NONE | ALL_OLD | UPDATED_OLD | ALL_NEW | UPDATED_NEW.
        """
        try:
            print(f"\nUpdating item in table {table_name} ...")
            table = self.dynamodb_resource.Table(table_name)
            params = {
                "Key": key,
                "UpdateExpression": update_expression,
                "ReturnValues": return_values,
            }
            if expr_attr_values:
                params["ExpressionAttributeValues"] = expr_attr_values
            if expr_attr_names:
                params["ExpressionAttributeNames"] = expr_attr_names
            if condition_expression:
                params["ConditionExpression"] = condition_expression

            resp = table.update_item(**params)
            attrs = resp.get("Attributes")
            if attrs:
                print("Returned attributes:")
                print(json.dumps(attrs, indent=2, default=str))
            else:
                print("Update succeeded. No attributes returned.")
            return True
        except ClientError as e:
            logging.error(e)
            return False  # [web:39][web:48]

    def delete_item(
        self,
        table_name: str,
        key: dict,
        condition_expression: Optional[str] = None,
        expr_attr_names: Optional[dict] = None,
        expr_attr_values: Optional[dict] = None,
        return_values: str = "ALL_OLD",
    ):
        """
        Delete a single item by key with optional ConditionExpression.
        ReturnValues: NONE | ALL_OLD.
        """
        try:
            print(f"\nDeleting item from table {table_name} ...")
            table = self.dynamodb_resource.Table(table_name)
            params = {"Key": key, "ReturnValues": return_values}
            if condition_expression:
                params["ConditionExpression"] = condition_expression
            if expr_attr_names:
                params["ExpressionAttributeNames"] = expr_attr_names
            if expr_attr_values:
                params["ExpressionAttributeValues"] = expr_attr_values

            resp = table.delete_item(**params)
            old = resp.get("Attributes")
            if old:
                print("Deleted item (old attributes):")
                print(json.dumps(old, indent=2, default=str))
            else:
                print("Delete succeeded.")
            return True
        except ClientError as e:
            logging.error(e)
            return False  # [web:60][web:61]

    def delete_table(self, table_name: str):
        try:
            print(f"\nDeleting table {table_name} ...")
            self.client.delete_table(TableName=table_name)
            print("Delete requested. Table will transition to DELETING and then be removed.")
            return True
        except ClientError as e:
            logging.error(e)
            return False  # [web:70]


def prompt_json(prompt_text: str):
    print(prompt_text)
    print("Enter JSON (single line) and press Enter, or leave blank to cancel:")
    line = input("> ").strip()
    if not line:
        return None
    try:
        return json.loads(line)
    except json.JSONDecodeError as e:
        print(f"Invalid JSON: {e}")
        return None  # general JSON parsing guidance

def default_music_schema():
    key_schema = [
        {"AttributeName": "artist", "KeyType": "HASH"},
        {"AttributeName": "song", "KeyType": "RANGE"},
    ]
    attribute_definitions = [
        {"AttributeName": "artist", "AttributeType": "S"},
        {"AttributeName": "song", "AttributeType": "S"},
    ]
    provisioned_throughput = {"ReadCapacityUnits": 1, "WriteCapacityUnits": 1}
    return key_schema, attribute_definitions, provisioned_throughput  # [web:3]

def print_menu():
    print("\nDynamoDB Menu")
    print("1) Create table")
    print("2) Put item")
    print("3) Get item")
    print("4) Describe table")
    print("5) List tables")
    print("6) Update item")
    print("7) Delete item")
    print("8) Delete table")
    print("9) Exit")  # [web:62]

def main():
    print("DynamoDB Demo (menu-driven)")
    region = input("Enter AWS region (default us-east-1): ").strip() or "us-east-1"
    app = DynamoDBDemo(region)

    default_table = "music"

    while True:
        print_menu()
        choice = input("Select an option: ").strip()

        if choice == "1":
            table_name = input(f"Table name (default {default_table}): ").strip() or default_table
            use_default = (input("Use default music schema? [Y/n]: ").strip().lower() or "y") == "y"
            if use_default:
                key_schema, attr_defs, thrpt = default_music_schema()
            else:
                key_schema = prompt_json('Enter KeySchema JSON (e.g., [{"AttributeName":"artist","KeyType":"HASH"},{"AttributeName":"song","KeyType":"RANGE"}]): ')
                if key_schema is None:
                    continue
                attr_defs = prompt_json('Enter AttributeDefinitions JSON (e.g., [{"AttributeName":"artist","AttributeType":"S"},{"AttributeName":"song","AttributeType":"S"}]): ')
                if attr_defs is None:
                    continue
                thrpt = prompt_json('Enter ProvisionedThroughput JSON (e.g., {"ReadCapacityUnits":1,"WriteCapacityUnits":1}): ')
                if thrpt is None:
                    continue
            app.create_table(table_name, key_schema, attr_defs, thrpt)

        elif choice == "2":
            table_name = input(f"Table name (default {default_table}): ").strip() or default_table
            print("Enter item as JSON (must include the table's key attributes):")
            item = prompt_json('Example: {"artist":"Michael Jackson","song":"Billie Jean","album":"Thriller","length_seconds":294}')
            if item is None:
                continue
            app.store_an_item(table_name, item)

        elif choice == "3":
            table_name = input(f"Table name (default {default_table}): ").strip() or default_table
            print("Enter key as JSON:")
            key = prompt_json('Example: {"artist":"Michael Jackson","song":"Billie Jean"}')
            if key is None:
                continue
            app.get_an_item(table_name, key)

        elif choice == "4":
            table_name = input(f"Table name (default {default_table}): ").strip() or default_table
            app.describe_table(table_name)

        elif choice == "5":
            app.list_tables()

        elif choice == "6":
            table_name = input(f"Table name (default {default_table}): ").strip() or default_table
            print("Enter key JSON for the item to update:")
            key = prompt_json('Example: {"artist":"Michael Jackson","song":"Billie Jean"}')
            if key is None:
                continue

            simple = (input("Use simple SET update? [Y/n]: ").strip().lower() or "y") == "y"
            if simple:
                attr = input("Attribute to set (e.g., album): ").strip()
                val_json = prompt_json('Enter value as JSON (e.g., "Thriller" or 294 or ["a","b"]): ')
                if val_json is None:
                    continue
                update_expression = "SET #a = :v"
                expr_attr_names = {"#a": attr}
                expr_attr_values = {":v": val_json}
                return_values = input("Return values [UPDATED_NEW|ALL_NEW|NONE] (default UPDATED_NEW): ").strip() or "UPDATED_NEW"
                app.update_item(
                    table_name,
                    key,
                    update_expression,
                    expr_attr_values=expr_attr_values,
                    expr_attr_names=expr_attr_names,
                    return_values=return_values,
                )
            else:
                print("Enter full UpdateExpression (e.g., 'SET #a=:v1, play_count = if_not_exists(play_count, :z) + :inc REMOVE old_attr'):")
                update_expression = input("> ").strip()
                expr_attr_names = prompt_json('ExpressionAttributeNames JSON or blank (e.g., {"#a":"album"}): ') or None
                expr_attr_values = prompt_json('ExpressionAttributeValues JSON or blank (e.g., {":v1":"Thriller",":z":0,":inc":1}): ') or None
                condition_expression = input("ConditionExpression or blank (e.g., attribute_exists(artist) AND attribute_exists(song)): ").strip() or None
                return_values = input("Return values [NONE|ALL_OLD|UPDATED_OLD|ALL_NEW|UPDATED_NEW] (default UPDATED_NEW): ").strip() or "UPDATED_NEW"

                app.update_item(
                    table_name,
                    key,
                    update_expression,
                    expr_attr_values=expr_attr_values,
                    expr_attr_names=expr_attr_names,
                    condition_expression=condition_expression,
                    return_values=return_values,
                )

        elif choice == "7":
            table_name = input(f"Table name (default {default_table}): ").strip() or default_table
            print("Enter key JSON for the item to delete:")
            key = prompt_json('Example: {"artist":"Michael Jackson","song":"Billie Jean"}')
            if key is None:
                continue

            guarded = (input("Use a ConditionExpression to guard delete? [y/N]: ").strip().lower() or "n") == "y"
            condition_expression = None
            expr_attr_names = None
            expr_attr_values = None
            if guarded:
                condition_expression = input("ConditionExpression (e.g., attribute_exists(artist) AND attribute_exists(song)): ").strip() or None
                expr_attr_names = prompt_json('ExpressionAttributeNames JSON or blank (e.g., {"#a":"album"}): ') or None
                expr_attr_values = prompt_json('ExpressionAttributeValues JSON or blank (e.g., {":val": "Thriller"}): ') or None

            return_values = input("Return values [NONE|ALL_OLD] (default ALL_OLD): ").strip() or "ALL_OLD"
            app.delete_item(
                table_name,
                key,
                condition_expression=condition_expression,
                expr_attr_names=expr_attr_names,
                expr_attr_values=expr_attr_values,
                return_values=return_values,
            )

        elif choice == "8":
            table_name = input(f"Table name to delete (default {default_table}): ").strip() or default_table
            confirm = input(f"Type DELETE to confirm deletion of table {table_name}: ").strip()
            if confirm == "DELETE":
                app.delete_table(table_name)
            else:
                print("Delete cancelled.")

        elif choice == "9":
            print("Exiting.")
            break

        else:
            print("Invalid choice. Please select a valid option.")

if __name__ == "__main__":
    main()