"""Cleanup: drop Supabase Activation rows for combos that exist in HF Hub.

Uses parallel workers for throughput. Each worker has its own DB connection.
"""
import json
import os
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Optional

from wisent.core.utils.config_tools.constants import (
    INDEX_FIRST,
    INDEX_LAST,
    COMBO_OFFSET,
    MIN_COMBO_PATH_PARTS,
    PG_STATEMENT_NO_LIMIT,
    RECURSION_INITIAL_DEPTH,
    PARALLEL_CLEANUP_WORKERS,
)
from wisent.core.reading.modules.utilities.data.sources.hf.hf_config import (
    HF_REPO_ID,
    HF_REPO_TYPE,
    safe_name_to_model,
)

_print_lock = threading.Lock()


def _get_db_connection(database_url: Optional[str] = None):
    """Get database connection."""
    import psycopg2
    from urllib.parse import urlparse, parse_qs, urlencode, urlunparse

    db_url = database_url or os.environ.get("DATABASE_URL")
    if not db_url:
        raise ValueError("No DATABASE_URL.")
    parsed = urlparse(db_url)
    pg = {
        k: v[next(iter(range(len(v))))]
        for k, v in parse_qs(parsed.query).items()
        if k not in ("pgbouncer", "connection_limit")
    }
    if "sslmode" not in pg:
        pg["sslmode"] = "require"
    db_url = urlunparse(parsed._replace(query=urlencode(pg)))
    conn = psycopg2.connect(db_url)
    cur = conn.cursor()
    cur.execute("SET statement_" + "timeout = %s", (PG_STATEMENT_NO_LIMIT,))
    return conn


def _download_hf_index() -> dict:
    """Download and parse index.json from HF Hub."""
    from huggingface_hub import hf_hub_download

    token = os.environ.get("HF_TOKEN") or os.environ.get(
        "HUGGING_FACE_HUB_TOKEN"
    )
    local_path = hf_hub_download(
        repo_id=HF_REPO_ID, filename="index.json",
        repo_type=HF_REPO_TYPE, token=token,
    )
    with open(local_path, "r") as f:
        return json.load(f)


def _parse_index_combos(index: dict) -> list:
    """Parse index.json into (model, task, strategy, layers) tuples."""
    combos = []
    for key, layers in index.items():
        parts = key.split("/")
        if len(parts) < MIN_COMBO_PATH_PARTS:
            continue
        safe_model = parts[INDEX_FIRST]
        strategy = parts[INDEX_LAST]
        task = "/".join(parts[COMBO_OFFSET:INDEX_LAST])
        model = safe_name_to_model(safe_model)
        combos.append((model, task, strategy, sorted(layers)))
    return combos


def _delete_layer(database_url, model_id, set_id, label, layer):
    """Delete one layer in its own connection. Returns rows deleted."""
    conn = _get_db_connection(database_url)
    try:
        cur = conn.cursor()
        cur.execute(
            """SELECT COUNT(*) FROM "Activation"
               WHERE "modelId" = %s AND "contrastivePairSetId" = %s
                 AND "extractionStrategy" = %s AND layer = %s""",
            (model_id, set_id, label, layer),
        )
        count = cur.fetchone()[INDEX_FIRST]
        if count == RECURSION_INITIAL_DEPTH:
            return RECURSION_INITIAL_DEPTH
        cur.execute(
            """DELETE FROM "Activation"
               WHERE "modelId" = %s AND "contrastivePairSetId" = %s
                 AND "extractionStrategy" = %s AND layer = %s""",
            (model_id, set_id, label, layer),
        )
        deleted = cur.rowcount
        conn.commit()
        return deleted
    except Exception as exc:
        try:
            conn.rollback()
        except Exception:
            pass
        conn = _get_db_connection(database_url)
        cur = conn.cursor()
        cur.execute(
            """DELETE FROM "Activation"
               WHERE "modelId" = %s AND "contrastivePairSetId" = %s
                 AND "extractionStrategy" = %s AND layer = %s""",
            (model_id, set_id, label, layer),
        )
        deleted = cur.rowcount
        conn.commit()
        with _print_lock:
            print(f"  RETRY OK layer={layer}: {exc}")
        return deleted
    finally:
        try:
            conn.close()
        except Exception:
            pass


def cleanup_migrated_activations(
    database_url: Optional[str] = None,
    dry_run: bool = True,
) -> dict:
    """Delete Supabase Activation rows that exist in HF Hub.

    Uses PARALLEL_CLEANUP_WORKERS threads for throughput.
    """
    print("Downloading HF index.json ...")
    index = _download_hf_index()
    combos = _parse_index_combos(index)
    print(f"Found {len(combos)} combos in HF index")

    conn = _get_db_connection(database_url)
    cur = conn.cursor()
    model_cache, set_cache = {}, {}
    total_deleted = RECURSION_INITIAL_DEPTH
    total_skipped = RECURSION_INITIAL_DEPTH
    results = []

    work_items = []
    for model, task, strategy, layers in combos:
        if model not in model_cache:
            cur.execute(
                'SELECT id FROM "Model" WHERE "huggingFaceId" = %s',
                (model,),
            )
            row = cur.fetchone()
            model_cache[model] = row[INDEX_FIRST] if row else None
        model_id = model_cache[model]
        if model_id is None:
            total_skipped += len(layers)
            continue
        if task not in set_cache:
            cur.execute(
                'SELECT id FROM "ContrastivePairSet" WHERE name = %s',
                (task,),
            )
            row = cur.fetchone()
            set_cache[task] = row[INDEX_FIRST] if row else None
        set_id = set_cache[task]
        if set_id is None:
            total_skipped += len(layers)
            continue
        for layer in layers:
            work_items.append(
                (model, task, strategy, layer, model_id, set_id)
            )
    conn.close()
    print(f"Work items: {len(work_items)} layers to check")
    print(f"Skipped (not in DB): {total_skipped} layer-slots")

    if dry_run:
        print("[DRY-RUN] parallel delete not supported, exiting")
        return {"dry_run": True, "combos_affected": RECURSION_INITIAL_DEPTH,
                "total_rows_deleted": RECURSION_INITIAL_DEPTH,
                "total_skipped": total_skipped, "details": []}

    workers = PARALLEL_CLEANUP_WORKERS
    print(f"Starting {workers} parallel workers ...")

    def _do_work(item):
        model, task, strategy, layer, m_id, s_id = item
        return (model, task, strategy, layer,
                _delete_layer(database_url, m_id, s_id, strategy, layer))

    done_count = RECURSION_INITIAL_DEPTH
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(_do_work, w): w for w in work_items}
        for future in as_completed(futures):
            model, task, strategy, layer, deleted = future.result()
            done_count += deleted
            if deleted > RECURSION_INITIAL_DEPTH:
                with _print_lock:
                    print(
                        f"  Deleted {deleted} rows: "
                        f"{model}/{task}/{strategy}/layer={layer}"
                        f"  [total={done_count}]"
                    )

    print(f"\nDone. Total deleted this run: {done_count}")
    return {
        "dry_run": False, "combos_affected": len(work_items),
        "total_rows_deleted": done_count,
        "total_skipped": total_skipped, "details": [],
    }
