import json
import os
import sys
import pg8000
from env_loader import load_env

# Read DB config from environment; provide sensible defaults
CONFIG = {
    'host': os.getenv('PG_HOST', '8.138.196.105'),
    'port': int(os.getenv('PG_PORT', '5432')),
    'database': os.getenv('PG_DATABASE', 'fs_law_risk'),
    'user': os.getenv('PG_USER', 'postgres'),
    'password': os.getenv('PG_PASSWORD', 'difyai123456'),
}

# Export file path can be overridden via env
OUTPUT = os.getenv('RISK_EXPORT_OUTPUT', 'risk_tables_export.json')

SQLS = {
    'risk_subject': "SELECT sub_id AS id, sub_name AS name FROM public.risk_subject ORDER BY sub_name;",
    'risk_permit': "SELECT per_id AS id, per_name AS name FROM public.risk_permit ORDER BY per_name;",
    'risk_sub_per': "SELECT sub_id, per_id FROM public.risk_sub_per;",
}

def fetch_all(cursor):
    cols = [d[0] for d in cursor.description]
    return [dict(zip(cols, row)) for row in cursor.fetchall()]

def main():
    load_env()
    try:
        conn = pg8000.connect(**CONFIG)
    except Exception as e:
        print('DB connect error:', e, file=sys.stderr)
        sys.exit(2)

    result = {}
    try:
        with conn.cursor() as cur:
            cur.execute(SQLS['risk_subject'])
            subjects = fetch_all(cur)
            cur.execute(SQLS['risk_permit'])
            permits = fetch_all(cur)
            cur.execute(SQLS['risk_sub_per'])
            rels = fetch_all(cur)
        conn.commit()
    finally:
        conn.close()

    # Build mapping: subject -> list of permit_ids
    sub_to_permit_ids = {}
    for r in rels:
        sub_id = r['sub_id']
        per_id = r['per_id']
        sub_to_permit_ids.setdefault(sub_id, set()).add(per_id)

    # Subjects with aggregated permit_ids
    subjects_out = []
    for s in subjects:
        subjects_out.append({
            'id': s['id'],
            'name': s['name'],
            'permit_ids': sorted(list(sub_to_permit_ids.get(s['id'], [])))
        })

    # Final JSON: keep full permit catalog (id+name), and subjects contain aggregated permit_ids
    result['risk_subject'] = subjects_out
    result['risk_permit'] = permits

    with open(OUTPUT, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print('Exported to', OUTPUT)

if __name__ == '__main__':
    main()