#!/usr/bin/env python3
"""
G-pre1 INCREMENTAL: Append new objects to object_groups via embedding similarity.

Non-destructive alternative to gpre1_object_clustering.py (which DELETEs and
rebuilds all groups via K-Means). This script:
  - Finds objects referenced in atomic controls that are NOT yet in
    object_groups.members
  - Embeds each unmatched object via bge-m3 (local embedding-service)
  - Nearest-neighbor search against existing object_groups.canonical_name
  - Cosine >= --threshold (default 0.85) → APPEND to existing group's members
  - Cosine <  --threshold → CREATE new object_group with next free group_id

Existing groups stay; only members get appended and new groups get added.

Usage (inside control-pipeline container):
  python3 /app/scripts/gpre1_object_groups_incremental.py --since 2026-05-18T02:53:00+00:00 --dry-run
  python3 /app/scripts/gpre1_object_groups_incremental.py --since 2026-05-18T02:53:00+00:00
  python3 /app/scripts/gpre1_object_groups_incremental.py --since 2026-05-18T02:53:00+00:00 --threshold 0.82
"""

import argparse
import json
import logging
import os
from datetime import datetime

import httpx
import numpy as np
from sqlalchemy import create_engine, text

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("gpre1_inc")

DB_URL = os.getenv("DATABASE_URL", "postgresql://breakpilot:breakpilot123@postgres:5432/breakpilot_db")
EMBEDDING_URL = os.getenv("EMBEDDING_URL", "http://embedding-service:8087")
BATCH_SIZE = 64


def embed_batch(texts: list[str]) -> np.ndarray:
    """Embed a list of strings via bge-m3 embedding-service."""
    with httpx.Client(timeout=120.0) as c:
        resp = c.post(f"{EMBEDDING_URL}/embed", json={"texts": texts, "normalize": True})
        resp.raise_for_status()
        return np.array(resp.json()["embeddings"], dtype=np.float32)


def embed_many(texts: list[str], label: str = "") -> np.ndarray:
    """Embed many strings in batches."""
    n = len(texts)
    out = np.zeros((n, 1024), dtype=np.float32)
    for i in range(0, n, BATCH_SIZE):
        batch = texts[i:i + BATCH_SIZE]
        out[i:i + len(batch)] = embed_batch(batch)
        if (i // BATCH_SIZE) % 20 == 0:
            logger.info("  %s: %d/%d embedded", label, i + len(batch), n)
    return out


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--since", required=True, help="ISO datetime — consider atomics from this date onwards")
    parser.add_argument("--threshold", type=float, default=0.85,
                        help="Cosine threshold for appending to existing group (default 0.85)")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    since_dt = datetime.fromisoformat(args.since.replace("Z", "+00:00"))
    logger.info("Incremental object_groups update since %s, threshold=%.2f, dry_run=%s",
                since_dt.isoformat(), args.threshold, args.dry_run)

    engine = create_engine(DB_URL, connect_args={"options": "-c search_path=compliance,public"})

    # 1. Load existing object_groups (id, canonical_name, members)
    with engine.connect() as c:
        rows = c.execute(text("""
            SELECT group_id, canonical_name, members FROM object_groups
        """)).fetchall()
    existing_groups = [(r[0], r[1], json.loads(r[2]) if isinstance(r[2], str) else r[2]) for r in rows]
    logger.info("Loaded %d existing object_groups", len(existing_groups))

    existing_members: set[str] = set()
    for _, _, members in existing_groups:
        for m in members:
            existing_members.add(m)
    logger.info("Existing union of members: %d distinct strings", len(existing_members))

    # 2. Find unmatched objects from atomics since `since`
    from services.control_dedup import normalize_object
    with engine.connect() as c:
        rows = c.execute(text("""
            SELECT DISTINCT split_part(generation_metadata->>'merge_group_hint', ':', 2) AS obj
            FROM canonical_controls
            WHERE decomposition_method = 'pass0b'
              AND created_at >= :since
              AND generation_metadata->>'merge_group_hint' IS NOT NULL
              AND generation_metadata->>'merge_group_hint' != ''
              AND release_state NOT IN ('deprecated', 'rejected', 'duplicate')
        """), {"since": since_dt}).fetchall()
    new_objects_raw = [r[0] for r in rows if r[0]]
    logger.info("Distinct objects in new atomics: %d", len(new_objects_raw))

    # Normalize each + dedupe; track originals → normalized
    normed_to_originals: dict[str, set[str]] = {}
    for obj in new_objects_raw:
        normed = normalize_object(obj)
        if not normed:
            continue
        if normed in existing_members or obj in existing_members:
            continue  # already in some group
        normed_to_originals.setdefault(normed, set()).update([normed, obj])

    unmatched_normed = list(normed_to_originals.keys())
    logger.info("Unmatched normalized objects: %d", len(unmatched_normed))

    if not unmatched_normed:
        logger.info("Nothing to do — all objects already mapped.")
        return

    # 3. Embed existing canonical_names + unmatched objects
    logger.info("Embedding %d existing canonical_names...", len(existing_groups))
    existing_emb = embed_many([g[1] for g in existing_groups], label="existing")
    logger.info("Embedding %d unmatched objects...", len(unmatched_normed))
    unmatched_emb = embed_many(unmatched_normed, label="unmatched")

    # 4. Nearest-neighbor: for each unmatched, find best existing match
    # cosine = dot product (both already L2-normalized)
    logger.info("Computing nearest-neighbor matches...")
    sims = unmatched_emb @ existing_emb.T  # (N_unmatched, N_existing)
    best_idx = sims.argmax(axis=1)
    best_score = sims.max(axis=1)

    appends: dict[int, list[str]] = {}  # group_id → list of new members
    new_groups: list[tuple[str, list[str]]] = []  # (canonical_name, members)

    for i, normed in enumerate(unmatched_normed):
        originals = sorted(normed_to_originals[normed])
        if best_score[i] >= args.threshold:
            gid = existing_groups[int(best_idx[i])][0]
            appends.setdefault(gid, []).extend(originals)
        else:
            # Create a new group with this object as canonical
            new_groups.append((normed, originals))

    # Stats
    distinct_groups_to_extend = len(appends)
    total_appends = sum(len(v) for v in appends.values())
    logger.info("Plan: extend %d existing groups (+%d members), create %d new groups",
                distinct_groups_to_extend, total_appends, len(new_groups))

    if args.dry_run:
        logger.info("DRY RUN — no writes")
        # Sample
        if appends:
            sample = list(appends.items())[:5]
            for gid, members in sample:
                gname = next((g[1] for g in existing_groups if g[0] == gid), "?")
                logger.info("  Extend group_id=%d (%s) with: %s", gid, gname, members[:3])
        if new_groups:
            for name, members in new_groups[:5]:
                logger.info("  NEW group: %s — members=%s", name, members[:3])
        return

    # 5. Write — pure INSERT/UPDATE
    with engine.begin() as c:
        c.execute(text("SET search_path TO compliance, public"))

        # UPDATE existing groups (append to members JSONB)
        for gid, new_members in appends.items():
            c.execute(text("""
                UPDATE object_groups
                SET members = (
                    SELECT jsonb_agg(DISTINCT m)
                    FROM jsonb_array_elements_text(members || CAST(:new_members AS jsonb)) AS x(m)
                ),
                member_count = (
                    SELECT count(DISTINCT m)
                    FROM jsonb_array_elements_text(members || CAST(:new_members AS jsonb)) AS x(m)
                )
                WHERE group_id = :gid
            """), {"gid": gid, "new_members": json.dumps(new_members)})

        # INSERT new groups with next free group_id
        next_gid_row = c.execute(text("SELECT COALESCE(MAX(group_id), 0) + 1 FROM object_groups")).fetchone()
        next_gid = next_gid_row[0] if next_gid_row else 1
        for name, members in new_groups:
            c.execute(text("""
                INSERT INTO object_groups (group_id, canonical_name, member_count, members, top_controls_count)
                VALUES (:gid, :name, :count, CAST(:members AS jsonb), 0)
            """), {
                "gid": next_gid,
                "name": name[:200],
                "count": len(members),
                "members": json.dumps(members),
            })
            next_gid += 1

    logger.info("DONE — extended %d existing groups (+%d members), created %d new groups",
                distinct_groups_to_extend, total_appends, len(new_groups))


if __name__ == "__main__":
    main()
