import os
import sys
from collections import OrderedDict
import psycopg2
import yaml

DEFAULT_PORT = 5432

def main():
    try:
        HOST = os.environ['PRODUCT_CATALOGUE_HOST']
        DB = os.environ['PRODUCT_CATALOGUE_DB']
        USER = os.environ['PRODUCT_CATALOGUE_USER']
        PASSWORD = os.environ['PRODUCT_CATALOGUE_PASSWORD']
        PORT = (
            int(os.environ['PRODUCT_CATALOGUE_PORT'])
            if 'PRODUCT_CATALOGUE_PORT' in os.environ
            else DEFAULT_PORT
        )
    except KeyError as e:
        print("ERROR: Environment variable not found: {}".format(e))
        sys.exit(1)

    with psycopg2.connect(host=HOST, dbname=DB, user=USER, password=PASSWORD, port=PORT) as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT DISTINCT stream, type, levtype, param::INTEGER FROM fields WHERE param != '' ORDER BY stream, type, levtype, param::INTEGER"
            )
            rows = cur.fetchall()

    index = OrderedDict()
    for row in rows:
        stream, type_, levtype, param = row
        key = (stream, type_, levtype)
        if key not in index:
            index[key] = []
        index[key].append(param)
        
    # Manually add type=tf parameters for PGEN
    index[("oper", "tf", "")] = [129, 999]
    index[("scda", "tf", "")] = [129, 999]
    index[("enfo", "tf", "")] = [129, 999]
    
    yaml_dump_data = []
    for key, vals in sorted(index.items()):
        if key[2]:
            yaml_dump_data.append(
                [dict(stream=key[0], type=key[1], levtype=key[2]), vals]
            )
        else:
            yaml_dump_data.append([dict(stream=key[0], type=key[1]), vals])

    with open("params.yaml", "w") as f:
        f.write(
            "# File automatically generated by %s\n# Do not edit\n\n"
            % (os.path.basename(__file__))
        )
        f.write(yaml.safe_dump(yaml_dump_data, default_flow_style=False))


if __name__ == "__main__":
    main()
