#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import ssl
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import pymysql


TERMINAL = ("processed", "skipped", "dead")


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


def iso(dt: datetime) -> str:
    return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def connect() -> pymysql.connections.Connection:
    return pymysql.connect(
        host=os.environ["TIDB_HOST"],
        port=int(os.environ.get("TIDB_PORT", "4000")),
        user=os.environ["TIDB_USER"],
        password=os.environ["TIDB_PASSWORD"],
        database=os.environ.get("TIDB_DATABASE", "nexuscrm"),
        autocommit=True,
        connect_timeout=10,
        read_timeout=120,
        write_timeout=120,
        ssl=ssl.create_default_context(),
        cursorclass=pymysql.cursors.DictCursor,
    )


def fetch_all(conn: pymysql.connections.Connection, sql: str, args: tuple[Any, ...] = ()) -> list[dict[str, Any]]:
    with conn.cursor() as cur:
        cur.execute(sql, args)
        return list(cur.fetchall())


def execute(conn: pymysql.connections.Connection, sql: str, args: tuple[Any, ...] = ()) -> int:
    with conn.cursor() as cur:
        cur.execute(sql, args)
        return int(cur.rowcount)


def scalar(conn: pymysql.connections.Connection, sql: str, args: tuple[Any, ...] = ()) -> int:
    rows = fetch_all(conn, sql, args)
    if not rows:
        return 0
    return int(next(iter(rows[0].values())) or 0)


def count_by(conn: pymysql.connections.Connection, table: str, field: str, where: str, args: tuple[Any, ...]) -> dict[str, int]:
    rows = fetch_all(conn, f"SELECT {field} AS k, COUNT(*) AS c FROM {table} WHERE {where} GROUP BY {field}", args)
    return {str(row["k"]): int(row["c"]) for row in rows}


def git_head(repo: str) -> str:
    try:
        return subprocess.check_output(["git", "-C", repo, "rev-parse", "HEAD"], text=True).strip()
    except Exception:
        return ""


def real_pack(conn: pymysql.connections.Connection) -> dict[str, Any]:
    rows = fetch_all(
        conn,
        """
        SELECT Id, PackInstallId, Name, Category, Status, SourceRevision, LastModifiedDate
          FROM _System_InstalledPack
         WHERE Name='hubspot-crm-connector'
           AND COALESCE(IsDeleted, 0)=0
         ORDER BY LastModifiedDate DESC
         LIMIT 1
        """,
    )
    if not rows:
        raise RuntimeError("hubspot-crm-connector installed pack not found")
    return rows[0]


def cleanup(conn: pymysql.connections.Connection, run_id: str, pack_id: str, generation: int) -> dict[str, int]:
    event_prefix = f"cw-matrix-{run_id}-event-%"
    work_prefix = f"cw-matrix-{run_id}-work-%"
    route_id = f"route-cw-matrix-{run_id}"
    out: dict[str, int] = {}
    out["workitems"] = execute(
        conn,
        "DELETE FROM _System_ConnectorWorkItem WHERE InstalledPackId=%s AND Generation=%s AND WorkId LIKE %s",
        (pack_id, generation, work_prefix),
    )
    out["events"] = execute(
        conn,
        "DELETE FROM _System_ConnectorEvent WHERE InstalledPackId=%s AND Id LIKE %s",
        (pack_id, event_prefix),
    )
    out["route"] = execute(
        conn,
        "DELETE FROM _System_ConnectorInboundRoute WHERE InstalledPackId=%s AND Id=%s",
        (pack_id, route_id),
    )
    # SchedulerLease is a pack/generation-wide runtime coordination table, not
    # scoped to this synthetic run. Leave it intact so the gate cannot disturb
    # unrelated staging workers.
    out["scheduler_lease"] = 0
    return out


def insert_route(conn: pymysql.connections.Connection, run_id: str, pack_id: str, queue_class: str) -> int:
    now = utc_now().strftime("%Y-%m-%d %H:%M:%S")
    return execute(
        conn,
        """
        INSERT INTO _System_ConnectorInboundRoute
          (Id, Namespace, InstalledPackId, EventKind, Version, FlowDeveloperName, IsCurrent, IsDeleted,
           CreatedDate, LastModifiedDate, InvocationMode, BatchInputName, BatchResultOutputName, MaxBatchSize,
           QueueClass, FlowExecutionFilter)
        VALUES
          (%s, 'hubspot-crm-connector', %s, 'noop.skip', 1, 'NoopFlow__c', 1, 0,
           %s, %s, 'per_event', 'events', 'eventResults', 100, %s,
           JSON_OBJECT('path', '.execute', 'equalsAny', JSON_ARRAY('yes')))
        ON DUPLICATE KEY UPDATE
          IsCurrent=1, IsDeleted=0, LastModifiedDate=VALUES(LastModifiedDate),
          QueueClass=VALUES(QueueClass), FlowExecutionFilter=VALUES(FlowExecutionFilter)
        """,
        (f"route-cw-matrix-{run_id}", pack_id, now, now, queue_class),
    )


def seed(conn: pymysql.connections.Connection, run_id: str, pack_id: str, events: int, lanes: int, queue_class: str) -> dict[str, int]:
    now = utc_now().strftime("%Y-%m-%d %H:%M:%S")
    event_rows = []
    work_rows = []
    lane_seq: dict[str, int] = {}
    for i in range(events):
        lane = f"cw-matrix-{run_id}-lane-{i % lanes:04d}"
        lane_seq[lane] = lane_seq.get(lane, 0) + 1
        seq = lane_seq[lane]
        event_id = f"cw-matrix-{run_id}-event-{i:06d}"
        work_id = f"cw-matrix-{run_id}-work-{i:06d}"
        payload = json.dumps(
            {"objectKind": "contact", "sourceID": lane, "payload": {"id": lane, "runId": run_id, "execute": "no"}},
            separators=(",", ":"),
            sort_keys=True,
        )
        state = "pending" if seq == 1 else "blocked"
        event_rows.append(
            (event_id, "hubspot-crm-connector", pack_id, "noop.skip", f"{run_id}-{i:06d}", "pending", queue_class, payload, seq, now, i, 1, "contact", lane, lane)
        )
        work_rows.append((work_id, event_id, pack_id, "connector_event", queue_class, state, now, 0, seq, lane, 0, now, now))
    with conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO _System_ConnectorEvent
              (Id, Namespace, InstalledPackId, EventKind, IdempotencyKey, Status, QueueClass, Payload,
               EnqueueOrdinal, CreatedDate, DeliverySequence, RouteVersion, SourceObjectKind, SourceObjectID, SourceLaneKey)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            event_rows,
        )
        inserted_events = int(cur.rowcount)
        cur.executemany(
            """
            INSERT INTO _System_ConnectorWorkItem
              (WorkId, EventId, InstalledPackId, WorkType, QueueClass, State, RunAfter, Attempt, LaneOrdinal,
               SourceLaneKey, Generation, CreatedDate, LastModifiedDate)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            work_rows,
        )
        inserted_work = int(cur.rowcount)
    return {"events": inserted_events, "workitems": inserted_work}


def snapshot(conn: pymysql.connections.Connection, run_id: str, pack_id: str, generation: int) -> dict[str, Any]:
    event_like = f"cw-matrix-{run_id}-event-%"
    work_like = f"cw-matrix-{run_id}-work-%"
    event_where = "InstalledPackId=%s AND Id LIKE %s"
    event_args = (pack_id, event_like)
    work_where = "InstalledPackId=%s AND Generation=%s AND WorkId LIKE %s"
    work_args = (pack_id, generation, work_like)
    return {
        "event_status": count_by(conn, "_System_ConnectorEvent", "Status", event_where, event_args),
        "workitem_state": count_by(conn, "_System_ConnectorWorkItem", "State", work_where, work_args),
        "event_terminal": scalar(conn, "SELECT COUNT(*) FROM _System_ConnectorEvent WHERE " + event_where + " AND Status IN ('processed','skipped','dead')", event_args),
        "work_terminal": scalar(conn, "SELECT COUNT(*) FROM _System_ConnectorWorkItem WHERE " + work_where + " AND State IN ('processed','skipped','dead')", work_args),
        "eligible_heads": scalar(
            conn,
            """
            SELECT COUNT(*) FROM _System_ConnectorWorkItem w
             WHERE w.InstalledPackId=%s AND w.Generation=%s AND w.WorkId LIKE %s
               AND w.State IN ('pending','retry') AND w.RunAfter <= UTC_TIMESTAMP()
               AND NOT EXISTS (
                 SELECT 1 FROM _System_ConnectorWorkItem older
                  WHERE older.InstalledPackId=w.InstalledPackId AND older.Generation=w.Generation
                    AND older.WorkId LIKE %s
                    AND COALESCE(NULLIF(older.TargetMutationLaneKey,''), older.SourceLaneKey, '') = COALESCE(NULLIF(w.TargetMutationLaneKey,''), w.SourceLaneKey, '')
                    AND older.State IN ('pending','retry','processing')
                    AND (older.LaneOrdinal < w.LaneOrdinal OR (older.LaneOrdinal=w.LaneOrdinal AND older.WorkId < w.WorkId))
               )
            """,
            (pack_id, generation, work_like, work_like),
        ),
        "claimed_or_terminal": scalar(conn, "SELECT COUNT(*) FROM _System_ConnectorWorkItem WHERE " + work_where + " AND State IN ('processing','processed','skipped','dead')", work_args),
        "lease_width": fetch_all(
            conn,
            """
            SELECT COUNT(DISTINCT LeaseOwner) AS owners, COUNT(DISTINCT LeaseToken) AS tokens,
                   COUNT(*) AS leased_work, COUNT(DISTINCT SourceLaneKey) AS leased_lanes
              FROM _System_ConnectorWorkItem
             WHERE InstalledPackId=%s AND Generation=%s AND WorkId LIKE %s AND State='processing'
            """,
            work_args,
        )[0],
        "scheduler_lease": fetch_all(
            conn,
            "SELECT InstalledPackId, Generation, OwnerId, FencingToken, LeaseUntil, RenewedAt FROM _System_ConnectorSchedulerLease WHERE InstalledPackId=%s AND Generation=%s",
            (pack_id, generation),
        ),
        "split_brain": scalar(
            conn,
            """
            SELECT COUNT(*)
              FROM _System_ConnectorEvent e
              JOIN _System_ConnectorWorkItem w ON w.EventId=e.Id
             WHERE e.InstalledPackId=%s AND e.Id LIKE %s AND w.Generation=%s
               AND e.Status IN ('processed','skipped','dead')
               AND w.State NOT IN ('processed','skipped','dead')
            """,
            (pack_id, event_like, generation),
        ),
    }


def explain_broad_scan(conn: pymysql.connections.Connection, run_id: str, pack_id: str, generation: int) -> tuple[list[str], bool]:
    rows = fetch_all(
        conn,
        """
        EXPLAIN SELECT w.WorkId
          FROM _System_ConnectorWorkItem w
          JOIN _System_ConnectorSchedulerLease l
            ON l.InstalledPackId = w.InstalledPackId
           AND l.Generation = w.Generation
           AND l.LeaseUntil > UTC_TIMESTAMP()
         WHERE w.InstalledPackId = %s
           AND w.Generation = %s
           AND w.WorkId LIKE %s
           AND w.State IN ('pending','retry')
           AND w.RunAfter <= UTC_TIMESTAMP()
         ORDER BY w.RunAfter ASC, w.WorkId ASC
         LIMIT 16
        """,
        (pack_id, generation, f"cw-matrix-{run_id}-work-%"),
    )
    text = [json.dumps(row, default=str, sort_keys=True) for row in rows]
    broad = any("TableFullScan" in line or ("Selection" in line and "_System_ConnectorEvent" in line) for line in text)
    return text, broad


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", required=True)
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--repo", default="")
    parser.add_argument("--events", type=int, default=800)
    parser.add_argument("--lanes", type=int, default=120)
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--queue-class", default="bulk")
    parser.add_argument("--timeout", type=int, default=720)
    parser.add_argument("--poll", type=float, default=2.0)
    parser.add_argument("--cleanup", action="store_true")
    args = parser.parse_args()

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    generation = 0
    conn = connect()
    report: dict[str, Any] = {
        "run_id": args.run_id,
        "head": git_head(args.repo) if args.repo else "",
        "events": args.events,
        "lanes": args.lanes,
        "workers_expected": args.workers,
        "queue_class": args.queue_class,
        "started_at": iso(utc_now()),
        "validation_mode": "deployed backend real installed hubspot pack with scoped materialized rows",
    }
    try:
        pack = real_pack(conn)
        pack_id = str(pack["PackInstallId"])
        report["pack_readback"] = {k: str(v) for k, v in pack.items()}
        report["precleanup"] = cleanup(conn, args.run_id, pack_id, generation)
        report["route_rows"] = insert_route(conn, args.run_id, pack_id, args.queue_class)
        report["seed"] = seed(conn, args.run_id, pack_id, args.events, args.lanes, args.queue_class)
        start = time.time()
        first_claim_at: float | None = None
        max_processing = max_owners = max_tokens = max_work = max_lanes = 0
        buckets: list[dict[str, Any]] = []
        while True:
            snap = snapshot(conn, args.run_id, pack_id, generation)
            elapsed = time.time() - start
            processing = int(snap["workitem_state"].get("processing", 0))
            terminal_work = int(snap["work_terminal"])
            terminal_event = int(snap["event_terminal"])
            lease_width = snap["lease_width"]
            max_processing = max(max_processing, processing)
            max_owners = max(max_owners, int(lease_width.get("owners") or 0))
            max_tokens = max(max_tokens, int(lease_width.get("tokens") or 0))
            max_work = max(max_work, int(lease_width.get("leased_work") or 0))
            max_lanes = max(max_lanes, int(lease_width.get("leased_lanes") or 0))
            if first_claim_at is None and (processing > 0 or terminal_work > 0):
                first_claim_at = elapsed
            bucket = {"elapsed_seconds": round(elapsed, 3), **snap}
            buckets.append(bucket)
            print(
                "sample",
                round(elapsed, 1),
                "event", snap["event_status"],
                "work", snap["workitem_state"],
                "eligible", snap["eligible_heads"],
                "claimed", snap["claimed_or_terminal"],
                "lease", snap["lease_width"],
                "split", snap["split_brain"],
                flush=True,
            )
            if terminal_work >= args.events and terminal_event >= args.events:
                break
            if elapsed >= args.timeout:
                report["error"] = "timeout"
                break
            time.sleep(args.poll)
        final = snapshot(conn, args.run_id, pack_id, generation)
        explain, broad = explain_broad_scan(conn, args.run_id, pack_id, generation)
        report.update(
            {
                "finished_at": iso(utc_now()),
                "duration_seconds": round(time.time() - start, 3),
                "first_claim_latency_ms": None if first_claim_at is None else round(first_claim_at * 1000, 3),
                "max_processing_observed": max_processing,
                "max_lease_width": {"owners": max_owners, "tokens": max_tokens, "work_items": max_work, "lanes": max_lanes},
                "final": final,
                "buckets": buckets[:: max(1, len(buckets) // 60)] if len(buckets) > 60 else buckets,
                "explain_access_summary": explain,
                "explain_has_broad_scan": broad,
                "verdict": "PASS" if "error" not in report and final["work_terminal"] >= args.events and final["event_terminal"] >= args.events else "BLOCK",
            }
        )
        if args.cleanup:
            report["cleanup"] = cleanup(conn, args.run_id, pack_id, generation)
    finally:
        conn.close()
    out_path.write_text(json.dumps(report, indent=2, sort_keys=True, default=str) + "\n")
    print(out_path, flush=True)
    print(report.get("verdict"), flush=True)
    return 0 if report.get("verdict") == "PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())