#!/usr/bin/env python3
import json
import os
import sqlite3
import sys
from typing import Iterable


def env(name: str, default: str = "") -> str:
    return os.getenv(name, default).strip().rstrip("/")


def infer_base_url(kind: str) -> str:
    hosts = env("DJANGO_ALLOWED_HOSTS") or env("CIE_AUTH_ALLOWED_HOSTS")
    items = [host.strip() for host in hosts.split(",") if host.strip() and host.strip() != "*"]
    for host in items:
        if kind in host:
            return f"https://{host}"
    defaults = {
        "fa": "http://federation-authority:8000",
        "provider": "http://provider:8002",
        "rp": "http://relying-party:8001",
    }
    return defaults[kind]


TA = env("DJANGO_TRUST_ANCHOR") or infer_base_url("fa")
PROVIDER_ISSUER = env("DJANGO_PROVIDER_ISSUER") or f"{infer_base_url('provider')}/oidc/op"
RP_BASE = env("DJANGO_RP_BASE_URL") or infer_base_url("rp")
RP_ISSUER = f"{RP_BASE}/oidc/rp"
PROVIDER_BASE = PROVIDER_ISSUER.removesuffix("/oidc/op")


def rewrite_string(value: str) -> str:
    return (
        value.replace("http://127.0.0.1:8000/oidc/op", TA + "/oidc/op")
        .replace("http://127.0.0.1:8001/oidc/rp", RP_ISSUER)
        .replace("http://127.0.0.1:8002/oidc/op", PROVIDER_ISSUER)
        .replace("http://127.0.0.1:8000", TA)
        .replace("http://127.0.0.1:8001", RP_BASE)
        .replace("http://127.0.0.1:8002", PROVIDER_BASE)
        .replace("http://provider:8002/oidc/op", PROVIDER_ISSUER)
        .replace("http://provider:8002", PROVIDER_BASE)
        .replace("http://relying-party:8001/oidc/rp", RP_ISSUER)
        .replace("http://relying-party:8001", RP_BASE)
        .replace("http://federation-authority:8000/oidc/op", TA + "/oidc/op")
        .replace("http://federation-authority:8000", TA)
    )


REPLACEMENTS = [
    ("http://127.0.0.1:8000/oidc/op", TA + "/oidc/op"),
    ("http://127.0.0.1:8001/oidc/rp", RP_ISSUER),
    ("http://127.0.0.1:8002/oidc/op", PROVIDER_ISSUER),
    ("http://127.0.0.1:8000", TA),
    ("http://127.0.0.1:8001", RP_BASE),
    ("http://127.0.0.1:8002", PROVIDER_BASE),
    ("http://provider:8002/oidc/op", PROVIDER_ISSUER),
    ("http://provider:8002", PROVIDER_BASE),
    ("http://relying-party:8001/oidc/rp", RP_ISSUER),
    ("http://relying-party:8001", RP_BASE),
    ("http://federation-authority:8000/oidc/op", TA + "/oidc/op"),
    ("http://federation-authority:8000", TA),
]


def deep_replace(value):
    if isinstance(value, dict):
        return {k: deep_replace(v) for k, v in value.items()}
    if isinstance(value, list):
        return [deep_replace(v) for v in value]
    if isinstance(value, str):
        return rewrite_string(value)
    return value


def replace_in_all_text_columns(cur, table: str):
    cols = cur.execute(f"PRAGMA table_info({table})").fetchall()
    text_cols = [c[1] for c in cols if c[2] and c[2].upper() in {"TEXT", "VARCHAR(255)", "VARCHAR(2048)", "VARCHAR(64)", "VARCHAR(33)", "JSON"}]
    for column in text_cols:
        for old, new in REPLACEMENTS:
            cur.execute(
                f"UPDATE {table} SET {column} = REPLACE({column}, ?, ?) WHERE {column} LIKE ?",
                (old, new, f"%{old}%"),
            )


def update_json_column(cur, table, pk, column):
    rows = cur.execute(f"SELECT id, {column} FROM {table}").fetchall()
    for row_id, raw in rows:
        if not raw:
            continue
        try:
            data = json.loads(raw)
        except Exception:
            continue
        updated = deep_replace(data)
        if updated != data:
            cur.execute(
                f"UPDATE {table} SET {column} = ? WHERE {pk} = ?",
                (json.dumps(updated, ensure_ascii=False), row_id),
            )


def update_text_url_column(cur, table, pk, column):
    rows = cur.execute(f"SELECT {pk}, {column} FROM {table}").fetchall()
    for row_id, raw in rows:
        if not raw:
            continue
        updated = rewrite_string(raw)
        if updated != raw:
            cur.execute(
                f"UPDATE {table} SET {column} = ? WHERE {pk} = ?",
                (updated, row_id),
            )


def main():
    if len(sys.argv) != 2:
        raise SystemExit("usage: normalize_db_urls.py <sqlite_db_path>")

    db_path = sys.argv[1]
    if not os.path.exists(db_path):
        raise SystemExit(0)

    conn = sqlite3.connect(db_path)
    try:
        cur = conn.cursor()

        tables = [
            row[0]
            for row in cur.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
            ).fetchall()
        ]

        for table in tables:
            replace_in_all_text_columns(cur, table)

        for table in (
            "spid_cie_oidc_entity_federationentityconfiguration",
            "spid_cie_oidc_entity_fetchedentitystatement",
            "spid_cie_oidc_entity_trustchain",
        ):
            if table not in tables:
                continue

            if table == "spid_cie_oidc_entity_federationentityconfiguration":
                update_text_url_column(cur, table, "id", "sub")
                update_json_column(cur, table, "id", "metadata")
                update_json_column(cur, table, "id", "authority_hints")
                update_json_column(cur, table, "id", "trust_marks")
                update_json_column(cur, table, "id", "trust_mark_issuers")
                update_json_column(cur, table, "id", "jwks_fed")
                update_json_column(cur, table, "id", "jwks_core")
            elif table == "spid_cie_oidc_entity_fetchedentitystatement":
                update_text_url_column(cur, table, "id", "sub")
                update_text_url_column(cur, table, "id", "iss")
                update_json_column(cur, table, "id", "statement")
                update_text_url_column(cur, table, "id", "jwt")
            elif table == "spid_cie_oidc_entity_trustchain":
                update_text_url_column(cur, table, "id", "sub")
                update_json_column(cur, table, "id", "chain")
                update_json_column(cur, table, "id", "metadata")
                update_json_column(cur, table, "id", "jwks")
                update_json_column(cur, table, "id", "parties_involved")
                update_json_column(cur, table, "id", "trust_marks")

        conn.commit()
    finally:
        conn.close()


if __name__ == "__main__":
    main()