tests/e2e_quality.py

"""
Quality regression test suite for DemoPrep.

Runs 6 pipeline tests against the live HF Space using the form UI:
  - 2 fixed         (same every run — regression baselines)
  - 2 random        (use case picked from pool, AI selects matching company)
  - 2 AI-generated  (AI picks vertical, line, function, and company)

Scoring (100 pts total):
  - Stage completion  → up to 25 pts  (research 5, ddl 7, data 8, thoughtspot 5)
  - Data quality      → up to 50 pts  (LLM grades model TML + Snowflake sample, 0-100 scaled)
  - Liveboard quality → up to 25 pts  (LLM grades liveboard TML, 0-100 scaled)

Usage:
    source demoprep/bin/activate
    python tests/e2e_quality.py
"""

import json
import os
import random
import re
import sys
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional

import requests
import yaml
from dotenv import load_dotenv
from playwright.sync_api import Page, sync_playwright

sys.path.insert(0, str(Path(__file__).parent.parent))

from llm_config import DEFAULT_LLM_MODEL

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
load_dotenv(Path(__file__).parent.parent / ".env")

BASE_URL      = os.getenv("TEST_TARGET_URL", "")  # may be overridden by --url flag at runtime
TEST_USER     = os.getenv("TEST_USER")
TEST_PASSWORD = os.getenv("TEST_PASSWORD")
TEST_NEW_PASSWORD = os.getenv("TEST_NEW_PASSWORD", "")

CONFIG_FILE = Path(__file__).parent / "quality_config.yaml"
RESULTS_DIR = Path(__file__).parent / "quality_results"
RESULTS_DIR.mkdir(exist_ok=True)

DRY_RUN = False  # set to True via --dry-run; fills form but does not click GO

STAGE_LABELS = {
    "research":    "Research",
    "ddl":         "DDL",
    "data":        "Data",
    "thoughtspot": "ThoughtSpot",
    "complete":    "Complete",
}

# ---------------------------------------------------------------------------
# Settings applied to every quality run via the Settings accordion in the UI.
# Change these here to adjust what the test runner uses.
# ---------------------------------------------------------------------------
RUN_SETTINGS = {
    "data_size":          "Medium",                     # Small=1k/50dim · Medium=10k/500dim
    "column_naming":      "Regular Case",
    "tag_name":           "TR",
    "object_prefix":      "tst",
    "share_with":         "mike.boone@thoughtspot.com",
    "geo_scope":          "USA Only",
    "ai_model":           "claude-sonnet-4-6",          # model used for this test run
    "ts_environment":     "sebe - se",                   # se-cloud having issues; run on sebe
}


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
def load_config() -> dict:
    with open(CONFIG_FILE) as f:
        return yaml.safe_load(f)


# ---------------------------------------------------------------------------
# LLM helper (uses app's configured LLM)
# ---------------------------------------------------------------------------
def _get_researcher():
    from main_research import MultiLLMResearcher
    from llm_config import map_llm_display_to_provider
    provider, model = map_llm_display_to_provider(RUN_SETTINGS["ai_model"])
    return MultiLLMResearcher(provider=provider, model=model)


def _llm(prompt: str, max_tokens: int = 300) -> str:
    researcher = _get_researcher()
    return (researcher.make_request(
        [{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        stream=False,
    ) or "").strip()


def _parse_json(text: str) -> dict:
    match = re.search(r'\{.*\}', text, re.DOTALL)
    if not match:
        raise ValueError(f"No JSON found in: {text[:200]}")
    return json.loads(match.group())


# ---------------------------------------------------------------------------
# Test case generation
# ---------------------------------------------------------------------------
def generate_ai_test_case(config: dict, exclude_list: list = None) -> dict:
    """AI picks vertical, line, function, and a matching company."""
    prompt = config["ai_generated"]["generation_prompt"]
    exclude_str = ", ".join(exclude_list) if exclude_list else "none"
    prompt = prompt.replace("{exclude_list}", exclude_str)
    data = _parse_json(_llm(prompt))
    return {
        "name":        f"AI: {data['company']} — {data['vertical']} / {data['line']} / {data['function']}",
        "type":        "ai_generated",
        "company":     data["company"],
        "company_url": data["company_url"],
        "vertical":    data["vertical"],
        "line":        data["line"],
        "function":    data["function"],
    }


def pick_random_test_case(config: dict, used_labels: set, exclude_companies: list = None) -> dict:
    """Pick a use case from the pool and select a company.

    If the pool entry has a `companies` list, pick one at random — no LLM call.
    Falls back to LLM company selection only when no list is present.
    """
    pool = config["random_pool"]["use_cases"]
    template = config["random_pool"]["company_prompt"]
    exclude_companies = exclude_companies or []

    available = [uc for uc in pool if uc["label"] not in used_labels]
    if not available:
        # All labels used — reset and allow repeats (different company still possible)
        available = pool
    uc = random.choice(available)
    used_labels.add(uc["label"])

    # Prefer the pre-seeded companies list — avoids LLM call and ensures variety
    if uc.get("companies"):
        candidates = [c for c in uc["companies"] if c["company"] not in exclude_companies]
        if not candidates:
            candidates = uc["companies"]  # all used — allow repeats rather than failing
        chosen = random.choice(candidates)
        company, company_url = chosen["company"], chosen["url"]
    else:
        exclude = ", ".join(exclude_companies)
        prompt = template.format(
            label=uc["label"],
            vertical=uc["vertical"],
            line=uc["line"],
            function=uc["function"],
            exclude_list=exclude,
        )
        try:
            data = _parse_json(_llm(prompt))
            company, company_url = data["company"], data["company_url"]
        except Exception as e:
            print(f"  ⚠️  Company selection failed ({e}), using fallback")
            company, company_url = uc["label"].split()[0], "example.com"

    return {
        "name":        f"Pool: {company} — {uc['label']}",
        "type":        "random",
        "company":     company,
        "company_url": company_url,
        "vertical":    uc["vertical"],
        "line":        uc["line"],
        "function":    uc["function"],
    }


def pick_custom_pool_test_case(config: dict, exclude_companies: list = None) -> dict:
    """Pick a Professional Services company from the custom_pool for the Custom tab."""
    pool = config.get("custom_pool", [])
    exclude_companies = exclude_companies or []
    candidates = [c for c in pool if c["company"] not in exclude_companies]
    if not candidates:
        candidates = pool
    entry = random.choice(candidates)
    return {
        "name":        f"Custom: {entry['company']} — Professional Services",
        "type":        "custom",
        "company":     entry["company"],
        "company_url": entry["url"],
        "vertical":    "* CUSTOM *",
        "line":        "",
        "function":    "",
        "context":     entry["context"].strip(),
    }


def build_test_suite(config: dict) -> list:
    """Build an 8-test suite: 2 fixed + 4 pool + 2 custom (Pro Services).

    Fixed baselines catch regressions — same companies every run.
    Pool rotates companies so the pipeline can't be tuned to specific names.
    Custom covers Professional Services (no app vertical match).
    """
    suite          = []
    used_labels    = set()
    used_companies: list[str] = []

    # 2 fixed baselines — same every run, regression anchors
    for tc in config.get("fixed_tests", []):
        suite.append({**tc, "type": "fixed"})
        used_companies.append(tc["company"])

    # 4 pool picks — random company from companies list, no repeats
    for _ in range(4):
        tc = pick_random_test_case(config, used_labels, exclude_companies=used_companies)
        used_companies.append(tc["company"])
        suite.append(tc)

    # 2 custom — Professional Services (40 customers, no app vertical match)
    for _ in range(2):
        tc = pick_custom_pool_test_case(config, exclude_companies=used_companies)
        used_companies.append(tc["company"])
        suite.append(tc)

    random.shuffle(suite)
    return suite


# ---------------------------------------------------------------------------
# Form interaction helpers
# ---------------------------------------------------------------------------
def select_gradio_dropdown(page: Page, label: str, value: str):
    """Select a value from a Gradio dropdown using aria-label (confirmed from DOM inspection)."""
    inp = page.locator(f'input[aria-label="{label}"]').first
    try:
        current_value = (inp.input_value(timeout=1000) or "").strip()
        if current_value == value:
            return
    except Exception:
        pass
    inp.click(timeout=5000)
    page.wait_for_timeout(300)
    option = page.get_by_role('option', name=value, exact=True)
    try:
        option.click(timeout=5000)
    except Exception as e:
        visible_options = page.locator('[role="option"]').all_inner_texts()
        raise RuntimeError(
            f"Dropdown {label!r} does not contain {value!r}. "
            f"Visible options: {visible_options}"
        ) from e
    page.wait_for_timeout(300)


def _fill_textbox(page: Page, placeholder: str, value: str):
    """Fill a Gradio Textbox by placeholder using JS native setter to trigger Svelte reactivity."""
    page.evaluate("""
        (args) => {
            const els = document.querySelectorAll('textarea[placeholder="' + args.placeholder + '"]');
            const el = Array.from(els).find(e => e.offsetParent !== null) || els[0];
            if (!el) return;
            const setter = Object.getOwnPropertyDescriptor(window.HTMLTextAreaElement.prototype, 'value').set;
            setter.call(el, args.value);
            el.dispatchEvent(new Event('input', { bubbles: true }));
            el.dispatchEvent(new Event('change', { bubbles: true }));
        }
    """, {"placeholder": placeholder, "value": value})
    page.wait_for_timeout(150)


def _select_dropdown_force(page: Page, label: str, value: str):
    """Select a Gradio dropdown value, forcing click even if hidden."""
    inp = page.locator(f'input[aria-label="{label}"]').first
    inp.click(force=True, timeout=5000)
    page.wait_for_timeout(300)
    page.get_by_role('option', name=value, exact=True).click(timeout=5000)
    page.wait_for_timeout(200)


def _open_settings_accordion(page: Page) -> bool:
    """
    Click the Settings accordion open. Returns True if open, False if it couldn't be opened.
    """
    data_size_input = page.locator('input[aria-label="Data Size"]').first
    try:
        if data_size_input.is_visible(timeout=500):
            return True  # already open
    except Exception:
        pass

    for selector in [
        'button:has-text("⚙️ Settings"):not([role=tab])',
        'button:has-text("⚙ Settings"):not([role=tab])',
        'button[aria-expanded]:has-text("Settings")',
    ]:
        try:
            btn = page.locator(selector).first
            if btn.count() > 0:
                btn.click(timeout=3000)
                page.wait_for_timeout(400)
                if data_size_input.is_visible(timeout=2000):
                    return True
        except Exception:
            continue

    return False


def apply_run_settings(page: Page, lb_name: str = "", tag_name: str = ""):
    """
    Open the Settings accordion and apply all RUN_SETTINGS values plus liveboard name.
    Called once, after all form dropdowns are filled — avoids the accordion being
    collapsed by a Gradio re-render triggered by Vertical/Line/Function selection.
    If the accordion can't be opened, skips settings rather than force-clicking hidden
    elements (which can open dangling dropdowns that block the GO button).
    """
    opened = _open_settings_accordion(page)
    if not opened:
        print("  ⚠️  Settings accordion could not be opened — skipping settings")
        return

    try:
        # Liveboard name — inside the accordion
        if lb_name:
            for placeholder in [
                "Auto from company URL if blank",
                "Auto-generated if blank",
            ]:
                try:
                    el = page.locator(f'textarea[placeholder="{placeholder}"]').first
                    if el.is_visible(timeout=1000):
                        el.click(click_count=3, timeout=2000)
                        el.fill(lb_name)
                        page.wait_for_timeout(200)
                        break
                except Exception:
                    continue

        _select_dropdown_force(page, "Data Size",           RUN_SETTINGS["data_size"])
        _select_dropdown_force(page, "Geographic Scope",    RUN_SETTINGS["geo_scope"])
        _select_dropdown_force(page, "Column Naming Style", RUN_SETTINGS["column_naming"])
        page.wait_for_timeout(500)  # let Svelte settle after dropdown changes
        _fill_textbox(page, "e.g. Sales_Demo (blank = no tag)",                  tag_name or RUN_SETTINGS["tag_name"])
        _fill_textbox(page, "e.g. ACME_ (blank = none)",                         RUN_SETTINGS["object_prefix"])
        _fill_textbox(page, "user@company.com or group-name (blank = no share)",  RUN_SETTINGS["share_with"])
    except Exception as e:
        print(f"  ⚠️  Settings error: {e}")


def _do_login(page: Page):
    """Fill and submit the login form, then wait for tabs."""
    page.fill('input[type=text]', TEST_USER)
    page.fill('input[type=password]', TEST_PASSWORD)
    page.click('button:has-text("Login")')
    _wait_for_visible_app_or_auth_control(page, include_login=False, timeout=90000)
    _handle_forced_password_change(page)
    page.wait_for_selector('button[role=tab]', timeout=90000)
    page.wait_for_timeout(3000)


def _wait_for_visible_app_or_auth_control(page: Page, *, include_login: bool, timeout: int):
    """Wait until a visible app tab or auth control is present."""
    page.wait_for_function(
        """
        ({ includeLogin }) => {
            const visible = (el) => !!(
                el &&
                (el.offsetWidth || el.offsetHeight || el.getClientRects().length)
            );
            const hasVisibleTab = Array.from(document.querySelectorAll('button[role="tab"]'))
                .some(visible);
            if (hasVisibleTab) return true;

            const buttons = Array.from(document.querySelectorAll('button'))
                .filter(visible)
                .map((button) => (button.textContent || '').trim());
            if (buttons.some((text) => text.includes('Change Password'))) return true;
            if (includeLogin && buttons.some((text) => text.includes('Login'))) return true;
            return false;
        }
        """,
        arg={"includeLogin": include_login},
        timeout=timeout,
    )


def _handle_forced_password_change(page: Page):
    """Handle or explicitly fail on the app's temporary-password gate."""
    try:
        gate = page.get_by_text("Change Password Required", exact=False)
        if not gate.is_visible(timeout=1500):
            return
    except Exception:
        return

    if not TEST_NEW_PASSWORD:
        raise RuntimeError(
            "Test user is blocked by the temporary-password gate. "
            "Clear must_change_password for TEST_USER or set TEST_NEW_PASSWORD "
            "so the harness can complete the required password change."
        )

    page.locator('input[placeholder="Enter the password you just used to sign in"]').first.fill(TEST_PASSWORD)
    page.locator('input[placeholder="At least 8 characters"]').first.fill(TEST_NEW_PASSWORD)
    page.locator('input[placeholder="Repeat new password"]').first.fill(TEST_NEW_PASSWORD)
    page.click('button:has-text("Change Password")', timeout=5000)
    page.wait_for_selector('button[role=tab]', timeout=30000)


def _navigate_and_ensure_logged_in(page: Page, max_wait_secs: int = 300):
    """
    Navigate to BASE_URL and ensure we're on the logged-in app.
    Handles: HF space sleeping/rebuilding after a long test, session expiry.
    Retries for up to max_wait_secs before raising.
    """
    deadline = time.time() + max_wait_secs
    attempt = 0
    while True:
        attempt += 1
        try:
            page.goto(BASE_URL, timeout=90000)
            # Wait for either the logged-in app (tabs) or the login form
            _wait_for_visible_app_or_auth_control(page, include_login=True, timeout=60000)
            break
        except Exception as nav_err:
            remaining = int(deadline - time.time())
            if remaining <= 0:
                raise RuntimeError(
                    f"Space not reachable after {max_wait_secs}s: {nav_err}"
                ) from nav_err
            print(f"  ⏳ App still loading (attempt {attempt}) — waiting 30s ({remaining}s left)...")
            time.sleep(30)

    # If we landed on the login page (session expired or space rebuilt), re-login
    try:
        if page.locator('button:has-text("Login")').is_visible(timeout=2000):
            print("  🔑 Session expired — re-logging in...")
            _do_login(page)
    except Exception:
        pass  # Already on the app — no login needed
    _handle_forced_password_change(page)


def submit_job(page: Page, test_case: dict):
    """Fill the form and click GO."""
    # Navigate, re-logging in if the session expired (e.g. after a long prior test)
    _navigate_and_ensure_logged_in(page)
    page.wait_for_timeout(2000)

    page.click('button[role=tab]:has-text("📱 App")', timeout=10000)
    page.wait_for_timeout(500)
    page.get_by_role('tab', name='App', exact=True).click(timeout=10000)
    page.wait_for_timeout(1000)

    lb_name = f"QA — {test_case['company']} {test_case.get('function', 'Demo')}"

    # AI Model and TS Environment — always visible, set before form dropdowns
    select_gradio_dropdown(page, "AI Model",       RUN_SETTINGS["ai_model"])
    select_gradio_dropdown(page, "TS Environment", RUN_SETTINGS["ts_environment"])

    # Select vertical (always set)
    select_gradio_dropdown(page, "Vertical", test_case["vertical"])

    if test_case["vertical"] == "* CUSTOM *":
        # Custom mode: wait for UI to settle after vertical dropdown change, then fill Context
        page.wait_for_timeout(1500)
        ctx_el = None
        for sel in [
            'textarea[placeholder="Describe your use case, industry, and key metrics..."]',
            '[placeholder="Describe your use case, industry, and key metrics..."]',
            'textarea[aria-label="Context *"]',
            'textarea[aria-label="Context"]',
        ]:
            try:
                el = page.locator(sel).first
                if el.is_visible(timeout=3000):
                    ctx_el = el
                    break
            except Exception:
                pass
        if ctx_el is None:
            raise Exception('Context textarea not found — tried placeholder and aria-label="Context"')
        ctx_el.click(click_count=3, timeout=5000)
        ctx_el.fill(test_case.get("context", ""))
        page.wait_for_timeout(300)
    else:
        select_gradio_dropdown(page, "Line",     test_case["line"])
        select_gradio_dropdown(page, "Function", test_case["function"])

    # Fill company URL — textarea with placeholder 'e.g. Amazon.com'
    url_el = page.locator('textarea[placeholder="e.g. Amazon.com"]')
    url_el.click(click_count=3, timeout=5000)
    url_el.fill(test_case["company_url"])
    page.wait_for_timeout(300)

    # Generate a unique tag for this test case. Logs are diagnostic-only; the
    # completed page's model/liveboard URLs identify the run under test.
    test_tag = f"TR-{uuid.uuid4().hex[:8].upper()}"
    test_case["_test_tag"] = test_tag  # store so run_single_test can pass it to diagnostics

    # Apply run settings + liveboard name — accordion opened once, after all form dropdowns
    apply_run_settings(page, lb_name=lb_name, tag_name=test_tag)

    # Click GO (skipped in dry-run mode)
    if DRY_RUN:
        print(f"  🔍 DRY RUN — form filled, pausing 120s so you can inspect the browser...")
        print(f"     Vertical={test_case['vertical']}  Line={test_case.get('line')}  Function={test_case.get('function')}")
        print(f"     URL={test_case['company_url']}  lb={lb_name}")
        print(f"     Settings: {RUN_SETTINGS}")
        time.sleep(120)
        print(f"  ⏭️  Skipping GO — dry run complete")
        return
    page.click('button:has-text("→ GO")', timeout=10000)
    print(f"  ✅ Form submitted: {test_case['vertical']} / {test_case['line']} / {test_case['function']} — {test_case['company_url']} | lb: {lb_name}")


# ---------------------------------------------------------------------------
# Pipeline stage detection
# ---------------------------------------------------------------------------
def read_progress(page: Page) -> dict:
    """
    Read pipeline progress from the right-side progress panel.
    Returns stage_key -> 'complete' | 'running' | 'not_started' | 'unknown'
    """
    # Stay on App tab — progress panel is on the right side
    try:
        page.click('button[role=tab]:has-text("📱 App")', timeout=5000)
        page.wait_for_timeout(500)
        page.get_by_role('tab', name='App', exact=True).click(timeout=3000)
        page.wait_for_timeout(300)
    except Exception:
        pass

    progress_text = page.inner_text('body')

    stages = {}
    for key, label in STAGE_LABELS.items():
        if f"✓ {label}" in progress_text or f"✅ {label}" in progress_text:
            stages[key] = "complete"
        elif f"▶ {label}" in progress_text:
            stages[key] = "running"
        elif f"○ {label}" in progress_text:
            stages[key] = "not_started"
        else:
            stages[key] = "unknown"
    return stages


def pipeline_finished(stages: dict) -> bool:
    # Done if app shows "Complete", OR if all 4 main stages are marked complete
    if stages.get("complete") == "complete":
        return True
    main_stages = ("research", "ddl", "data", "thoughtspot")
    return all(stages.get(s) == "complete" for s in main_stages)


# ---------------------------------------------------------------------------
# Post-run GUID extraction
# ---------------------------------------------------------------------------
def extract_run_context(page: Page) -> dict:
    """
    After completion, find model and liveboard URLs in the page.
    Uses visible text plus full HTML so GUIDs in href attributes are also
    matched. This is the source of truth for the run under test.
    """
    try:
        visible = page.inner_text("body", timeout=5000)
    except Exception:
        visible = ""
    body = f"{visible}\n{page.content()}"  # full HTML catches GUIDs in href attrs too

    model_match = re.search(r'(https://[^\s"\'<>#]+)/#/data/tables/([a-f0-9-]{36})', body)
    lb_match    = re.search(r'(https://[^\s"\'<>#]+)/#/pinboard/([a-f0-9-]{36})',    body)
    ts_base = None
    if model_match:
        ts_base = model_match.group(1)
    elif lb_match:
        ts_base = lb_match.group(1)

    return {
        "ts_base_url":    ts_base,
        "model_guid":     model_match.group(2) if model_match else None,
        "liveboard_guid": lb_match.group(2)    if lb_match    else None,
        "source":          "page",
    }


# ---------------------------------------------------------------------------
# ThoughtSpot API helpers
# ---------------------------------------------------------------------------
def _find_ts_key_for_url(ts_base_url: str) -> str:
    target = (ts_base_url or "").rstrip("/")
    for i in range(1, 10):
        url = os.getenv(f"TS_ENV_{i}_URL", "").rstrip("/")
        key = os.getenv(f"TS_ENV_{i}_KEY_VAR", "")
        if url and key and url == target:
            return key
    return os.getenv("TS_ENV_1_KEY_VAR", "")


def ts_authenticate(ts_base_url: str, username: str = None) -> requests.Session:
    secret_key = _find_ts_key_for_url(ts_base_url)
    if not secret_key:
        raise RuntimeError(f"No trusted auth key found for {ts_base_url}")
    auth_user = username or TEST_USER
    session = requests.Session()
    session.headers["Accept"] = "application/json"
    resp = session.post(
        f"{ts_base_url}/api/rest/2.0/auth/token/full",
        json={"username": auth_user, "secret_key": secret_key, "validity_time_in_sec": 3600},
        timeout=30,
    )
    resp.raise_for_status()
    token = resp.json().get("token")
    if token:
        session.headers["Authorization"] = f"Bearer {token}"
    return session


def export_tml(ts_base_url: str, session: requests.Session, guid: str) -> str:
    """Export TML for a liveboard or answer (JSON format)."""
    resp = session.post(
        f"{ts_base_url}/api/rest/2.0/metadata/tml/export",
        json={"metadata": [{"identifier": guid}], "export_associated": False, "export_fqn": True},
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    return data[0].get("edoc", "") if data else ""


def export_model_tml(ts_base_url: str, session: requests.Session, guid: str) -> str:
    """Export TML for a model (LOGICAL_TABLE) in YAML format — returns db/schema in tables[]."""
    resp = session.post(
        f"{ts_base_url}/api/rest/2.0/metadata/tml/export",
        json={
            "metadata": [{"identifier": guid, "type": "LOGICAL_TABLE"}],
            "export_associated": False,
            "export_fqn": True,
            "format_type": "YAML",
        },
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    edoc = data[0].get("edoc", "") if data else ""
    return edoc


def export_model_related_tmls(ts_base_url: str, session: requests.Session, guid: str) -> list[str]:
    """Export model TML plus associated table TMLs so physical db/schema can be read directly."""
    resp = session.post(
        f"{ts_base_url}/api/rest/2.0/metadata/tml/export",
        json={
            "metadata": [{"identifier": guid, "type": "LOGICAL_TABLE"}],
            "export_associated": True,
            "export_fqn": True,
            "format_type": "YAML",
        },
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    return [item.get("edoc", "") for item in (data or []) if item.get("edoc")]


def _parse_db_schema_from_fqn(fqn: str) -> tuple[str, str]:
    if not fqn or "." not in fqn:
        return "", ""
    quoted = re.findall(r'"([^"]+)"', fqn)
    if len(quoted) >= 2:
        return quoted[0], quoted[1]
    parts = [p.strip().strip('"') for p in str(fqn).split(".") if p.strip()]
    if len(parts) >= 3:
        return parts[0], parts[1]
    if len(parts) == 2:
        return parts[0], parts[1]
    return "", ""


def _walk_dicts(value):
    if isinstance(value, dict):
        yield value
        for child in value.values():
            yield from _walk_dicts(child)
    elif isinstance(value, list):
        for child in value:
            yield from _walk_dicts(child)


def extract_db_schema(model_tml_str: str) -> tuple:
    return extract_db_schema_from_tml_docs([model_tml_str])


def extract_db_schema_from_tml_docs(tml_docs: list[str]) -> tuple:
    """
    Extract physical Snowflake db/schema from model or associated table TML.
    This intentionally does not derive schema from naming convention.
    """
    try:
        for tml_str in tml_docs:
            if not tml_str:
                continue
            tml = yaml.safe_load(tml_str) or {}
            for node in _walk_dicts(tml):
                table_node = node.get("table") if isinstance(node.get("table"), dict) else {}
                db = (
                    node.get("db")
                    or node.get("database")
                    or node.get("database_name")
                    or node.get("db_name")
                    or table_node.get("db")
                    or table_node.get("database")
                    or table_node.get("database_name")
                    or table_node.get("db_name")
                    or ""
                )
                schema = (
                    node.get("schema")
                    or node.get("schema_name")
                    or table_node.get("schema")
                    or table_node.get("schema_name")
                    or ""
                )
                if db and schema:
                    return str(db), str(schema)

                for key in ("fqn", "table_fqn", "physical_table", "db_table"):
                    fqn = node.get(key) or table_node.get(key)
                    db, schema = _parse_db_schema_from_fqn(str(fqn or ""))
                    if db and schema:
                        return db, schema
    except Exception:
        pass
    return "", ""


def resolve_schema_from_model(run_context: dict) -> dict:
    """
    Resolve the Snowflake schema from the exact ThoughtSpot model printed by
    the app. No prefix/date guessing.
    """
    ts_base = run_context.get("ts_base_url")
    model_guid = run_context.get("model_guid")
    if not ts_base or not model_guid:
        return {"found": False, "reason": "missing model URL"}
    try:
        session = ts_authenticate(ts_base)
        tml_docs = export_model_related_tmls(ts_base, session, model_guid)
        db, schema = extract_db_schema_from_tml_docs(tml_docs)
        if not schema:
            return {"found": False, "reason": "model/associated table TML did not expose physical schema"}
        return {"found": True, "database": db, "schema": schema}
    except Exception as e:
        return {"found": False, "reason": str(e)}


def get_snowflake_sample(db: str, schema: str) -> str:
    try:
        from snowflake_auth import get_snowflake_connection
        conn   = get_snowflake_connection()
        cursor = conn.cursor()
        cursor.execute(f'SHOW TABLES IN SCHEMA "{db}"."{schema}"')
        tables = [row[1] for row in cursor.fetchall()]

        # Count rows in every table first — so we can prioritize the fact table
        row_counts = {}
        for table in tables:
            try:
                cursor.execute(f'SELECT COUNT(*) FROM "{db}"."{schema}"."{table}"')
                row_counts[table] = cursor.fetchone()[0]
            except Exception:
                row_counts[table] = 0

        # Sort descending — fact table (most rows) sampled first
        tables_sorted = sorted(tables, key=lambda t: row_counts.get(t, 0), reverse=True)

        # Build row-count summary header so grader knows what's populated
        header = ["Table row counts:"]
        for t in tables_sorted:
            header.append(f"  {t}: {row_counts.get(t, 0)} rows")
        empty_tables = [t for t in tables_sorted if row_counts.get(t, 0) == 0]
        if empty_tables:
            header.append(
                f"\n⚠️  WARNING: {len(empty_tables)} table(s) have 0 rows: "
                f"{', '.join(empty_tables)}"
            )

        parts = ["\n".join(header)]

        # Sample from tables that actually have data (up to 6); fall back to first 3 if all empty
        tables_with_data = [t for t in tables_sorted if row_counts.get(t, 0) > 0]
        to_sample = tables_with_data[:6] if tables_with_data else tables_sorted[:3]

        synthetic_hits = []
        for table in to_sample:
            try:
                cursor.execute(f'SELECT * FROM "{db}"."{schema}"."{table}" LIMIT 200')
                cols = [d[0] for d in cursor.description]
                rows = cursor.fetchall()
                parts.append(
                    f"\nTable: {table}  "
                    f"({row_counts.get(table, 0)} total rows, {min(len(rows), 15)} displayed / {len(rows)} scanned)"
                )
                parts.append(f"Columns: {', '.join(cols)}")
                for row in rows:
                    for col, value in zip(cols, row):
                        if isinstance(value, str) and SYNTHETIC_NUMERIC_SUFFIX_RE.search(value.strip()):
                            synthetic_hits.append((table, col, value.strip()))
                for row in rows[:15]:
                    parts.append("  " + str(dict(zip(cols, row))))
            except Exception as e:
                parts.append(f"\nTable: {table} — error: {e}")

        if synthetic_hits:
            parts.append("\nDATA QUALITY HARD FAIL CANDIDATES:")
            for table, col, value in synthetic_hits[:25]:
                parts.append(f"  {{'TABLE': '{table}', 'COLUMN': '{col}', 'SYNTHETIC_VALUE': '{value}'}}")

        cursor.close()
        conn.close()
        return "\n".join(parts) if parts else "No tables found"
    except Exception as e:
        return f"Snowflake connection failed: {e}"


# ---------------------------------------------------------------------------
# AI quality grading
# ---------------------------------------------------------------------------
def _extract_grader_json(text: str) -> dict:
    """Find the first complete JSON object containing a 'score' key."""
    decoder = json.JSONDecoder()
    idx = 0
    while idx < len(text):
        brace = text.find('{', idx)
        if brace == -1:
            break
        try:
            obj, _ = decoder.raw_decode(text, brace)
            if isinstance(obj, dict) and 'score' in obj:
                return obj
        except json.JSONDecodeError:
            pass
        idx = brace + 1
    raise ValueError(f"No JSON with 'score' key in: {text[:200]}")


def _call_grader(prompt: str, max_retries: int = 3) -> dict:
    last_raw = ""
    for attempt in range(max_retries):
        try:
            raw = _llm(prompt, max_tokens=2000)
            last_raw = raw
            return _extract_grader_json(raw)
        except (json.JSONDecodeError, ValueError, Exception):
            pass
        if attempt < max_retries - 1:
            time.sleep(3)
    print(f"    ❌ Grader parse failed — raw response: {last_raw[:300]!r}")
    return {"score": 0, "reasoning": "Could not parse response after retries",
            "strengths": [], "weaknesses": [last_raw[:200]]}


SYNTHETIC_NUMERIC_SUFFIX_RE = re.compile(
    r"\b(?:"
    r"north|south|east|west|central|northeast|northwest|southeast|southwest|"
    r"route|corridor|express|lane|zone|region|market|segment|category|"
    r"customer|account|vendor|supplier|warehouse|store|location|product|"
    r"service|plan|item|team"
    r")\b(?:[\w\s&/-]{0,80})\s+\d{1,4}$",
    re.IGNORECASE,
)
SYNTHETIC_DISTINCT_FAIL_THRESHOLD = 5
SYNTHETIC_OCCURRENCE_FAIL_THRESHOLD = 10
SYNTHETIC_WARNING_SCORE_CAP = 80
SYNTHETIC_FAILURE_SCORE = 45


def detect_synthetic_dimension_values(sample_data: str) -> dict:
    """Find generic dimension values like 'North Corridor Route 31'."""
    if not sample_data:
        return {
            "fail": False,
            "warn": False,
            "examples": [],
            "count": 0,
            "occurrences": 0,
        }

    offenders = []
    for match in re.finditer(r":\s*'([^']+)'", sample_data):
        value = match.group(1).strip()
        if SYNTHETIC_NUMERIC_SUFFIX_RE.search(value):
            offenders.append(value)
    for match in re.finditer(r':\s*"([^"]+)"', sample_data):
        value = match.group(1).strip()
        if SYNTHETIC_NUMERIC_SUFFIX_RE.search(value):
            offenders.append(value)

    distinct_offenders = sorted(set(offenders))
    fail = (
        len(distinct_offenders) >= SYNTHETIC_DISTINCT_FAIL_THRESHOLD
        or len(offenders) >= SYNTHETIC_OCCURRENCE_FAIL_THRESHOLD
    )
    return {
        "fail": fail,
        "warn": bool(offenders),
        "examples": distinct_offenders[:12],
        "count": len(distinct_offenders),
        "occurrences": len(offenders),
    }


def grade_data_quality(company: str, vertical: str, line: str, function: str,
                       model_tml: str, sample_data: str) -> dict:
    synthetic_check = detect_synthetic_dimension_values(sample_data)
    if synthetic_check["fail"]:
        examples = ", ".join(synthetic_check["examples"])
        return {
            "score": SYNTHETIC_FAILURE_SCORE,
            "reasoning": (
                "Automatic data-quality failure: Snowflake sample contains generic "
                f"synthetic dimension values ending in numbers, such as {examples}. "
                f"Detected {synthetic_check['occurrences']} occurrences across "
                f"{synthetic_check['count']} distinct values. This is a data realism failure."
            ),
            "strengths": [],
            "weaknesses": [
                "Synthetic numeric-suffix dimension values detected",
                *synthetic_check["examples"],
            ],
            "synthetic_dimension_failure": synthetic_check,
        }

    synthetic_warning = ""
    if synthetic_check["warn"]:
        examples = ", ".join(synthetic_check["examples"])
        synthetic_warning = (
            "\n\nDETERMINISTIC DATA QUALITY WARNING:\n"
            f"Detected {synthetic_check['occurrences']} generic numeric-suffix "
            f"dimension value(s), including {examples}. This is not an automatic "
            "failure at this volume, but it should reduce realism/story quality.\n"
        )

    today = datetime.now().strftime("%Y-%m-%d")
    prompt = f"""You are grading a ThoughtSpot demo dataset.

Company: {company}
Vertical: {vertical} / {line}
Analytics function: {function}
Today's date is {today}. Treat any date on or before today as HISTORICAL — do NOT
penalize current-year or recent dates as "future-dated"; the demo is built to run today.

The goal is a compelling demo with realistic data, outliers that drive a narrative,
and a schema that supports the key KPIs for this use case.

MODEL TML (full schema, column definitions, and relationships):
{model_tml[:10000]}

SNOWFLAKE DATA — actual row counts and sample rows:
{sample_data[:6000]}
{synthetic_warning}

Grade 0–100 using the actual data above. Do NOT hedge with phrases like "constrained by
partial TML" or "missing sample data" — the full TML and real row counts are provided.
Score based on what you can observe.

1. REALISM (20 pts): Values look like real {company} data at realistic scale and ranges.
2. STORY POTENTIAL (30 pts): Outliers, trends, or anomalies exist that anchor a demo narrative.
3. TIME COVERAGE (20 pts): 12–24 months of history with meaningful trends over time.
   Judge coverage relative to today's date above — recent/current-year data is historical,
   not "future"; only genuinely implausible far-future dates should count against this.
4. SCHEMA FITNESS (15 pts): Star schema design supports the key KPIs for {line} {function}.
5. COMPLETENESS (15 pts): Fact tables are well-populated (thousands of rows) with variation
   across dimensions. Dimensions have REALISTIC cardinality for what they represent — a
   handful of values for a naturally-small dimension (channel, region, tier, segment) is
   CORRECT and must NOT be penalized; entity dimensions (products, customers, accounts,
   stores) may have many. Do NOT require any fixed member count.
   RULE: If the row counts above show any key table at 0 rows, score COMPLETENESS = 0 for
   that criteria. If the fact table is 0 rows, also deduct heavily from STORY POTENTIAL.
   PENALTY: If there are generated-looking dimension labels with numeric suffixes
   such as "North Corridor Route 31", "Customer 17", or "Product 42", penalize realism.
   If this pattern is repeated or widespread, the data should fail.

Return ONLY valid JSON:
{{"score": 0, "reasoning": "...", "strengths": ["..."], "weaknesses": ["..."]}}"""
    result = _call_grader(prompt)
    if synthetic_check["warn"]:
        score = max(0, min(100, int(result.get("score", 0))))
        if score > SYNTHETIC_WARNING_SCORE_CAP:
            result["score"] = SYNTHETIC_WARNING_SCORE_CAP
            result["reasoning"] = (
                f"{result.get('reasoning', '')} Capped at {SYNTHETIC_WARNING_SCORE_CAP} "
                "because isolated synthetic numeric-suffix dimension values were detected."
            ).strip()
        result.setdefault("weaknesses", [])
        result["weaknesses"].append("Synthetic numeric-suffix dimension values detected at low volume")
        result["synthetic_dimension_warning"] = synthetic_check
    return result


def grade_liveboard_quality(company: str, vertical: str, line: str, function: str,
                            liveboard_tml: str, viz_count: int = None) -> dict:
    viz_note = ""
    if viz_count is not None:
        if viz_count < 3:
            viz_note = f"\n⚠️  WARNING: This liveboard has only {viz_count} visualization(s). Penalize heavily under Visualization Variety."
        else:
            viz_note = f"\nNote: Liveboard contains {viz_count} visualizations."

    prompt = f"""You are grading a ThoughtSpot liveboard.

Company: {company}
Vertical: {vertical} / {line}
Analytics function: {function}{viz_note}

A great liveboard opens with KPIs, shows trends with clear directionality,
and breaks down performance by dimensions — telling a story a presenter can walk through.

LIVEBOARD TML:
{liveboard_tml[:15000]}

Grade 0–100:
1. DATA COVERAGE (25 pts): All vizzes have backing data, questions use real column names.
2. TREND COHERENCE (20 pts): Line charts produce coherent time series; KPIs have time grains.
3. STORY STRUCTURE (25 pts): Flows KPIs → trends → breakdowns; walkable in a demo.
4. VISUALIZATION VARIETY (15 pts): Mix of KPIs, line charts, bar charts. Fewer than 3 vizzes = 0 pts here.
5. USE CASE ALIGNMENT (15 pts): Titles and questions match {line} {function} at {company}.

Return ONLY valid JSON:
{{"score": 0, "reasoning": "...", "strengths": ["..."], "weaknesses": ["..."]}}"""
    return _call_grader(prompt)


def run_ai_grading(run_context: dict, company: str, vertical: str, line: str, function: str,
                   schema_override: str = None, username: str = None) -> dict:
    result = {
        "data_score": None, "data_points": 0.0,
        "data_reasoning": "Not graded", "data_strengths": [], "data_weaknesses": [],
        "liveboard_score": None, "liveboard_points": 0.0,
        "liveboard_reasoning": "Not graded", "liveboard_strengths": [], "liveboard_weaknesses": [],
        "grading_errors": [],
    }

    ts_base    = run_context.get("ts_base_url")
    model_guid = run_context.get("model_guid")
    lb_guid    = run_context.get("liveboard_guid")

    if not ts_base or not model_guid:
        result["grading_errors"].append("No model URL found — pipeline may not have completed")
        return result

    try:
        session = ts_authenticate(ts_base, username=username)
    except Exception as e:
        result["grading_errors"].append(f"ThoughtSpot auth failed: {e}")
        return result

    # Data quality
    try:
        print("    🔍 Exporting model TML...")
        model_tml = export_model_tml(ts_base, session, model_guid)
        db, schema = extract_db_schema(model_tml)
        if not db or not schema:
            tml_docs = export_model_related_tmls(ts_base, session, model_guid)
            related_db, related_schema = extract_db_schema_from_tml_docs(tml_docs)
            db = db or related_db
            schema = schema or related_schema
        if (not db or not schema) and schema_override:
            from snowflake_auth import get_demo_database
            db, schema = get_demo_database(), schema_override
            print(f"    ℹ️  Using schema resolved from model: {schema_override}")
        sample = get_snowflake_sample(db, schema) if db and schema else "Could not determine db/schema"
        print("    🤖 Grading data quality...")
        dg    = grade_data_quality(company, vertical, line, function, model_tml, sample)
        score = max(0, min(100, int(dg.get("score", 0))))
        result.update({
            "data_score": score, "data_points": round(score * 0.50, 1),
            "data_reasoning": dg.get("reasoning", ""),
            "data_strengths": dg.get("strengths", []),
            "data_weaknesses": dg.get("weaknesses", []),
        })
        if dg.get("synthetic_dimension_failure"):
            result["grading_errors"].append("Synthetic numeric-suffix dimension values detected")
        reasoning = dg.get("reasoning", "")
        reasoning_short = reasoning[:400].rstrip() + ("…" if len(reasoning) > 400 else "")
        print(f"    📊 Data: {score}/100 → {result['data_points']} pts  |  {reasoning_short}")
    except Exception as e:
        result["grading_errors"].append(f"Data grading failed: {e}")

    # Liveboard quality
    if lb_guid:
        try:
            print("    🔍 Exporting liveboard TML...")
            lb_tml = export_tml(ts_base, session, lb_guid)

            # Count visualizations — a liveboard with 0 vizzes scores 0, no AI needed
            try:
                lb_parsed = yaml.safe_load(lb_tml)
                viz_count = len(lb_parsed.get("liveboard", {}).get("visualizations") or [])
            except Exception:
                viz_count = None
            result["liveboard_viz_count"] = viz_count
            print(f"    📋 Liveboard vizzes: {viz_count}")

            if viz_count == 0:
                lb_score = 0
                result.update({
                    "liveboard_score": 0, "liveboard_points": 0.0,
                    "liveboard_reasoning": "Liveboard has 0 visualizations — empty board.",
                    "liveboard_strengths": [],
                    "liveboard_weaknesses": ["No visualizations in liveboard TML"],
                })
                result["grading_errors"].append("Liveboard has 0 visualizations")
                print("    ❌ Liveboard is empty (0 vizzes) — score 0")
            else:
                print("    🤖 Grading liveboard quality...")
                lg       = grade_liveboard_quality(company, vertical, line, function,
                                                   lb_tml, viz_count=viz_count)
                lb_score = max(0, min(100, int(lg.get("score", 0))))
                result.update({
                    "liveboard_score": lb_score, "liveboard_points": round(lb_score * 0.25, 1),
                    "liveboard_reasoning": lg.get("reasoning", ""),
                    "liveboard_strengths": lg.get("strengths", []),
                    "liveboard_weaknesses": lg.get("weaknesses", []),
                })
                print(f"    📊 Liveboard: {lb_score}/100 → {result['liveboard_points']} pts")
        except Exception as e:
            result["grading_errors"].append(f"Liveboard grading failed: {e}")
    else:
        result["grading_errors"].append("No liveboard GUID — liveboard may not have been created")

    return result


# ---------------------------------------------------------------------------
# Group 1 settings verification
# ---------------------------------------------------------------------------
def verify_group1_settings(result: dict) -> dict:
    """
    Verify Group 1 settings are reflected in the run output.
    Checks: row counts, naming prefix, tag, share_with, column_naming_style, geo_scope.
    Returns dict of {setting: {expected, actual, pass, note}}.
    """
    checks = {}
    sf          = result.get("snowflake_check", {})
    run_ctx     = result.get("run_context", {})
    ts_base     = run_ctx.get("ts_base_url")
    model_guid  = run_ctx.get("model_guid")
    lb_guid     = run_ctx.get("liveboard_guid")

    # Load testrunner settings for expected values
    try:
        from supabase_client import SupabaseSettings
        raw = SupabaseSettings().load_all_settings(TEST_USER)
    except Exception:
        raw = {}

    # ── 1. fact_table_size ────────────────────────────────────────
    expected_fact = int(raw.get("fact_table_size") or 1000)
    if sf.get("found"):
        # Prefer name-based detection, fall back to largest table
        fact_table = next(
            (t for t in sf["tables"] if "SALES" in t["table"].upper() or "FACT" in t["table"].upper()),
            None
        )
        if fact_table is None and sf["tables"]:
            fact_table = max(sf["tables"], key=lambda t: t["rows"])
        if fact_table:
            checks["fact_table_size"] = {
                "expected": expected_fact, "actual": fact_table["rows"],
                "pass": fact_table["rows"] == expected_fact,
            }

    # ── 2. dim_table_size ────────────────────────────────────────
    expected_dim = int(raw.get("dim_table_size") or 100)
    if sf.get("found"):
        # Exclude fact-sized tables by row count — avoids hardcoded name list failures
        dim_tables = [t for t in sf["tables"] if t["rows"] < expected_fact]
        if dim_tables:
            mismatches = [t for t in dim_tables if t["rows"] != expected_dim]
            checks["dim_table_size"] = {
                "expected": expected_dim,
                "actual": {t["table"]: t["rows"] for t in dim_tables},
                "pass": len(mismatches) == 0,
                "note": f"{len(mismatches)} dim tables don't match" if mismatches else "all match",
            }

    # ── 3. object_naming_prefix ───────────────────────────────────
    expected_prefix = (raw.get("object_naming_prefix") or "").upper()
    if sf.get("found") and sf.get("schema"):
        schema = sf["schema"]
        if expected_prefix:
            passed = schema.upper().startswith(expected_prefix)
        else:
            passed = True  # no prefix expected, anything goes
        checks["object_naming_prefix"] = {
            "expected": expected_prefix or "(blank)",
            "actual": schema, "pass": passed,
        }

    # ── 4. geo_scope ─────────────────────────────────────────────
    expected_geo = raw.get("geo_scope", "USA Only")
    if sf.get("found") and sf.get("schema"):
        try:
            from snowflake_auth import get_snowflake_connection
            conn   = get_snowflake_connection()
            cursor = conn.cursor()
            schema = sf["schema"]
            database = sf["database"]
            # Look for a column named COUNTRY, REGION, or STATE
            cursor.execute(f'SHOW TABLES IN SCHEMA "{database}"."{schema}"')
            tables = [row[1] for row in cursor.fetchall()]
            foreign_found = False
            checked = False
            for tname in tables:
                cursor.execute(f'SHOW COLUMNS IN TABLE "{database}"."{schema}"."{tname}"')
                cols = [row[2].upper() for row in cursor.fetchall()]
                if "COUNTRY" in cols:
                    cursor.execute(f'SELECT DISTINCT "COUNTRY" FROM "{database}"."{schema}"."{tname}" LIMIT 20')
                    countries = [row[0] for row in cursor.fetchall() if row[0]]
                    non_us = [c for c in countries if c not in ("USA", "US", "United States", "United States of America")]
                    foreign_found = len(non_us) > 0
                    checked = True
                    break
            cursor.close(); conn.close()
            if checked:
                if expected_geo == "USA Only":
                    checks["geo_scope"] = {
                        "expected": "USA Only", "actual": f"foreign countries: {non_us}" if foreign_found else "USA only",
                        "pass": not foreign_found,
                    }
                else:
                    checks["geo_scope"] = {
                        "expected": "International", "actual": f"foreign countries found: {not foreign_found}",
                        "pass": foreign_found,
                    }
        except Exception as e:
            checks["geo_scope"] = {"pass": None, "note": f"geo check failed: {e}"}

    # ── 5. column_naming_style ────────────────────────────────────
    expected_style = raw.get("column_naming_style", "Regular Case")
    if ts_base and model_guid:
        try:
            session = ts_authenticate(ts_base)
            model_tml_str = export_tml(ts_base, session, model_guid)
            tml = yaml.safe_load(model_tml_str)
            columns = []
            for tbl in (tml.get("model", {}).get("tables") or []):
                for col in (tbl.get("columns") or []):
                    name = col.get("name", "")
                    if name:
                        columns.append(name)
            if columns:
                snake_count = sum(1 for c in columns if "_" in c and c == c.lower())
                is_snake = snake_count > len(columns) * 0.5
                actual_style = "snake_case" if is_snake else "Regular Case"
                checks["column_naming_style"] = {
                    "expected": expected_style, "actual": actual_style,
                    "pass": actual_style == expected_style,
                    "sample": columns[:5],
                }
        except Exception as e:
            checks["column_naming_style"] = {"pass": None, "note": f"TML check failed: {e}"}

    # ── 6. tag_name ───────────────────────────────────────────────
    expected_tag = raw.get("tag_name", "")
    if expected_tag and ts_base and model_guid:
        try:
            session = ts_authenticate(ts_base)
            resp = session.get(
                f"{ts_base}/tspublic/v1/metadata/list",
                params={"type": "LOGICAL_TABLE", "batchsize": 1,
                        "offset": 0, "pattern": model_guid},
            )
            body = resp.text.strip()
            if not body:
                checks["tag_name"] = {"pass": None, "note": "tag check skipped: empty API response"}
            else:
                try:
                    data = resp.json()
                except Exception:
                    data = None
                if data is None:
                    checks["tag_name"] = {"pass": None, "note": "tag check skipped: non-JSON API response"}
                else:
                    headers_data = data.get("headers", []) if isinstance(data, dict) else []
                    obj_tags = []
                    for h in headers_data:
                        if h.get("id") == model_guid:
                            obj_tags = [t.get("name", "") for t in (h.get("tags") or [])]
                            break
                    checks["tag_name"] = {
                        "expected": expected_tag, "actual": obj_tags,
                        "pass": expected_tag in obj_tags,
                    }
        except Exception as e:
            checks["tag_name"] = {"pass": None, "note": f"tag check failed: {e}"}

    # ── 7. share_with ────────────────────────────────────────────
    expected_share = raw.get("share_with", "")
    if expected_share and ts_base and model_guid:
        try:
            session = ts_authenticate(ts_base)
            resp = session.post(
                f"{ts_base}/api/rest/2.0/security/metadata/fetch",
                json={"metadata": [{"type": "LOGICAL_TABLE", "identifier": model_guid}]},
                timeout=15,
            )
            body = resp.text.strip()
            if not body:
                checks["share_with"] = {"pass": None, "note": "share check skipped: empty API response"}
            else:
                perms = resp.json()
                principals = []
                for item in (perms if isinstance(perms, list) else []):
                    for p in (item.get("permissions") or []):
                        principals.append(p.get("principal", {}).get("name", ""))
                checks["share_with"] = {
                    "expected": expected_share, "actual": principals,
                    "pass": any(expected_share.lower() in p.lower() for p in principals),
                }
        except Exception as e:
            checks["share_with"] = {"pass": None, "note": f"share check failed: {e}"}

    return checks


def print_settings_verification(checks: dict):
    if not checks:
        return
    print("  ── Settings Verification ────────────────────────────")
    for setting, result in checks.items():
        if result.get("pass") is True:
            icon = "✅"
        elif result.get("pass") is False:
            icon = "❌"
        else:
            icon = "⚠️ "
        exp = result.get("expected", "")
        act = result.get("actual", result.get("note", ""))
        print(f"  {icon} {setting}: expected={exp!r}  actual={str(act)[:60]}")
    print("  ─────────────────────────────────────────────────────")


# ---------------------------------------------------------------------------
# Stage grading
# ---------------------------------------------------------------------------
def grade_stages(stages: dict, config: dict) -> dict:
    weights   = config["grading"]["stages"]
    breakdown = {}
    total     = 0
    for key, weight in weights.items():
        status = stages.get(key, "unknown")
        earned = weight if status == "complete" else (weight // 2 if status == "running" else 0)
        breakdown[key] = {"weight": weight, "earned": earned, "status": status}
        total += earned
    return {"stage_total": total, "breakdown": breakdown}


def reconcile_stages_with_logs(stages: dict, diag: dict) -> dict:
    """
    If session_logs confirms stages completed that the UI monitor missed
    (e.g. slow DDL that triggered the 20-min bail-out), upgrade those stages to 'complete'.
    Returns a new dict — does not mutate the original.
    """
    if not diag.get("found"):
        return stages

    completed_in_logs = set(diag.get("stages_completed", []))

    # Map session_log stage names → UI progress keys
    log_to_ui = {
        "research":    "research",
        "ddl":         "ddl",
        "deploy":      "data",    # app deploy stage creates/loads Snowflake data
        "populate":    "data",    # session calls it 'populate', UI shows 'Data'
        "data":        "data",
        "thoughtspot": "thoughtspot",
    }

    stages = dict(stages)  # copy — don't mutate
    for log_stage, ui_key in log_to_ui.items():
        if log_stage in completed_in_logs and stages.get(ui_key) != "complete":
            old = stages.get(ui_key, "unknown")
            stages[ui_key] = "complete"
            print(f"    ℹ️  Stage '{ui_key}' upgraded to complete via session_logs (monitor saw: {old})")

    # Synthetic 'complete' key — set if all main stages now complete
    main = ("research", "ddl", "data", "thoughtspot")
    if all(stages.get(s) == "complete" for s in main):
        stages["complete"] = "complete"

    return stages


def _effective_stuck_threshold(stages: dict) -> int:
    """Return stage-stale timeout seconds for the currently running UI stage."""
    if any(k == "data" and v == "running" for k, v in stages.items()):
        # Complex schemas with multiple fact tables can take 25-30 min.
        return 35 * 60
    if any(k == "thoughtspot" and v == "running" for k, v in stages.items()):
        # ThoughtSpot table import + model semantics + MCP liveboard creation can
        # legitimately sit on the same UI stage for 30+ minutes.
        return 40 * 60
    return 20 * 60


def compute_grade(score: float, config: dict) -> str:
    grade = "F"
    for letter, threshold in sorted(config["grading"]["thresholds"].items(), key=lambda x: -x[1]):
        if score >= threshold:
            grade = letter
            break
    return grade


# ---------------------------------------------------------------------------
# Failure diagnostics — Supabase session_logs + Snowflake verification
# ---------------------------------------------------------------------------
def fetch_run_diagnostics(start_time: float, company: str = "", test_tag: str = "") -> dict:
    """
    Query session_logs for entries by testrunner after start_time.
    Matches only by exact test tag. This is intentionally not used to identify
    the model/schema under test because concurrent runs can contaminate logs.
    """
    try:
        from supabase_client import SupabaseSettings
        from datetime import datetime, timezone

        start_iso = datetime.fromtimestamp(start_time, tz=timezone.utc).isoformat()
        s = SupabaseSettings()
        result = (
            s.client.table("session_logs")
            .select("*")
            .eq("user_email", TEST_USER)
            .gte("ts", start_iso)
            .order("ts", desc=False)
            .execute()
        )

        if not result.data:
            return {"found": False, "reason": "No session_logs entries found after test start"}

        logs = result.data

        # Group by session_id
        sessions = {}
        for log in logs:
            sid = log["session_id"]
            sessions.setdefault(sid, []).append(log)

        if test_tag:
            tag_matching = {
                sid: entries for sid, entries in sessions.items()
                if any((l.get("meta") or {}).get("test_tag") == test_tag for l in entries)
            }
            if not tag_matching:
                return {"found": False, "reason": f"Exact test tag not found in session_logs: {test_tag}"}
            session_logs = max(tag_matching.values(), key=len)
        else:
            return {"found": False, "reason": "No test tag provided; refusing to guess session"}

        # Summarise
        completed  = [l["stage"] for l in session_logs if "completed" in (l.get("event") or "")]
        errors     = [l for l in session_logs if l.get("error")]
        last       = session_logs[-1]

        # Pull everything useful from meta fields
        total_rows = 0
        tables_populated = 0
        schema_name = None
        model_guid = None
        liveboard_guid = None
        ts_base_url = None
        for l in session_logs:
            meta = l.get("meta") or {}
            if "total_rows" in meta:
                total_rows = meta["total_rows"]
                tables_populated = meta.get("tables", 0)
            if "schema" in meta:
                schema_name = meta["schema"]
            if "schema_name" in meta:
                schema_name = meta["schema_name"]
            if "model_guid" in meta and meta["model_guid"]:
                model_guid = meta["model_guid"]
            if "liveboard_guid" in meta and meta["liveboard_guid"]:
                liveboard_guid = meta["liveboard_guid"]
            if "ts_url" in meta and meta["ts_url"]:
                ts_base_url = meta["ts_url"]

        return {
            "found":            True,
            "session_id":       session_logs[0]["session_id"],
            "stages_completed": completed,
            "last_stage":       last.get("stage"),
            "last_event":       last.get("event"),
            "snowflake_schema": schema_name,
            "total_rows":       total_rows,
            "tables_populated": tables_populated,
            "model_guid":       model_guid,
            "liveboard_guid":   liveboard_guid,
            "ts_base_url":      ts_base_url,
            "errors": [
                {
                    "stage": l["stage"],
                    "event": l.get("event"),
                    "error": (l["error"] or "")[:300],
                }
                for l in errors
            ],
            "log_count": len(session_logs),
        }

    except Exception as e:
        return {"found": False, "reason": f"Diagnostics query failed: {e}"}


def fetch_monitor_diagnostics(start_time: float, test_tag: str = "") -> dict:
    """
    Lightweight exact-tag session log check used while the UI monitor is running.
    This prevents the harness from declaring a stage stuck while the backend is
    still logging progress for the same run.
    """
    try:
        from supabase_client import SupabaseSettings
        from datetime import datetime, timezone

        if not test_tag:
            return {"found": False, "reason": "No test tag provided"}

        start_iso = datetime.fromtimestamp(start_time, tz=timezone.utc).isoformat()
        s = SupabaseSettings()
        result = (
            s.client.table("session_logs")
            .select("session_id,ts,stage,event,error,meta")
            .eq("user_email", TEST_USER)
            .gte("ts", start_iso)
            .order("ts", desc=False)
            .execute()
        )
        logs = result.data or []
        if not logs:
            return {"found": False, "reason": "No session_logs entries found after test start"}

        sessions = {}
        for log in logs:
            sid = log.get("session_id")
            if sid:
                sessions.setdefault(sid, []).append(log)

        tag_matching = {
            sid: entries for sid, entries in sessions.items()
            if any((l.get("meta") or {}).get("test_tag") == test_tag for l in entries)
        }
        if not tag_matching:
            return {"found": False, "reason": f"Exact test tag not found in session_logs: {test_tag}"}

        session_logs = max(tag_matching.values(), key=len)
        last = session_logs[-1]
        completed = [l["stage"] for l in session_logs if "completed" in (l.get("event") or "")]
        errors = [l for l in session_logs if l.get("error")]
        model_guid = None
        liveboard_guid = None
        ts_base_url = None
        for l in session_logs:
            meta = l.get("meta") or {}
            model_guid = meta.get("model_guid") or model_guid
            liveboard_guid = meta.get("liveboard_guid") or liveboard_guid
            ts_base_url = meta.get("ts_url") or ts_base_url

        last_ts = datetime.fromisoformat(str(last.get("ts", "")).replace("Z", "+00:00"))
        if last_ts.tzinfo is None:
            last_ts = last_ts.replace(tzinfo=timezone.utc)
        age_s = int((datetime.now(timezone.utc) - last_ts).total_seconds())

        return {
            "found": True,
            "session_id": session_logs[0].get("session_id"),
            "last_stage": last.get("stage"),
            "last_event": last.get("event"),
            "last_ts": last.get("ts"),
            "last_age_s": age_s,
            "stages_completed": completed,
            "errors": errors,
            "model_guid": model_guid,
            "liveboard_guid": liveboard_guid,
            "ts_base_url": ts_base_url,
            "log_count": len(session_logs),
        }
    except Exception as e:
        return {"found": False, "reason": f"Monitor diagnostics query failed: {e}"}


def check_snowflake_schema(company: str, start_time: float, schema_override: str = None,
                           database_override: str = None) -> dict:
    """
    Get row counts for the Snowflake schema created during this test run.
    Requires an explicit database + schema resolved from the exact ThoughtSpot
    model (demos live in a rotating <base>_<YYYY_MM> database, so the database
    must come from the model TML too). Never guesses by company/date prefix.
    """
    try:
        from snowflake_auth import get_snowflake_connection

        if not schema_override:
            return {"found": False, "reason": "No explicit schema provided; refusing to guess"}
        if not database_override:
            return {"found": False, "reason": "No explicit database provided; refusing to guess"}
        conn   = get_snowflake_connection()
        cursor = conn.cursor()

        schema = schema_override
        database = database_override
        cursor.execute(f'SHOW TABLES IN SCHEMA "{database}"."{schema}"')
        tables = cursor.fetchall()

        table_info = []
        total_rows = 0
        for t in tables:
            tname = t[1]
            try:
                cursor.execute(f'SELECT COUNT(*) FROM "{database}"."{schema}"."{tname}"')
                count = cursor.fetchone()[0]
                total_rows += count
                table_info.append({"table": tname, "rows": count})
            except Exception:
                table_info.append({"table": tname, "rows": "error"})

        cursor.close()
        conn.close()
        return {"found": True, "database": database, "schema": schema,
                "tables": table_info, "total_rows": total_rows}

    except Exception as e:
        return {"found": False, "error": str(e)}


def choose_snowflake_schema_for_check(schema_resolution: dict, diag: dict) -> tuple[Optional[str], str]:
    """Choose the safest explicit schema for Snowflake verification.

    Prefer the schema resolved from the exact ThoughtSpot model. If model
    creation failed, fall back to the exact-tag session log schema. Never guess
    by company or date prefix.
    """
    if schema_resolution.get("found") and schema_resolution.get("schema"):
        return schema_resolution.get("schema"), "model"
    if diag.get("snowflake_schema"):
        return diag.get("snowflake_schema"), "session_logs"
    return None, "none"


def print_diagnostics(diag: dict, sf: dict):
    """Print a human-readable failure summary."""
    print("  ── Diagnostics ──────────────────────────────────────")
    if diag.get("found"):
        print(f"  Session:  {diag['session_id']}")
        print(f"  Last:     [{diag['last_stage']}] {diag['last_event']}")
        if diag["stages_completed"]:
            print(f"  Done:     {', '.join(diag['stages_completed'])}")
        if diag["total_rows"]:
            print(f"  Snowflake (from logs): {diag['tables_populated']} tables, {diag['total_rows']} rows")
        for err in diag.get("errors", []):
            msg = (err["error"] or "")[:120].replace("\n", " ")
            print(f"  ❌ [{err['stage']}] {msg}")
    else:
        print(f"  Supabase: {diag.get('reason', 'no data')}")

    if sf.get("found"):
        print(f"  Snowflake schema: {sf['schema']}  ({sf['total_rows']} rows across {len(sf['tables'])} tables)")
        for t in sf["tables"]:
            print(f"    {t['table']}: {t['rows']} rows")
    elif sf:
        print(f"  Snowflake: {sf.get('reason') or sf.get('error') or 'schema not found'}")
    print("  ─────────────────────────────────────────────────────")


# ---------------------------------------------------------------------------
# Single test runner
# ---------------------------------------------------------------------------
def run_single_test(page: Page, test_case: dict, config: dict) -> dict:
    timeout_sec = config["grading"]["timeout_minutes"] * 60
    start       = time.time()

    result = {
        "name": test_case["name"], "type": test_case["type"],
        "company": test_case.get("company", ""),
        "vertical": test_case.get("vertical", ""), "line": test_case.get("line", ""),
        "function": test_case.get("function", ""), "company_url": test_case.get("company_url", ""),
        "stages": {}, "run_context": {}, "stage_grading": {}, "ai_grading": {},
        "total_score": 0.0, "grade": "F",
        "error": None, "timed_out": False, "late_complete": False, "duration_seconds": 0,
        "diagnostics": {}, "snowflake_check": {},
        "liveboard_viz_count": None,
    }

    try:
        submit_job(page, test_case)
        print(f"  ⏳ Monitoring pipeline (timeout: {config['grading']['timeout_minutes']}min)...")

        poll_interval    = 15
        last_stages      = {}
        last_change_time = time.time()
        STUCK_THRESHOLD  = 20 * 60  # 20 min with no stage change → bail early
        BACKEND_STALE_THRESHOLD = 15 * 60
        last_monitor_log_count = 0
        PIPELINE_ERROR_INDICATORS = [
            "Research failed",
            "pipeline has been interrupted",
            "An unexpected error occurred",
            "Population failed",
            "Pipeline failed",
            "Something went wrong during the pipeline",
            "Traceback (most recent call last)",
            "NameError:",
        ]
        while time.time() - start < timeout_sec:
            time.sleep(poll_interval)
            stages = read_progress(page)
            if stages != last_stages:
                done    = [k for k, v in stages.items() if v == "complete"]
                running = [k for k, v in stages.items() if v == "running"]
                print(f"    ✓ {done}  ▶ {running}")
                last_stages      = stages
                last_change_time = time.time()
            if pipeline_finished(stages):
                print("  ✅ Pipeline complete")
                time.sleep(10)  # let final output (URLs) finish rendering before extraction
                break
            # Detect hard pipeline failure in the chat output
            try:
                page_text = page.inner_text('body')
                if any(ind in page_text for ind in PIPELINE_ERROR_INDICATORS):
                    result["timed_out"] = True
                    print("  ❌ Pipeline error detected in page — stopping early")
                    break
            except Exception:
                pass
            # Bail if stages have been stuck for too long, with stage-specific
            # allowances for slow data generation and ThoughtSpot object creation.
            stuck_secs = time.time() - last_change_time
            effective_threshold = _effective_stuck_threshold(stages)
            if stuck_secs > effective_threshold and any(v == "running" for v in stages.values()):
                monitor_diag = fetch_monitor_diagnostics(start, test_case.get("_test_tag", ""))
                if monitor_diag.get("found"):
                    completed = set(monitor_diag.get("stages_completed", []))
                    errors = monitor_diag.get("errors") or []
                    backend_fresh = monitor_diag.get("last_age_s", 999999) <= BACKEND_STALE_THRESHOLD
                    backend_has_model = bool(monitor_diag.get("model_guid"))
                    if backend_has_model:
                        result["run_context"] = {
                            "ts_base_url": monitor_diag.get("ts_base_url"),
                            "model_guid": monitor_diag.get("model_guid"),
                            "liveboard_guid": monitor_diag.get("liveboard_guid"),
                            "source": "session_logs_monitor",
                        }
                        print("  ℹ️  Backend completed via session_logs while UI was still stale")
                        break
                    if backend_fresh and not errors:
                        last_change_time = time.time()
                        if monitor_diag.get("log_count", 0) != last_monitor_log_count:
                            last_monitor_log_count = monitor_diag.get("log_count", 0)
                            print(
                                "    ℹ️  UI stage stale, but backend still active: "
                                f"[{monitor_diag.get('last_stage')}] {monitor_diag.get('last_event')}"
                            )
                        continue
                    if "thoughtspot" in completed and not errors:
                        last_change_time = time.time()
                        print("    ℹ️  ThoughtSpot stage completed in logs; waiting for model/liveboard GUIDs")
                        continue
                result["timed_out"] = True
                print(f"  ⏰ Stage stuck for {int(stuck_secs/60)}min — treating as failure")
                break
        else:
            result["timed_out"] = True
            print(f"  ⏰ Timed out after {config['grading']['timeout_minutes']} min")

        result["stages"] = last_stages or read_progress(page)

        # GUIDs normally come from the final page. The monitor may also set
        # run_context if session_logs prove the backend completed after the UI
        # became stale.
        result["run_context"] = result.get("run_context") or {}

    except Exception as e:
        result["error"] = str(e)
        print(f"  ❌ Error: {e}")
        try:
            result["stages"] = read_progress(page)
        except Exception:
            pass

    result["duration_seconds"] = round(time.time() - start)

    # The completed page prints the exact model/liveboard URLs for this run.
    # Treat that as authoritative; session_logs are diagnostics only.
    page_ctx = extract_run_context(page)
    if page_ctx.get("model_guid"):
        result["run_context"] = page_ctx
    if page_ctx.get("model_guid"):
        lb_note = page_ctx.get("liveboard_guid", "")
        print(f"  🔑 GUIDs from page: model={page_ctx['model_guid'][:8]}… lb={lb_note[:8] if lb_note else 'none'}…")
    elif result["run_context"].get("model_guid"):
        lb_note = result["run_context"].get("liveboard_guid", "")
        print(
            f"  🔑 GUIDs from {result['run_context'].get('source', 'session logs')}: "
            f"model={result['run_context']['model_guid'][:8]}… "
            f"lb={lb_note[:8] if lb_note else 'none'}…"
        )
    else:
        print("  ⚠️  No model GUID found on final page — AI grading will be skipped")

    # Fetch exact-tag diagnostics only for supplemental errors/stage reconciliation.
    diag = fetch_run_diagnostics(start, company=result.get("company", ""),
                                 test_tag=test_case.get("_test_tag", ""))
    result["diagnostics"] = diag
    if not diag.get("found"):
        print(f"  ℹ️  Session logs not used for identity: {diag.get('reason')}")
    elif not result["run_context"].get("model_guid") and diag.get("model_guid"):
        result["run_context"] = {
            "ts_base_url": diag.get("ts_base_url"),
            "model_guid": diag.get("model_guid"),
            "liveboard_guid": diag.get("liveboard_guid"),
            "source": "session_logs_exact_tag",
        }
        lb_note = result["run_context"].get("liveboard_guid", "")
        print(
            f"  🔑 GUIDs from exact-tag session logs: "
            f"model={result['run_context']['model_guid'][:8]}… "
            f"lb={lb_note[:8] if lb_note else 'none'}…"
        )

    schema_resolution = resolve_schema_from_model(result["run_context"])
    known_schema, schema_source = choose_snowflake_schema_for_check(schema_resolution, diag)
    if schema_source == "model":
        print(f"  🧭 Schema from ThoughtSpot model: {known_schema}")
    else:
        print(f"  ⚠️  Could not resolve schema from model: {schema_resolution.get('reason')}")
        if schema_source == "session_logs":
            print(f"  🧭 Schema from exact-tag session logs: {known_schema}")
    known_database = schema_resolution.get("database")
    if not known_database and known_schema:
        # Schema came from session logs (model export failed): this run just
        # wrote to the active rotating demo database, so that IS its database.
        from snowflake_auth import get_demo_database
        known_database = get_demo_database()
    sf = check_snowflake_schema(result["company"], start, schema_override=known_schema,
                                database_override=known_database)
    result["snowflake_check"] = sf
    if sf.get("found"):
        print(f"  📦 Snowflake: {sf['schema']}  ({sf['total_rows']} rows, {len(sf['tables'])} tables)")
        for t in sf["tables"]:
            print(f"    {t['table']}: {t['rows']} rows")
    else:
        print(f"  📦 Snowflake: schema not checked ({sf.get('reason') or sf.get('error') or 'unknown'})")

    # Reconcile stages with session_logs — catches cases where the monitor bailed early
    # but the pipeline actually completed (e.g. slow DDL that took >20 min)
    result["stages"] = reconcile_stages_with_logs(result["stages"], diag)
    main_stages = ("research", "ddl", "data", "thoughtspot")
    if result["timed_out"] and all(result["stages"].get(s) == "complete" for s in main_stages):
        result["late_complete"] = True
        print("  ℹ️  Pipeline completed late — all stages confirmed via session_logs")

    # Print diagnostics on timeout/error — skip for late_complete since pipeline did finish
    if result["error"] or (result["timed_out"] and not result["late_complete"]):
        print_diagnostics(diag, {})  # sf already printed above

    # Stage scoring
    sg = grade_stages(result["stages"], config)
    result["stage_grading"] = sg

    # AI grading
    ag = {"data_points": 0.0, "liveboard_points": 0.0, "grading_errors": []}
    if result["run_context"].get("model_guid"):
        print("  🔬 Running AI quality grading...")
        ag = run_ai_grading(
            result["run_context"],
            result["company"], result["vertical"], result["line"], result["function"],
            schema_override=known_schema,
        )
    else:
        ag["grading_errors"].append("Skipped — no model GUID (pipeline did not complete)")
    result["ai_grading"] = ag
    result["liveboard_viz_count"] = ag.get("liveboard_viz_count")

    total = sg["stage_total"] + ag.get("data_points", 0) + ag.get("liveboard_points", 0)
    result["total_score"] = round(total, 1)
    result["grade"]       = compute_grade(total, config)
    return result


# ---------------------------------------------------------------------------
# Results
# ---------------------------------------------------------------------------
def save_to_postgres(run: dict, env_name: str = "") -> None:
    """Write one row per test result into ts_quality_results in Supabase."""
    try:
        sys.path.insert(0, str(Path(__file__).parent.parent))
        from supabase_client import SupabaseSettings
        ss = SupabaseSettings()

        rows = []
        for t in run["tests"]:
            ag   = t.get("ai_grading", {})
            sg   = t.get("stage_grading", {})
            diag = t.get("diagnostics", {})
            ctx  = t.get("run_context") or {}
            sf   = t.get("snowflake_check", {})
            rows.append({
                "run_id":               run["run_id"],
                "run_timestamp":        run["timestamp"],
                "environment":          env_name or ("prod" if "test" not in run.get("target_url","") else "test"),
                "target_url":           run.get("target_url",""),
                "run_avg_score":        run["avg_score"],
                "run_overall_grade":    run["overall_grade"],
                "company":              t.get("company", t["name"]),
                "vertical":             t.get("vertical",""),
                "line":                 t.get("line",""),
                "function":             t.get("function",""),
                "test_type":            t.get("type",""),
                "total_score":          t.get("total_score", 0),
                "grade":                t.get("grade","F"),
                "stage_score":          sg.get("stage_total", 0),
                "data_score":           ag.get("data_score"),
                "data_points":          ag.get("data_points"),
                "liveboard_score":      ag.get("liveboard_score"),
                "liveboard_points":     ag.get("liveboard_points"),
                "timed_out":            t.get("timed_out", False),
                "late_complete":        t.get("late_complete", False),
                "error":                t.get("error"),
                "duration_seconds":     t.get("duration_seconds"),
                "session_id":           diag.get("session_id",""),
                "model_guid":           diag.get("model_guid","") or ctx.get("model_guid",""),
                "liveboard_guid":       diag.get("liveboard_guid","") or ctx.get("liveboard_guid",""),
                "ts_base_url":          diag.get("ts_base_url","") or ctx.get("ts_base_url",""),
                "snowflake_schema":     sf.get("schema",""),
                "total_rows":           sf.get("total_rows"),
                "tables_populated":     diag.get("tables_populated"),
                "liveboard_viz_count":  t.get("liveboard_viz_count"),
                "data_reasoning":       ag.get("data_reasoning",""),
                "liveboard_reasoning":  ag.get("liveboard_reasoning",""),
            })

        ss.client.table("ts_quality_results").insert(rows).execute()
        print(f"📊 Saved {len(rows)} rows to ts_quality_results")
    except Exception as e:
        print(f"⚠️  Postgres write failed (non-fatal): {e}")


def save_results(run: dict) -> Path:
    ts   = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    path = RESULTS_DIR / f"{ts}_quality_run.json"
    with open(path, "w") as f:
        json.dump(run, f, indent=2, default=str)
    print(f"\n💾 Results: {path}")
    return path


def save_summary_md(run: dict, json_path: Path, env_name: str = "") -> Path:
    """
    Save a compact markdown summary alongside the JSON.
    Also writes to latest_{env_name}_summary.md (or latest_summary.md) for nightly use.
    env_name: 'test' | 'prod' | '' (default)
    """
    W = 32
    lines = [
        f"# DemoPrep Quality Run — {run['timestamp'][:16]}",
        f"**Target:** {run.get('target_url', 'unknown')}  |  "
        f"**Avg:** {run['avg_score']}/100  Grade: {run['overall_grade']}",
        "",
        "| Company | Use Case | Data | LB | Total | Note |",
        "|---------|----------|------|----|-------|------|",
    ]
    for r in run["tests"]:
        ag      = r.get("ai_grading", {})
        ds      = ag.get("data_score", "n/a")
        ls      = ag.get("liveboard_score", "n/a")
        t_str   = f"{r['total_score']}/{r['grade']}"
        company = r.get("company", r["name"])
        parts   = [p for p in [r.get("vertical",""), r.get("line",""), r.get("function","")]
                   if p and p != "* CUSTOM *"]
        uc      = (" / ".join(parts) if parts else "Custom")[:W]
        if r.get("error"):            note = "❌ network err"
        elif r.get("late_complete"):  note = "⚠️ slow (complete)"
        elif r.get("timed_out"):      note = "⏰ timeout"
        elif r["grade"] in ("A","B"): note = "🏆 great"
        elif r["grade"] == "C":       note = "✅ solid"
        else:                         note = ""
        ctx     = r.get("run_context") or {}
        lb_guid = ctx.get("liveboard_guid","")
        lb_base = (ctx.get("ts_base_url","") or "").rstrip("/")
        lb_link = f"[lb]({lb_base}/#/pinboard/{lb_guid})" if lb_guid and lb_base else "—"
        lines.append(f"| {company} | {uc} | {ds} | {ls} | {t_str} {lb_link} | {note} |")

    # Issues and errors
    issues = []
    for r in run["tests"]:
        if r.get("timed_out") and not r.get("late_complete"):
            last = r.get("diagnostics",{}).get("last_event","unknown")
            issues.append(f"- **{r.get('company',r['name'])}**: TIMEOUT — last event: {last}")
        for err in r.get("ai_grading",{}).get("grading_errors",[]):
            if "skip" not in err.lower():
                issues.append(f"- **{r.get('company',r['name'])}**: {err}")
    if issues:
        lines += ["", "## Issues", ""] + issues

    # Top data weaknesses
    weaknesses = []
    for r in run["tests"]:
        ag = r.get("ai_grading",{})
        ww = ag.get("data_weaknesses",[])
        if ww:
            weaknesses.append(f"**{r.get('company',r['name'])}** (data={ag.get('data_score','?')}/100):")
            for w in ww[:2]:
                weaknesses.append(f"  - {w[:120]}")
    if weaknesses:
        lines += ["", "## Data Quality Weaknesses", ""] + weaknesses

    lines += ["", "---", f"*JSON: {json_path.name}*"]

    md_text = "\n".join(lines) + "\n"
    md_path = json_path.with_suffix(".md")
    md_path.write_text(md_text)

    latest_name = f"latest_{env_name}_summary.md" if env_name else "latest_summary.md"
    latest = RESULTS_DIR / latest_name
    latest.write_text(md_text)

    print(f"📋 Summary: {md_path}")
    print(f"📋 Latest:  {latest}")
    return md_path


def print_handoff_block(run: dict, json_path: Path, md_path: Path, env_name: str = ""):
    latest_name = f"latest_{env_name}_summary.md" if env_name else "latest_summary.md"
    latest_path = RESULTS_DIR / latest_name
    timestamp = run.get("timestamp", "")
    run_id = run.get("run_id", "")
    target = run.get("target_url", "")
    avg = run.get("avg_score", "")
    grade = run.get("overall_grade", "")

    print("\n📌 Agent handoff")
    print(f"  Run ID:    {run_id}")
    print(f"  Timestamp: {timestamp}")
    print(f"  Target:    {target}")
    print(f"  Results:   {json_path}")
    print(f"  Summary:   {md_path}")
    print(f"  Latest:    {latest_path}")
    print(
        "  Paste this: "
        f"DemoPrep quality run {run_id} ({timestamp}) "
        f"avg={avg}/{grade} target={target} "
        f"results={json_path} summary={md_path}"
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def run_quality_suite(max_tests: int = None, env_name: str = "", suite_override: list[dict] = None):
    if not TEST_USER or not TEST_PASSWORD:
        raise RuntimeError("TEST_USER and TEST_PASSWORD must be set in .env")

    config = load_config()
    suite  = suite_override or build_test_suite(config)
    if max_tests:
        suite = suite[:max_tests]
    run_id  = str(uuid.uuid4())[:8]

    print(f"\n{'='*62}")
    print(f"  DemoPrep Quality Run — {datetime.now().strftime('%Y-%m-%d %H:%M')}")
    print(f"  Run ID: {run_id}")
    print(f"  Target: {BASE_URL}")
    print(f"  {len(suite)} tests  |  "
          f"{sum(1 for t in suite if t['type']=='fixed')} fixed  "
          f"{sum(1 for t in suite if t['type']=='random')} random  "
          f"{sum(1 for t in suite if t['type']=='ai_generated')} AI-generated  "
          f"{sum(1 for t in suite if t['type']=='custom')} custom")
    print(f"  Scoring: stages(25) + data(50) + liveboard(25) = 100 pts")
    print(f"{'='*62}")
    for i, tc in enumerate(suite, 1):
        label = {"fixed": "🔒", "random": "🎲", "ai_generated": "🤖", "custom": "✏️"}[tc["type"]]
        print(f"  [{i}] {label} {tc['name']}")

    results = []

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=not DRY_RUN)
        ctx     = browser.new_context(viewport={"width": 1280, "height": 900})

        print(f"\n🔐 Logging in as {TEST_USER}...")
        page = ctx.new_page()
        page.goto(BASE_URL, timeout=90000)
        page.wait_for_selector('input[type=password], button[role=tab]', timeout=90000)
        if page.locator('input[type=password]').is_visible(timeout=2000):
            _do_login(page)
        print("✅ Logged in\n")

        for i, test_case in enumerate(suite, 1):
            label = {"fixed": "🔒", "random": "🎲", "ai_generated": "🤖", "custom": "✏️"}[test_case["type"]]
            print(f"{'─'*62}")
            print(f"[{i}/{len(suite)}] {label} {test_case['name']}")

            result = run_single_test(page, test_case, config)
            results.append(result)

            ag = result["ai_grading"]
            sg = result["stage_grading"]
            print(f"  Stages: {sg.get('stage_total', 0)}/25")
            if ag.get("data_score") is not None:
                print(f"  Data:   {ag['data_score']}/100 → {ag['data_points']} pts")
            if ag.get("liveboard_score") is not None:
                print(f"  Board:  {ag['liveboard_score']}/100 → {ag['liveboard_points']} pts")
            for err in ag.get("grading_errors", []):
                print(f"  ⚠️  {err}")
            ctx_r  = result.get("run_context", {})
            ts_url = (ctx_r.get("ts_base_url") or "").rstrip("/")
            m_guid = ctx_r.get("model_guid", "")
            l_guid = ctx_r.get("liveboard_guid", "")
            if m_guid and ts_url:
                print(f"  Model:     {ts_url}/#/data/tables/{m_guid}")
            if l_guid and ts_url:
                viz_n = result.get("liveboard_viz_count")
                viz_note = f"  ({viz_n} vizzes)" if viz_n is not None else ""
                print(f"  Liveboard: {ts_url}/#/pinboard/{l_guid}{viz_note}")
            if result.get("late_complete"):
                timeout_tag = "  ⚠️  SLOW (completed late)"
            elif result.get("timed_out"):
                timeout_tag = "  ⏰ TIMEOUT"
            else:
                timeout_tag = ""
            print(f"  TOTAL: {result['total_score']}/100  Grade: {result['grade']}"
                  f"  ({result['duration_seconds']}s)"
                  f"{timeout_tag}"
                  f"{'  ❌ ERROR' if result['error'] else ''}")

        ctx.close()
        browser.close()

    avg   = round(sum(r["total_score"] for r in results) / len(results), 1) if results else 0
    grade = compute_grade(avg, config)

    run = {
        "run_id": run_id, "timestamp": datetime.now().isoformat(),
        "target_url": BASE_URL,
        "avg_score": avg, "overall_grade": grade,
        "test_count": len(results), "tests": results,
    }
    path = save_results(run)
    md_path = save_summary_md(run, path, env_name=env_name)
    save_to_postgres(run, env_name=env_name)
    print_handoff_block(run, path, md_path, env_name=env_name)

    # --- Summary table ---
    try:
        W_CO, W_UC, W_DA, W_LB, W_TO, W_NO = 16, 22, 6, 6, 9, 14
        B = "│"

        def _link(url, label):
            return f"\033]8;;{url}\033\\{label}\033]8;;\033\\"

        def _row(co, uc, da, lb, to, lk, no):
            return (f"{B} {co:<{W_CO}} {B} {uc:<{W_UC}} {B} {da:>{W_DA}} {B}"
                    f" {lb:>{W_LB}} {B} {to:>{W_TO}} {B} {lk} {B} {no:<{W_NO}} {B}")

        def _div(l, m, r):
            s = "─"
            return (f"{l}{s*(W_CO+2)}{m}{s*(W_UC+2)}{m}{s*(W_DA+2)}{m}"
                    f"{s*(W_LB+2)}{m}{s*(W_TO+2)}{m}{s*6}{m}{s*(W_NO+2)}{r}")

        print(f"\n{'='*62}")
        print(f"  COMPLETE  —  Avg: {avg}/100  Grade: {grade}  |  {BASE_URL}")
        print()
        print(_div("┌", "┬", "┐"))
        print(_row("Company", "Use Case", " Data", "  LB", "  Total", " Link", "Note"))
        print(_div("├", "┼", "┤"))

        for r in results:
            ag    = r.get("ai_grading", {})
            ds    = ag.get("data_score")
            ls    = ag.get("liveboard_score")
            d_str = str(ds) if ds is not None else "n/a"
            l_str = str(ls) if ls is not None else "n/a"
            t_str = f"{r['total_score']}/{r['grade']}"
            ctx   = r.get("run_context") or {}
            lb_url, lb_base = ctx.get("liveboard_guid",""), ctx.get("ts_base_url","")
            lk = _link(f"{lb_base}/#/pinboard/{lb_url}", " 📋 ") if lb_url and lb_base else " —  "
            parts = [p for p in [r.get("vertical",""), r.get("line",""), r.get("function","")] if p and p != "* CUSTOM *"]
            uc = (" / ".join(parts) if parts else r.get("context","")[:W_UC] or "Custom")[:W_UC]
            errs = ag.get("grading_errors", [])
            if r.get("error"):                                        note = "❌ network err"
            elif r.get("late_complete"):                              note = "⚠️ slow"
            elif r.get("timed_out"):                                  note = "⏰ timeout"
            elif any("auth failed" in (e or "").lower() for e in errs): note = "❌ auth fail"
            elif any("parse" in (e or "").lower() for e in errs):    note = "❌ parse fail"
            elif ds == 0 and ls is not None:                          note = "⚠️ data fail"
            elif r["grade"] in ("A","B"):                             note = "🏆 great"
            elif r["grade"] == "C":                                   note = "✅ solid"
            else:                                                     note = ""
            company = r.get("company", r["name"])[:W_CO]
            print(_row(company, uc, d_str, l_str, t_str, lk, note))

        print(_div("└", "┴", "┘"))

        # Liveboard links — plain text for easy copy/click
        print("\n  Liveboards:")
        for i, r in enumerate(results, 1):
            ctx     = r.get("run_context") or {}
            lb_url  = ctx.get("liveboard_guid", "")
            lb_base = ctx.get("ts_base_url", "")
            company = r.get("company", r["name"])
            url     = f"{lb_base}/#/pinboard/{lb_url}" if lb_url and lb_base else "— not created"
            print(f"  {i}. {company:<30}  {url}")

        print(f"\n{'='*62}\n")

    except Exception as _table_err:
        print(f"\n⚠️  Summary table failed: {_table_err}")
        print(f"{'='*62}")
        print(f"  COMPLETE  —  Avg: {avg}/100  Grade: {grade}")
        for r in results:
            ag = r.get("ai_grading", {})
            ds = ag.get("data_score")
            ls = ag.get("liveboard_score")
            print(f"  {r['name']}: {r['total_score']}/100 {r['grade']}  data={ds}  lb={ls}")
        print(f"{'='*62}\n")
    return run


def test_quality_run():
    run = run_quality_suite()
    assert run["overall_grade"] != "F", f"Quality run averaged {run['avg_score']}% — too many failures."


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--count", type=int, default=0,
                        help="Run only N tests (default: all 8)")
    parser.add_argument("--datadog-saas-sales", action="store_true",
                        help="Run one targeted Datadog Software as a Service / Sales test")
    parser.add_argument("--dataset-first-five", action="store_true",
                        help="Run five curated dataset-first smoke tests: four defined flows plus one custom")
    parser.add_argument("--dry-run", action="store_true",
                        help="Fill form but do not click GO — browser opens visibly for inspection")
    parser.add_argument("--url", type=str, default="",
                        help="Override TEST_TARGET_URL (e.g. --url https://thoughtspot-dp-demoprep.hf.space)")
    parser.add_argument("--env-name", type=str, default="",
                        help="Tag for summary filename: 'test' → latest_test_summary.md, 'prod' → latest_prod_summary.md")
    parser.add_argument("--ts-environment", type=str, default="",
                        help="Override the TS Environment dropdown value for this run")
    parser.add_argument("--test-user", type=str, default="",
                        help="Override TEST_USER for this run only")
    parser.add_argument("--test-password", type=str, default="",
                        help="Override TEST_PASSWORD for this run only; prefer --test-password-env")
    parser.add_argument("--test-password-env", type=str, default="",
                        help="Environment variable containing the password for --test-user")
    args = parser.parse_args()
    if args.test_user:
        TEST_USER = args.test_user
        if not args.test_password and not args.test_password_env:
            raise SystemExit(
                "--test-user requires --test-password or --test-password-env; "
                "otherwise the runner would use the default TEST_PASSWORD for a different user."
            )
    if args.test_password_env:
        TEST_PASSWORD = os.getenv(args.test_password_env, "")
    elif args.test_password:
        TEST_PASSWORD = args.test_password
    if args.dry_run:
        DRY_RUN = True
    if args.url:
        BASE_URL = args.url
    if args.ts_environment:
        RUN_SETTINGS["ts_environment"] = args.ts_environment
    if not BASE_URL:
        raise ValueError("No target URL — set TEST_TARGET_URL in .env or pass --url <url>")
    suite_override = None
    if args.datadog_saas_sales:
        suite_override = [{
            "name": "datadog_saas_sales_dataset_first",
            "type": "fixed",
            "company": "Datadog",
            "company_url": "datadog.com",
            "vertical": "Technology",
            "line": "Software as a Service",
            "function": "Sales",
        }]
    if args.dataset_first_five:
        suite_override = [
            {
                "name": "ey_professional_services_dataset_first",
                "type": "custom",
                "company": "EY",
                "company_url": "ey.com",
                "vertical": "* CUSTOM *",
                "line": "",
                "function": "Custom",
                "context": (
                    "Create a professional services analytics demo for EY. "
                    "Focus on client engagements, service lines, industries, consultants, billable hours, "
                    "utilization, realization, pipeline, project margin, delivery risk, and client satisfaction. "
                    "Use consulting and assurance terminology only. Do not create sports, venue, ticketing, "
                    "fan engagement, or entertainment analytics."
                ),
            },
            {
                "name": "datadog_saas_sales_dataset_first",
                "type": "fixed",
                "company": "Datadog",
                "company_url": "datadog.com",
                "vertical": "Technology",
                "line": "Software as a Service",
                "function": "Sales",
            },
            {
                "name": "nike_retail_sales_dataset_first",
                "type": "fixed",
                "company": "Nike",
                "company_url": "nike.com",
                "vertical": "Retail & Consumer Goods",
                "line": "Fashion/Apparel",
                "function": "Sales",
            },
            {
                "name": "delta_airline_operations_dataset_first",
                "type": "fixed",
                "company": "Delta",
                "company_url": "delta.com",
                "vertical": "Transportation & Logistics",
                "line": "Air Transport",
                "function": "Sales",
            },
            {
                "name": "wells_fargo_banking_marketing_dataset_first",
                "type": "fixed",
                "company": "Wells Fargo",
                "company_url": "wellsfargo.com",
                "vertical": "Financial Services",
                "line": "Banking",
                "function": "Marketing",
            },
            {
                "name": "starbucks_custom_store_operations_dataset_first",
                "type": "custom",
                "company": "Starbucks",
                "company_url": "starbucks.com",
                "vertical": "* CUSTOM *",
                "line": "",
                "function": "Custom",
                "context": (
                    "Create a store operations analytics demo for Starbucks. "
                    "Focus on store-day and daypart performance, transactions, net sales, "
                    "labor hours, order channel, product category, wait times, and customer satisfaction. "
                    "Keep values realistic for coffee retail: transactions must reconcile to sales, "
                    "refunds must stay below transactions, wait times should be measured in minutes, "
                    "and channels should be in-store, drive-thru, mobile order, or delivery."
                ),
            },
        ]
    run_quality_suite(
        max_tests=args.count or (1 if args.dry_run else None),
        env_name=args.env_name,
        suite_override=suite_override,
    )