#!/usr/bin/env python3
"""
tests/ts_import_sync_vs_async.py — sync vs async TML import, apples-to-apples.

Compares the app's current SYNCHRONOUS table-import path against the
ThoughtSpot ASYNC import API, on the real two-phase flow the app uses:

    Phase 1  create all tables (no joins)          create_new=True
    Phase 2  re-import the tables WITH joins_with   create_new=False (update)

Both modes build a small star schema (SALES fact + PRODUCT/STORE/CUSTOMER
dims, 3 FK joins) via the app's real ThoughtSpotDeployer.create_table_tml(),
so the join wiring is identical to production. Each mode is timed as
WALL-CLOCK-TO-VERIFIED:
    - phase 1 done  = all 4 tables visible under the connection
    - phase 2 done  = SALES export contains its 3 joins_with

SYNC   : POST /metadata/tml/import  (timeout 360; may 504 at the ~300s gateway)
         then poll the connection until verified — mirrors deploy_all.
ASYNC  : POST /metadata/tml/async/import -> task_id (returns instantly),
         then poll /metadata/tml/async/status until completed_at>0.

Everything created is TAGGED 'PERF_TEST_DELETE' and CLEANED UP by default
(all logical tables on the connection are deleted, the Snowflake schema is
dropped, and the connection is deleted). Pass --keep to leave objects.

SAFETY: refuses prod unless --allow-prod. Prints a plan; executes with --yes.

USAGE:
    source ./demoprep/bin/activate
    python tests/ts_import_sync_vs_async.py --env "sebe - se"          # preview
    python tests/ts_import_sync_vs_async.py --env "sebe - se" --yes    # run+clean
    python tests/ts_import_sync_vs_async.py --env "sebe - se" --yes --keep
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path

from dotenv import load_dotenv

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT))
load_dotenv(REPO_ROOT / ".env")

from thoughtspot_deployer import ThoughtSpotDeployer  # noqa: E402
from snowflake_auth import get_snowflake_connection  # noqa: E402

RESULTS_DIR = REPO_ROOT / "tests" / "perf_results"
TAG = "PERF_TEST_DELETE"
DB, SCHEMA = "DEMOBUILD", "ZPERFTEST"
CONNECTION = "ZPERFTEST_DONOTDELETE"          # reused if it already exists
POLL_DEADLINE_S = 600
POLL_INTERVAL_S = 5


# ---- env resolution (inlined; see chat_interface.get_ts_env_*) --------------

def resolve_env(label: str) -> tuple[str, str]:
    i = 1
    while True:
        env_label = os.getenv(f"TS_ENV_{i}_LABEL", "").strip()
        if not env_label:
            break
        if env_label == label:
            return (os.getenv(f"TS_ENV_{i}_URL", "").strip().rstrip("/"),
                    os.getenv(f"TS_ENV_{i}_KEY_VAR", "").strip())
        i += 1
    raise SystemExit(f"TS environment {label!r} not configured in .env")


def is_prod(base_url: str, label: str) -> bool:
    return "se-thoughtspot-cloud" in base_url.lower() or label.lower().startswith("secloud")


# ---- star schema builder ----------------------------------------------------

def build_star(prefix: str) -> tuple[dict, list]:
    """Return (tables{name: columns}, foreign_keys[]) for a prefixed star schema."""
    P = prefix
    tables = {
        f"{P}SALES": [
            {"name": "SALE_ID", "type": "INT"},
            {"name": "PRODUCT_ID", "type": "INT"},
            {"name": "STORE_ID", "type": "INT"},
            {"name": "CUSTOMER_ID", "type": "INT"},
            {"name": "AMOUNT", "type": "NUMBER(18,2)"},
            {"name": "SALE_DATE", "type": "DATE"},
        ],
        f"{P}PRODUCT": [
            {"name": "PRODUCT_ID", "type": "INT"},
            {"name": "PRODUCT_NAME", "type": "VARCHAR"},
            {"name": "CATEGORY", "type": "VARCHAR"},
        ],
        f"{P}STORE": [
            {"name": "STORE_ID", "type": "INT"},
            {"name": "STORE_NAME", "type": "VARCHAR"},
            {"name": "REGION", "type": "VARCHAR"},
        ],
        f"{P}CUSTOMER": [
            {"name": "CUSTOMER_ID", "type": "INT"},
            {"name": "CUSTOMER_NAME", "type": "VARCHAR"},
            {"name": "SEGMENT", "type": "VARCHAR"},
        ],
    }
    fks = [
        {"from_table": f"{P}SALES", "to_table": f"{P}PRODUCT",
         "from_column": "PRODUCT_ID", "to_column": "PRODUCT_ID"},
        {"from_table": f"{P}SALES", "to_table": f"{P}STORE",
         "from_column": "STORE_ID", "to_column": "STORE_ID"},
        {"from_table": f"{P}SALES", "to_table": f"{P}CUSTOMER",
         "from_column": "CUSTOMER_ID", "to_column": "CUSTOMER_ID"},
    ]
    return tables, fks


def ensure_physical(tables: dict) -> None:
    conn = get_snowflake_connection()
    cur = conn.cursor()
    try:
        cur.execute(f'CREATE SCHEMA IF NOT EXISTS "{DB}"."{SCHEMA}"')
        for name, cols in tables.items():
            ddl_cols = ", ".join(f'"{c["name"].upper()}" {c["type"]}' for c in cols)
            cur.execute(f'CREATE TABLE IF NOT EXISTS "{DB}"."{SCHEMA}"."{name.upper()}" ({ddl_cols})')
    finally:
        cur.close()
        conn.close()


# ---- import helpers ---------------------------------------------------------

def _normalize(result):
    def _one(o):
        if not isinstance(o, dict):
            return o
        return o if "response" in o else ({"response": o} if ("status" in o or "header" in o) else o)
    if isinstance(result, list):
        return [_one(o) for o in result]
    if isinstance(result, dict) and "object" in result:
        return [_one(o) for o in result["object"]]
    return None


def sync_import_call(dep, tmls, create_new) -> dict:
    """One synchronous batch import — mirrors deploy_all._import_tml_chunk."""
    t0 = time.time()
    try:
        r = dep.session.post(f"{dep.base_url}/api/rest/2.0/metadata/tml/import",
                             json={"metadata_tmls": tmls, "import_policy": "PARTIAL",
                                   "create_new": create_new}, timeout=360)
    except Exception as e:
        return {"call_s": round(time.time() - t0, 2), "http": None, "gw": False, "err": str(e)[:300]}
    return {"call_s": round(time.time() - t0, 2), "http": r.status_code,
            "gw": r.status_code in (502, 503, 504),
            "err": None if r.status_code == 200 else r.text[:200]}


def async_import_call(dep, tmls, create_new) -> dict:
    """Submit async import and poll status until completed_at>0 or deadline."""
    t0 = time.time()
    r = dep.session.post(f"{dep.base_url}/api/rest/2.0/metadata/tml/async/import",
                         json={"metadata_tmls": tmls, "import_policy": "PARTIAL",
                               "create_new": create_new}, timeout=60)
    submit_s = round(time.time() - t0, 2)
    if r.status_code not in (200, 202):
        return {"submit_s": submit_s, "complete_s": None, "http": r.status_code,
                "task_status": "SUBMIT_FAILED", "err": r.text[:200]}
    task_id = (r.json() or {}).get("task_id")
    final = None
    deadline = time.time() + POLL_DEADLINE_S
    while time.time() < deadline:
        s = dep.session.post(f"{dep.base_url}/api/rest/2.0/metadata/tml/async/status",
                             json={"task_ids": [task_id], "include_import_response": True}, timeout=60)
        if s.status_code == 200:
            sl = ((s.json() or {}).get("status_list") or [{}])[0]
            final = sl
            if sl.get("completed_at") or sl.get("task_status") in (
                    "COMPLETED", "SUCCESS", "FAILED", "ERROR", "PARTIAL_SUCCESS"):
                break
        time.sleep(POLL_INTERVAL_S)
    return {"submit_s": submit_s, "complete_s": round(time.time() - t0, 2),
            "http": r.status_code, "task_status": (final or {}).get("task_status"),
            "processed": (final or {}).get("object_processed_count"),
            "err": None}


# ---- connection scan / verification -----------------------------------------

def scan(dep, conn_guid) -> dict:
    """name -> guid for logical tables under the connection."""
    resolved = dep.search_logical_tables_for_connection(conn_guid, CONNECTION)
    return {n: v["response"]["header"]["id_guid"] for n, v in resolved.items()}


def wait_visible(dep, conn_guid, names) -> tuple[float, dict]:
    want = {n.upper() for n in names}
    t0 = time.time()
    deadline = t0 + POLL_DEADLINE_S
    while True:
        have = scan(dep, conn_guid)
        if want <= set(have):
            return round(time.time() - t0, 2), have
        if time.time() >= deadline:
            return round(time.time() - t0, 2), have
        time.sleep(POLL_INTERVAL_S)


def joins_on(dep, guid) -> int:
    r = dep.session.post(f"{dep.base_url}/api/rest/2.0/metadata/tml/export",
                         json={"metadata": [{"identifier": guid, "type": "LOGICAL_TABLE"}],
                               "export_associated": False, "format_type": "YAML"}, timeout=60)
    if r.status_code != 200:
        return -1
    data = r.json() or []
    if not data or "edoc" not in data[0]:
        return -1
    tml = dep._parse_tml_edoc(data[0]["edoc"])
    return len((tml.get("table", {}) or {}).get("joins_with") or [])


def wait_joins(dep, sales_guid, expected) -> tuple[float, int]:
    t0 = time.time()
    deadline = t0 + POLL_DEADLINE_S
    while True:
        n = joins_on(dep, sales_guid)
        if n >= expected:
            return round(time.time() - t0, 2), n
        if time.time() >= deadline:
            return round(time.time() - t0, 2), n
        time.sleep(POLL_INTERVAL_S)


# ---- one full mode run (create + join) --------------------------------------

def run_mode(dep, conn_guid, conn_fqn, mode, prefix) -> dict:
    tables, fks = build_star(prefix)
    names = list(tables.keys())
    sales = f"{prefix}SALES"
    print(f"\n{'='*70}\nMODE={mode.upper()}  prefix={prefix}  tables={names}\n{'='*70}")
    ensure_physical(tables)
    print("  physical tables ready")

    # ---- Phase 1: create (no joins) ----
    p1_tmls = [dep.create_table_tml(n, cols, CONNECTION, DB, SCHEMA,
                                    all_tables=None, foreign_keys=fks, connection_fqn=conn_fqn)
               for n, cols in tables.items()]
    t_phase1 = time.time()
    if mode == "sync":
        call = sync_import_call(dep, p1_tmls, create_new=True)
        print(f"  P1 sync call: {call['call_s']}s http={call['http']} 504={call['gw']}")
    else:
        call = async_import_call(dep, p1_tmls, create_new=True)
        print(f"  P1 async: submit={call['submit_s']}s complete={call['complete_s']}s "
              f"status={call['task_status']}")
    vis_s, have = wait_visible(dep, conn_guid, names)
    p1_total = round(time.time() - t_phase1, 2)
    guids = {n.upper(): have[n.upper()] for n in names if n.upper() in have}
    print(f"  P1 verified {len(guids)}/{len(names)} tables visible; phase1 total={p1_total}s")

    # ---- Phase 2: add joins (update) ----
    p2_tmls, p2_names = [], []
    for n, cols in tables.items():
        g = guids.get(n.upper())
        if not g:
            continue
        p2_tmls.append(dep.create_table_tml(n, cols, CONNECTION, DB, SCHEMA,
                                            all_tables=tables, table_guid=g,
                                            foreign_keys=fks, connection_fqn=conn_fqn))
        p2_names.append(n.upper())
    t_phase2 = time.time()
    if mode == "sync":
        call2 = sync_import_call(dep, p2_tmls, create_new=False)
        print(f"  P2 sync call: {call2['call_s']}s http={call2['http']} 504={call2['gw']}")
    else:
        call2 = async_import_call(dep, p2_tmls, create_new=False)
        print(f"  P2 async: submit={call2['submit_s']}s complete={call2['complete_s']}s "
              f"status={call2['task_status']}")
    join_s, njoins = (0.0, -1)
    if sales.upper() in guids:
        join_s, njoins = wait_joins(dep, guids[sales.upper()], expected=len(fks))
    p2_total = round(time.time() - t_phase2, 2)
    print(f"  P2 verified joins on {sales}: {njoins}/{len(fks)}; phase2 total={p2_total}s")

    # tag everything PERF_TEST_DELETE
    try:
        dep.assign_tags_to_objects(list(guids.values()), "LOGICAL_TABLE", TAG)
        print(f"  tagged {len(guids)} tables '{TAG}'")
    except Exception as e:
        print(f"  tag warning: {e}")

    return {
        "mode": mode, "prefix": prefix,
        "p1_call_s": call.get("call_s"), "p1_submit_s": call.get("submit_s"),
        "p1_http": call.get("http"), "p1_gw": call.get("gw", False),
        "p1_total_s": p1_total, "p1_tables": f"{len(guids)}/{len(names)}",
        "p2_call_s": call2.get("call_s"), "p2_submit_s": call2.get("submit_s"),
        "p2_http": call2.get("http"), "p2_gw": call2.get("gw", False),
        "p2_total_s": p2_total, "joins": f"{njoins}/{len(fks)}",
        "grand_total_s": round(p1_total + p2_total, 2),
        "guids": list(guids.values()),
    }


# ---- cleanup ----------------------------------------------------------------

def cleanup(dep, conn_guid, run_guids, drop_sandbox) -> None:
    """Delete ONLY this run's throwaway logical tables. The sandbox (connection
    + schema + physical tables) is KEPT so we can keep testing — unless
    --drop-sandbox is explicitly given."""
    print(f"\n{'='*70}\nCLEANUP (this run's tables only)\n{'='*70}")
    print(f"  deleting {len(run_guids)} logical tables created by this run…")
    for guid in run_guids:
        r = dep.session.post(f"{dep.base_url}/api/rest/2.0/metadata/delete",
                             json={"metadata": [{"type": "LOGICAL_TABLE", "identifier": guid}]}, timeout=120)
        print(f"    - {guid}: http={r.status_code}")
    if not drop_sandbox:
        print(f"  KEPT sandbox: connection {CONNECTION} + schema {DB}.{SCHEMA} (+ physical tables)")
        return
    # Full teardown (explicit opt-in only).
    try:
        conn = get_snowflake_connection()
        cur = conn.cursor()
        cur.execute(f'DROP SCHEMA IF EXISTS "{DB}"."{SCHEMA}" CASCADE')
        cur.close()
        conn.close()
        print(f"  dropped Snowflake schema {DB}.{SCHEMA}")
    except Exception as e:
        print(f"  snowflake drop warning: {e}")
    # NOTE: metadata/delete rejects type CONNECTION ("does not exist in
    # DeleteMetadatatype enum"). Connections use the dedicated v2 endpoint
    # /connection/delete with a singular string `connection_identifier`.
    r = dep.session.post(f"{dep.base_url}/api/rest/2.0/connection/delete",
                         json={"connection_identifier": conn_guid}, timeout=120)
    print(f"  deleted connection {CONNECTION}: http={r.status_code}")


# ---- reporting --------------------------------------------------------------

def report(stamp, meta, sync_r, async_r) -> Path:
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    md = RESULTS_DIR / f"{stamp}_sync_vs_async.md"

    def _speedup(a, b):
        return f"{a / b:.1f}×" if (a and b) else "n/a"

    lines = [
        f"# Sync vs Async TML import — {stamp}", "",
        f"- Environment: **{meta['env']}** ({meta['base_url']})",
        f"- Connection: {CONNECTION}  •  Snowflake: {DB}.{SCHEMA}",
        f"- Flow: 4-table star schema, phase 1 create + phase 2 joins (3 FKs on SALES)",
        f"- Timing = wall-clock to VERIFIED (tables visible / joins present)", "",
        "## Comparison", "",
        "| metric | SYNC | ASYNC | note |",
        "|---|--:|--:|---|",
        f"| Phase 1 — import call/submit | {sync_r['p1_call_s']}s | {async_r['p1_submit_s']}s "
        f"| async submit returns instantly |",
        f"| Phase 1 — 504? | {'⚠️ yes' if sync_r['p1_gw'] else 'no'} | "
        f"{'⚠️ yes' if async_r['p1_gw'] else 'no'} | sync hits the ~300s gateway wall |",
        f"| Phase 1 — total to verified | {sync_r['p1_total_s']}s | {async_r['p1_total_s']}s | |",
        f"| Phase 2 — import call/submit | {sync_r['p2_call_s']}s | {async_r['p2_submit_s']}s | |",
        f"| Phase 2 — total to verified | {sync_r['p2_total_s']}s | {async_r['p2_total_s']}s | |",
        f"| Tables created | {sync_r['p1_tables']} | {async_r['p1_tables']} | |",
        f"| Joins verified on SALES | {sync_r['joins']} | {async_r['joins']} | |",
        f"| **GRAND TOTAL** | **{sync_r['grand_total_s']}s** | **{async_r['grand_total_s']}s** "
        f"| **{_speedup(sync_r['grand_total_s'], async_r['grand_total_s'])} faster** |",
        "", "## Raw", "", "```json",
        json.dumps({"sync": sync_r, "async": async_r}, indent=2, default=str), "```", "",
    ]
    md.write_text("\n".join(lines) + "\n")
    return md


# ---- main -------------------------------------------------------------------

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--env", required=True)
    ap.add_argument("--allow-prod", action="store_true")
    ap.add_argument("--cleanup", action="store_true",
                    help="Delete THIS run's logical tables (keeps connection + schema + physical)")
    ap.add_argument("--drop-sandbox", action="store_true",
                    help="With --cleanup, also drop the schema and delete the connection")
    ap.add_argument("--yes", action="store_true")
    args = ap.parse_args()

    base_url, key = resolve_env(args.env)
    user = os.getenv("TEST_USER") or os.getenv("THOUGHTSPOT_USERNAME")
    if not base_url or not key or not user:
        raise SystemExit("Missing env URL/key or TEST_USER in .env")
    if is_prod(base_url, args.env) and not args.allow_prod:
        raise SystemExit(f"REFUSING prod ({base_url}) — pass --allow-prod to override")

    print("=" * 70)
    print("SYNC vs ASYNC TML import comparison — PLAN")
    print("=" * 70)
    print(f"  Env         : {args.env} ({base_url})")
    print(f"  Connection  : {CONNECTION}  Snowflake: {DB}.{SCHEMA}")
    print(f"  Flow        : 4-table star schema; P1 create + P2 joins; both sync & async")
    print(f"  Tag         : {TAG} on all objects")
    _cl = "KEEP everything (sandbox persists)"
    if args.cleanup:
        _cl = "delete this run's tables" + (
            " + DROP SANDBOX (schema + connection)" if args.drop_sandbox else " (keep sandbox)")
    print(f"  Cleanup     : {_cl}")
    if not args.yes:
        print("\nDRY PREVIEW — nothing created. Re-run with --yes.")
        return 0

    dep = ThoughtSpotDeployer(base_url=base_url, username=user, secret_key=key)
    if not dep.authenticate():
        raise SystemExit(f"Auth failed: {dep.last_auth_error}")
    print("\n✅ authenticated")

    print(f"▶ Ensuring connection {CONNECTION}…")
    conn_guid, conn_fqn = dep.create_connection_with_reconcile(
        CONNECTION, DB, log_progress=lambda m: print(f"    {m}"))
    print(f"  connection guid={conn_guid}")
    try:
        dep.assign_tags_to_objects([conn_guid], "DATA_SOURCE", TAG)
    except Exception as e:
        print(f"  connection tag warning: {e}")

    # Run ASYNC first (fast) then SYNC (slow), distinct prefixes → clean creates.
    async_r = run_mode(dep, conn_guid, conn_fqn, "async", "PT_ASYNC_")
    sync_r = run_mode(dep, conn_guid, conn_fqn, "sync", "PT_SYNC_")

    stamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    meta = {"env": args.env, "base_url": base_url}
    md = report(stamp, meta, sync_r, async_r)

    print(f"\n{'='*70}\nCOMPARISON\n{'='*70}")
    print(f"  SYNC  grand total: {sync_r['grand_total_s']}s "
          f"(P1 {sync_r['p1_total_s']}s, P2 {sync_r['p2_total_s']}s, 504s: "
          f"{int(sync_r['p1_gw']) + int(sync_r['p2_gw'])})")
    print(f"  ASYNC grand total: {async_r['grand_total_s']}s "
          f"(submit P1 {async_r['p1_submit_s']}s, P2 {async_r['p2_submit_s']}s, 504s: 0)")
    print(f"  Report: {md}")

    if args.cleanup:
        cleanup(dep, conn_guid, sync_r["guids"] + async_r["guids"], args.drop_sandbox)
    else:
        print("\n  Objects KEPT (connection + schema + tables), tagged PERF_TEST_DELETE.")
        print("  Use --cleanup to delete a run's tables; --cleanup --drop-sandbox to remove all.")
    return 0


if __name__ == "__main__":
    sys.exit(main())