import re
import sys

# Function to convert z/OS Db2 DDL into RDS Db2 DDL
def convert_ddl(zos_ddl, target_schema, target_data_tablespace, target_index_tablespace):
    
    # Remove CREATE TABLESPACE statements from the z/OS DDL
    zos_ddl = re.sub(r"CREATE TABLESPACE.*?;", "", zos_ddl, flags=re.DOTALL)
    
    # Remove clauses specific to z/OS that are unnecessary in Db2 for Linux/Unix/Windows
    zos_ddl = re.sub(r"SET CURRENT.*?;", "", zos_ddl, flags=re.DOTALL)
    zos_ddl = re.sub(r"USING STOGROUP \w+", "", zos_ddl)  # Remove storage group clause
    zos_ddl = re.sub(r"BUFFERPOOL \w+", "", zos_ddl)  # Remove buffer pool clause
    zos_ddl = re.sub(r"PIECESIZE \d+ \w+", "", zos_ddl)  # Remove piecesize settings
    zos_ddl = re.sub(r"CCSID\s+\w+", "", zos_ddl)  # Remove CCSID clause
    zos_ddl = re.sub(r"FOR SBCS DATA", "", zos_ddl)  # Remove single-byte data clause
    zos_ddl = re.sub(r"WITHOUT TIME ZONE", "", zos_ddl)  # Remove time zone restrictions
    zos_ddl = re.sub(r"WITH RESTRICT ON DROP", "", zos_ddl)  # Remove drop restrictions
    zos_ddl = re.sub(r"PARTITION BY SIZE", "", zos_ddl)  # Remove partitioning by size
    zos_ddl = re.sub(r"AUDIT NONE", "", zos_ddl)  # Remove audit settings
    zos_ddl = re.sub(r"DATA CAPTURE NONE|DATA CAPTURE CHANGES", "", zos_ddl)  # Remove data capture settings
    zos_ddl = re.sub(r"APPEND NO", "", zos_ddl)  # Remove append settings
    zos_ddl = re.sub(r"FREEPAGE \d+", "", zos_ddl)  # Remove freepage settings
    zos_ddl = re.sub(r"PCTFREE \d+", "", zos_ddl)  # Remove percentage free settings
    zos_ddl = re.sub(r"GBPCACHE \w+", "", zos_ddl)  # Remove group buffer pool cache settings
    zos_ddl = re.sub(r"NOT PADDED", "", zos_ddl)  # Remove padding settings
    zos_ddl = re.sub(r"CLOSE NO|CLOSE YES", "", zos_ddl)  # Remove close settings
    zos_ddl = re.sub(r"COPY NO", "", zos_ddl)  # Remove copy settings
    zos_ddl = re.sub(r"DEFER NO", "", zos_ddl)  # Remove defer settings
    zos_ddl = re.sub(r"DEFINE YES", "", zos_ddl)  # Remove define settings
    zos_ddl = re.sub(r"VOLATILE|NOT VOLATILE", "", zos_ddl)  # Remove volatility settings
    zos_ddl = re.sub(r"ERASE\s+NO", "", zos_ddl)  # Remove erase settings
    zos_ddl = re.sub(r"PRIQTY\s+\d+", "", zos_ddl)  # Remove primary quantity settings
    zos_ddl = re.sub(r"SECQTY\s+-?\d+", "", zos_ddl)  # Remove secondary quantity settings
    zos_ddl = re.sub(r"CLUSTER|NOT CLUSTER", "", zos_ddl)  # Remove clustering settings
    zos_ddl = re.sub(r"COMMIT;", "", zos_ddl)  # Remove commit statements

    # Uncomment the following line if you want to remove GRANT statements as well
    # zos_ddl = re.sub(r"GRANT.*?;", "", zos_ddl, flags=re.DOTALL)

    # Replace schema name 
    if target_schema:
        zos_ddl = re.sub(r"(\b\w+\.)", target_schema + '.', zos_ddl)

    # Replace data and index tablespace names
    if target_data_tablespace and target_index_tablespace:
        zos_ddl = re.sub(r"IN\s+\w+\.\w+", f"IN {target_data_tablespace} INDEX IN {target_index_tablespace}", zos_ddl)

    # Add COMPRESS YES clause to all tables to enable compression
    zos_ddl = re.sub(r'(CREATE TABLE.*?)\s*;', r"\1\n    COMPRESS YES;", zos_ddl, flags=re.DOTALL | re.IGNORECASE)

    # Remove consecutive blank lines to make the output cleaner
    zos_ddl = re.sub(r'\n\s*\n', '\n', zos_ddl)

    return zos_ddl

# Main function to handle user input and call conversion logic
def main():
    # Check for the correct number of command line arguments
    if len(sys.argv) != 6:
         print("Invalid Syntax") 
         print("Usage: python zos_to_luw_ddl_conversion.py <zos_ddl_file> <target_schema> <target_data_tablespace> <target_index_tablespace> <output_file>")
         sys.exit(1)

    # Extract command line arguments
    zos_ddl_file = sys.argv[1]
    target_schema = sys.argv[2]
    target_data_tablespace = sys.argv[3]
    target_index_tablespace = sys.argv[4]
    out_file = sys.argv[5]

    # Read the z/OS DDL from the input file
    with open(zos_ddl_file, 'r') as file:
        zos_ddl = file.read()

    # Convert the z/OS DDL to Db2 for Linux DDL
    db2_linux_ddl = convert_ddl(zos_ddl, target_schema, target_data_tablespace, target_index_tablespace)
    
    # Output the converted DDL
    try:
        with open(out_file, 'w') as file:
            file.write(db2_linux_ddl)
        print(f"Converted DDL written to {out_file}")
    except IOError as e:
        print(f"Error writing to output file: {e}")
        sys.exit(1)

# Entry point of the script
if __name__ == "__main__":
    main()
