"""Submit and poll the test-sweep-batch2 notebook job on Databricks."""

import os, time, json

# Load creds
env_path = os.path.expanduser("~/.config/burning-cost/databricks.env")
with open(env_path) as f:
    for line in f:
        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 as jobs_svc

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

run_resp = w.jobs.submit(
    run_name="test-sweep-batch2",
    tasks=[
        jobs_svc.SubmitTask(
            task_key="run-tests",
            notebook_task=jobs_svc.NotebookTask(
                notebook_path="/Workspace/Shared/nb-test-sweep-batch2",
                base_parameters={},
            ),
            new_cluster=jobs_svc.ClusterSpec(
                spark_version="15.4.x-scala2.12",
                node_type_id="m5d.large",
                num_workers=0,
                spark_conf={"spark.master": "local[*]"},
                data_security_mode=jobs_svc.DataSecurityMode.SINGLE_USER,
            ),
        )
    ],
)

run_id = run_resp.run_id
print(f"Submitted. Run ID: {run_id}")
print(f"Tracking URL: {os.environ['DATABRICKS_HOST']}#job/runs/{run_id}")

# Poll
while True:
    run_state = w.jobs.get_run(run_id=run_id)
    lc = run_state.state.life_cycle_state.value if run_state.state.life_cycle_state else "?"
    rs = run_state.state.result_state.value if run_state.state.result_state else ""
    print(f"  [{time.strftime('%H:%M:%S')}] lifecycle={lc}  result={rs}")
    if lc in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
        break
    time.sleep(20)

print(f"\nFinal: lifecycle={lc}  result={rs}")

# Get output
for task_run in (run_state.tasks or []):
    out = w.jobs.get_run_output(run_id=task_run.run_id)
    if out.notebook_output and out.notebook_output.result:
        print("\n=== NOTEBOOK EXIT VALUE ===")
        raw = out.notebook_output.result
        print(raw)
        try:
            data = json.loads(raw)
            print("\n=== PARSED SUMMARY ===")
            for repo, res in data.items():
                print(f"\n[{repo}]")
                if not res.get("install"):
                    print(f"  Install: FAILED  error={res.get('error','')[:200]}")
                else:
                    rc = res.get("returncode", -1)
                    print(f"  Install:  OK")
                    print(f"  Pytest:   {'PASSED' if rc==0 else 'FAILED/ERROR'}  rc={rc}")
                    print(f"  Summary:  {res.get('summary','')}")
                    print(f"  Output tail:\n")
                    for line in res.get("output_tail","").split("\n")[-30:]:
                        print(f"    {line}")
        except Exception as e:
            print(f"(could not parse JSON: {e})")
    if out.error:
        print(f"\n=== TASK ERROR ===\n{out.error}")
    if out.error_trace:
        print(f"\n=== ERROR TRACE ===\n{out.error_trace[:4000]}")
