"""
Test sweep batch 2: insurance-optimise, insurance-deploy, insurance-distill
Runs each test suite on Databricks serverless via sequential notebook jobs.
"""

import os, base64, time, json
from pathlib import Path

# Load creds
env_path = Path.home() / ".config" / "burning-cost" / "databricks.env"
for line in env_path.read_text().splitlines():
    line = line.strip()
    if line and not line.startswith("#") and "=" in line:
        k, v = line.split("=", 1)
        os.environ[k.strip()] = v.strip()

from databricks.sdk import WorkspaceClient
from databricks.sdk.service import jobs, compute
from databricks.sdk.service.workspace import ImportFormat, Language

w = WorkspaceClient()
print(f"Connected: {os.environ['DATABRICKS_HOST']}")

# ── Each repo: git URL, pip package name with [dev], package source path ──────
REPOS = [
    {
        "name":       "insurance-optimise",
        "github_url": "https://github.com/burning-cost/insurance-optimise",
        "install":    ".[dev]",          # installed from cloned dir
        "tests_dir":  "tests/",
    },
    {
        "name":       "insurance-deploy",
        "github_url": "https://github.com/burning-cost/insurance-deploy",
        "install":    ".[dev]",
        "tests_dir":  "tests/",
    },
    {
        "name":       "insurance-distill",
        "github_url": "https://github.com/burning-cost/insurance-distill",
        "install":    ".[dev]",
        "tests_dir":  "tests/",
    },
]

ENV_KEY = "sweep-batch2-env"

def make_notebook(repo: dict) -> str:
    """Generate notebook source for a single repo."""
    return f'''# Databricks notebook source
import subprocess, sys, os, shutil, json

name       = {json.dumps(repo["name"])}
github_url = {json.dumps(repo["github_url"])}
tests_dir  = {json.dumps(repo["tests_dir"])}

work_dir = f"/tmp/work-{{name}}"
shutil.rmtree(work_dir, ignore_errors=True)

# Clone
print(f"Cloning {{github_url}}...")
r_clone = subprocess.run(
    ["git", "clone", "--depth=1", github_url, work_dir],
    capture_output=True, text=True, timeout=120
)
if r_clone.returncode != 0:
    dbutils.notebook.exit(json.dumps({{
        "status": "CLONE_FAILED",
        "output": r_clone.stderr[-2000:]
    }}))

print("Clone OK")

# Install deps using venv
venv_dir = f"/tmp/venv-{{name}}"
shutil.rmtree(venv_dir, ignore_errors=True)

r_venv = subprocess.run([sys.executable, "-m", "venv", venv_dir],
    capture_output=True, text=True, timeout=60)
if r_venv.returncode != 0:
    dbutils.notebook.exit(json.dumps({{
        "status": "VENV_FAILED",
        "output": r_venv.stderr[-1000:]
    }}))

pip    = f"{{venv_dir}}/bin/pip"
python = f"{{venv_dir}}/bin/python"

subprocess.run([pip, "install", "--upgrade", "pip", "-q"], capture_output=True)

print(f"Installing {{name}}[dev]...")
r_install = subprocess.run(
    [pip, "install", "-e", f"{{work_dir}}[dev]"],
    capture_output=True, text=True, timeout=300, cwd=work_dir
)
install_out = (r_install.stdout + r_install.stderr).strip()
install_tail = "\\n".join(install_out.split("\\n")[-10:])
print(f"Install output (last 10 lines):\\n{{install_tail}}")

if r_install.returncode != 0:
    dbutils.notebook.exit(json.dumps({{
        "status": "INSTALL_FAILED",
        "output": install_tail
    }}))

print("Install OK")

# Run pytest
print("Running pytest...")
log_file = f"/tmp/pytest-{{name}}.txt"
with open(log_file, "w") as f:
    r_pytest = subprocess.run(
        [python, "-m", "pytest", tests_dir, "-v", "--tb=short", "-p", "no:cacheprovider"],
        stdout=f, stderr=subprocess.STDOUT,
        timeout=360, cwd=work_dir
    )

with open(log_file) as f:
    output = f.read()

all_lines = output.strip().split("\\n")
tail = "\\n".join(all_lines[-60:])
print(f"--- pytest tail ---\\n{{tail}}")

# Parse summary line
summary = ""
for line in reversed(all_lines):
    if "passed" in line or "failed" in line or "error" in line:
        summary = line.strip()
        break

status = "PASSED" if r_pytest.returncode == 0 else "FAILED"
dbutils.notebook.exit(json.dumps({{
    "status": status,
    "returncode": r_pytest.returncode,
    "summary": summary,
    "tail": tail,
}}))
'''

# ── Submit one job per repo ───────────────────────────────────────────────────
WS_BASE = "/Workspace/Shared/test-sweep-batch2"
try:
    w.workspace.mkdirs(WS_BASE)
except Exception:
    pass

all_results = {}

for repo in REPOS:
    name = repo["name"]
    print(f"\n{'='*60}")
    print(f"Submitting: {name}")
    print(f"{'='*60}")

    nb_path = f"{WS_BASE}/test_{name.replace('-', '_')}"
    nb_content = make_notebook(repo)
    encoded = base64.b64encode(nb_content.encode()).decode()

    # Upload notebook
    try:
        w.workspace.delete(nb_path)
    except Exception:
        pass
    w.workspace.import_(
        path=nb_path,
        content=encoded,
        overwrite=True,
        format=ImportFormat.SOURCE,
        language=Language.PYTHON,
    )
    print(f"Notebook uploaded: {nb_path}")

    # Submit serverless job
    run_waiter = w.jobs.submit(
        run_name=f"test-sweep-{name}",
        environments=[
            jobs.JobEnvironment(
                environment_key=ENV_KEY,
                spec=compute.Environment(client="2"),
            )
        ],
        tasks=[
            jobs.SubmitTask(
                task_key="run-tests",
                environment_key=ENV_KEY,
                notebook_task=jobs.NotebookTask(
                    notebook_path=nb_path,
                ),
            )
        ],
    )
    run_id = run_waiter.run_id
    print(f"Run ID: {run_id}")

    # Poll
    while True:
        state = w.jobs.get_run(run_id=run_id)
        life = state.state.life_cycle_state
        life_str = life.value if hasattr(life, "value") else str(life)
        result_str = ""
        if state.state.result_state:
            result_str = state.state.result_state.value if hasattr(state.state.result_state, "value") else str(state.state.result_state)
        print(f"  [{time.strftime('%H:%M:%S')}] {life_str}  {result_str}")
        if any(s in life_str for s in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR")):
            break
        time.sleep(15)

    # Fetch output
    nb_result = None
    error_msg = ""
    error_trace = ""
    if state.tasks:
        for task in state.tasks:
            try:
                out = w.jobs.get_run_output(run_id=task.run_id)
                if out.notebook_output and out.notebook_output.result:
                    nb_result = out.notebook_output.result
                if out.error:
                    error_msg = out.error
                if out.error_trace:
                    error_trace = out.error_trace[-3000:]
            except Exception as e:
                print(f"  Could not fetch output: {e}")

    # Parse notebook exit value
    repo_result = {
        "job_state": result_str,
        "error": error_msg,
        "error_trace": error_trace,
    }
    if nb_result:
        try:
            parsed = json.loads(nb_result)
            repo_result.update(parsed)
        except Exception:
            repo_result["raw_output"] = nb_result

    all_results[name] = repo_result

    print(f"\n[{name}] Result: {repo_result.get('status','?')} | {repo_result.get('summary','')}")
    if repo_result.get("tail"):
        print("Tail:")
        for line in repo_result["tail"].split("\n")[-30:]:
            print(f"  {line}")

# ── Final summary ─────────────────────────────────────────────────────────────
print("\n\n" + "="*70)
print("FINAL SUMMARY")
print("="*70)
for name, res in all_results.items():
    status = res.get("status", "UNKNOWN")
    summary = res.get("summary", "")
    rc = res.get("returncode", "?")
    print(f"\n[{name}]")
    print(f"  Status:    {status}")
    print(f"  Returncode:{rc}")
    print(f"  Summary:   {summary}")
    if res.get("error"):
        print(f"  Error:     {res['error'][:300]}")

print("\n\nDone.")

# Save full JSON for KB
print("\n\n=== JSON FOR KB ===")
print(json.dumps({k: {
    "status": v.get("status"),
    "returncode": v.get("returncode"),
    "summary": v.get("summary",""),
    "tail": v.get("tail","")[-2000:],
    "error": v.get("error",""),
} for k, v in all_results.items()}, indent=2))
