#  File:    load.py
#  Path:    apps/etl/src/airiskindex_etl/load.py
#  Project: AI Risk Index — airiskindex.io
#  Author:  Simon-Pierre Boucher
#  Contact: contact@spboucher.ai
#  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
#
#  Description: Load derived CSVs into Postgres (idempotent upserts).

"""Load derived O*NET + OEWS CSVs into Postgres (idempotent upserts).

Occupations are keyed by O*NET-SOC code; wages join by SOC prefix
("15-1252.00" -> "15-1252"). Demo-seeded rows are overwritten by real data
where codes collide, which is the desired direction of truth.
"""

from __future__ import annotations

import csv
import os
from pathlib import Path

import psycopg

REPO_ROOT = Path(__file__).resolve().parents[4]
DERIVED = REPO_ROOT / "data" / "derived" / "onet"


def read(name: str) -> list[dict[str, str]]:
    with (DERIVED / name).open(encoding="utf-8", newline="") as fh:
        return list(csv.DictReader(fh))


def main() -> None:
    dsn = os.environ["DATABASE_URL"]
    occupations = read("occupations.csv")
    tasks = read("tasks.csv")
    wages = {row["soc"]: row for row in read("wages.csv")}

    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        for occ in occupations:
            wage = wages.get(occ["code"].split(".")[0], {})
            median = wage.get("median_wage_cents") or None
            employment = wage.get("employment") or None
            cur.execute(
                """
                INSERT INTO "Occupation" (code, title, description, "medianWageCents",
                                          "wageCurrency", employment)
                VALUES (%s, %s, %s, %s, 'USD', %s)
                ON CONFLICT (code) DO UPDATE SET
                  title = EXCLUDED.title,
                  description = EXCLUDED.description,
                  "medianWageCents" = COALESCE(EXCLUDED."medianWageCents", "Occupation"."medianWageCents"),
                  employment = COALESCE(EXCLUDED.employment, "Occupation".employment)
                """,
                (occ["code"], occ["title"], occ["description"], median, employment),
            )

        for task in tasks:
            cur.execute(
                """
                INSERT INTO "Task" (id, "occupationCode", statement, importance)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET
                  statement = EXCLUDED.statement,
                  importance = EXCLUDED.importance
                """,
                (
                    task["task_id"],
                    task["code"],
                    task["statement"],
                    float(task["importance"]) if task["importance"] else None,
                ),
            )

        cur.execute('SELECT count(*) FROM "Occupation"')
        occ_count = cur.fetchone()[0]
        cur.execute('SELECT count(*) FROM "Task"')
        task_count = cur.fetchone()[0]

    print(f"loaded: {occ_count} occupations, {task_count} tasks (incl. pre-existing)")


if __name__ == "__main__":
    main()