"""
Databricks job script: run pytest for insurance-optimise, insurance-deploy, insurance-distill
Uploads each repo's source + tests to DBFS, installs in a temp venv per repo, runs pytest.
"""

import os
import sys
import subprocess
import tempfile
import shutil
import base64
import json
from pathlib import Path

# ── Databricks SDK bootstrap ──────────────────────────────────────────────────
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import (
    NotebookTask, Task, JobCluster, ClusterSpec,
    PythonWheelTask, SparkPythonTask,
)
from databricks.sdk.service import jobs as jobs_svc
from databricks.sdk.service.compute import (
    AutoScale, ClusterLogConf, DbfsStorageInfo,
    RuntimeEngine,
)

# 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()

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

# ── Upload repos to workspace ─────────────────────────────────────────────────
REPOS = {
    "insurance-optimise": "/home/ralph/burning-cost/repos/insurance-optimise",
    "insurance-deploy":   "/home/ralph/burning-cost/repos/insurance-deploy",
    "insurance-distill":  "/home/ralph/burning-cost/repos/insurance-distill",
}

WS_BASE = "/Workspace/Shared/test-sweep-batch2"

def upload_dir(local_dir: str, ws_path: str, w: WorkspaceClient):
    """Recursively upload a local directory to Databricks workspace as files."""
    local = Path(local_dir)
    ws = ws_path.rstrip("/")

    # Gather all files (skip .venv, __pycache__, *.pyc, .git, dist, *.egg-info)
    skip_dirs = {".venv", "__pycache__", ".git", "dist", ".eggs"}
    skip_ext  = {".pyc", ".pyo"}

    files_to_upload = []
    for path in local.rglob("*"):
        if path.is_dir():
            continue
        # Check if any parent is in skip_dirs
        parts = set(path.relative_to(local).parts[:-1])
        if parts & skip_dirs:
            continue
        if path.suffix in skip_ext:
            continue
        files_to_upload.append(path)

    print(f"  Uploading {len(files_to_upload)} files to {ws}...")

    for path in files_to_upload:
        rel = path.relative_to(local)
        target_ws = f"{ws}/{rel}".replace("\\", "/")

        # Create parent directory
        parent = "/".join(target_ws.split("/")[:-1])
        try:
            w.workspace.mkdirs(parent)
        except Exception:
            pass  # may already exist

        # Upload file
        content = path.read_bytes()
        encoded = base64.b64encode(content).decode()
        try:
            w.workspace.import_(
                path=target_ws,
                content=encoded,
                format="AUTO",
                overwrite=True,
            )
        except Exception as e:
            # Try as RAW for binary files
            try:
                w.workspace.import_(
                    path=target_ws,
                    content=encoded,
                    format="RAW",
                    overwrite=True,
                )
            except Exception as e2:
                print(f"    WARN: could not upload {rel}: {e2}")

print("Uploading repos to Databricks workspace...")
for name, local_path in REPOS.items():
    print(f"\n[{name}]")
    upload_dir(local_path, f"{WS_BASE}/{name}", w)

print("\nAll uploads complete.")

# ── Build a notebook that runs all three test suites ─────────────────────────
NOTEBOOK_CONTENT = '''# Databricks notebook source
# MAGIC %md
# MAGIC # Test Sweep Batch 2: optimise, deploy, distill

# COMMAND ----------

import subprocess, sys, os, json
from pathlib import Path

results = {}

REPOS = [
    ("insurance-optimise", "/Workspace/Shared/test-sweep-batch2/insurance-optimise"),
    ("insurance-deploy",   "/Workspace/Shared/test-sweep-batch2/insurance-deploy"),
    ("insurance-distill",  "/Workspace/Shared/test-sweep-batch2/insurance-distill"),
]

# COMMAND ----------

for (name, repo_path) in REPOS:
    print(f"\\n{'='*60}")
    print(f"REPO: {name}")
    print(f"{'='*60}")

    venv_dir = f"/tmp/venv-{name}"

    # Create venv
    r = subprocess.run(
        [sys.executable, "-m", "venv", venv_dir],
        capture_output=True, text=True
    )
    if r.returncode != 0:
        print(f"VENV CREATION FAILED: {r.stderr}")
        results[name] = {"install": False, "error": r.stderr}
        continue

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

    # Install package with dev deps
    print(f"Installing {name}...")
    install_r = subprocess.run(
        [pip, "install", "-e", f"{repo_path}[dev]", "--quiet"],
        capture_output=True, text=True, cwd=repo_path
    )
    install_tail = "\\n".join((install_r.stdout + install_r.stderr).strip().split("\\n")[-5:])
    print(f"Install stdout/stderr (last 5 lines):\\n{install_tail}")

    if install_r.returncode != 0:
        print(f"INSTALL FAILED (rc={install_r.returncode})")
        results[name] = {"install": False, "error": install_tail}
        # cleanup
        import shutil; shutil.rmtree(venv_dir, ignore_errors=True)
        continue

    print(f"Install succeeded.")

    # Run pytest
    print(f"Running pytest for {name}...")
    pytest_r = subprocess.run(
        [python, "-m", "pytest", "--tb=short", "-v", "tests/"],
        capture_output=True, text=True, cwd=repo_path, timeout=300
    )
    output = pytest_r.stdout + pytest_r.stderr
    tail = "\\n".join(output.strip().split("\\n")[-40:])
    print(f"--- pytest output (last 40 lines) ---")
    print(tail)

    results[name] = {
        "install": True,
        "returncode": pytest_r.returncode,
        "output_tail": tail,
    }

    # Clean up venv
    import shutil; shutil.rmtree(venv_dir, ignore_errors=True)

# COMMAND ----------

print("\\n\\n" + "="*70)
print("SUMMARY")
print("="*70)
for name, res in results.items():
    print(f"\\n[{name}]")
    if not res.get("install"):
        print(f"  Install: FAILED")
        print(f"  Error: {res.get('error','')[:200]}")
    else:
        print(f"  Install: OK")
        rc = res.get("returncode", -1)
        print(f"  Pytest return code: {rc} ({'PASSED' if rc == 0 else 'FAILURES/ERRORS'})")
        # Extract summary line
        for line in res["output_tail"].split("\\n"):
            if "passed" in line or "failed" in line or "error" in line:
                print(f"  Summary line: {line.strip()}")

# COMMAND ----------
import json
dbutils.notebook.exit(json.dumps(results))
'''

# Upload the notebook
print("\nUploading test runner notebook...")
try:
    w.workspace.mkdirs(WS_BASE)
except Exception:
    pass

notebook_path = f"{WS_BASE}/run_tests"
w.workspace.import_(
    path=notebook_path,
    content=base64.b64encode(NOTEBOOK_CONTENT.encode()).decode(),
    format="SOURCE",
    language="PYTHON",
    overwrite=True,
)
print(f"Notebook uploaded to: {notebook_path}")

# ── Submit notebook job ───────────────────────────────────────────────────────
print("\nSubmitting job run...")

run = w.jobs.submit(
    run_name="test-sweep-batch2",
    tasks=[
        jobs_svc.SubmitTask(
            task_key="run-tests",
            notebook_task=jobs_svc.NotebookTask(
                notebook_path=notebook_path,
                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,
            ),
        )
    ],
).result()

print(f"Run submitted. Run ID: {run.run_id}")
print(f"Waiting for completion...")

# Poll until done
import time
while True:
    run_state = w.jobs.get_run(run_id=run.run_id)
    life_cycle = run_state.state.life_cycle_state
    print(f"  State: {life_cycle}")
    if life_cycle in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
        break
    time.sleep(15)

result_state = run_state.state.result_state
print(f"\nFinal result state: {result_state}")

# Fetch output
for task_run in run_state.tasks or []:
    output = w.jobs.get_run_output(run_id=task_run.run_id)
    if output.notebook_output:
        print("\n=== NOTEBOOK OUTPUT ===")
        print(output.notebook_output.result)
    if output.error:
        print(f"\n=== ERROR ===\n{output.error}")
    if output.error_trace:
        print(f"\n=== TRACE ===\n{output.error_trace[:3000]}")

print("\nDone.")
