"""Submit insurance-copula pytest suite to Databricks serverless compute.

Uses the REST API directly — no cluster spec required. Embeds all source
and test files in the notebook via JSON so there is no dependency on
workspace file structure.

Usage:
    python run_databricks_tests.py
"""

from __future__ import annotations

import base64
import json
import os
import sys
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path

# ---------------------------------------------------------------------------
# Load credentials
# ---------------------------------------------------------------------------

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

api_base = os.environ["DATABRICKS_HOST"].rstrip("/")
token = os.environ["DATABRICKS_TOKEN"]
headers = {
    "Authorization": f"Bearer {token}",
    "Content-Type": "application/json",
}

RUN_ID = uuid.uuid4().hex[:8]
WORKSPACE_FOLDER = "/Workspace/insurance-copula"
NOTEBOOK_PATH = f"{WORKSPACE_FOLDER}/run_pytest_{RUN_ID}"

# ---------------------------------------------------------------------------
# Read source and test files
# ---------------------------------------------------------------------------

BASE = Path(__file__).parent
SKIP = {".venv", "__pycache__", ".git", ".pytest_cache", "notebooks",
        "run_tests_databricks.py", "run_databricks_tests.py", "uv.lock"}


def read_py_files(directory: Path, prefix: str = "") -> dict[str, str]:
    """Recursively read .py files, keyed by relative path."""
    result: dict[str, str] = {}
    for item in sorted(directory.iterdir()):
        if item.name in SKIP or item.name.startswith("."):
            continue
        if item.is_dir():
            result.update(read_py_files(item, prefix=f"{prefix}{item.name}/"))
        elif item.suffix == ".py":
            result[f"{prefix}{item.name}"] = item.read_text()
    return result


src_files = read_py_files(BASE / "src" / "insurance_copula")
test_files = read_py_files(BASE / "tests")
pyproject_content = (BASE / "pyproject.toml").read_text()

files_json = json.dumps({**src_files, **test_files})
pyproject_json = json.dumps(pyproject_content)
src_file_names_json = json.dumps(list(src_files.keys()))

# ---------------------------------------------------------------------------
# Build notebook source
#
# Pattern that works on Databricks serverless:
# 1. %pip for standard deps (scipy, numpy, etc) — restarts kernel
# 2. After restart: subprocess pip --only-binary for pyvinecopulib
# 3. Write files from embedded JSON, install package, run pytest
# ---------------------------------------------------------------------------

NOTEBOOK_SOURCE = f"""# Databricks notebook source
# MAGIC %pip install scipy numpy pandas matplotlib pytest pytest-cov hatchling --quiet

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

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

# Install pyvinecopulib via binary wheel — must happen after %pip kernel restart
r = subprocess.run(
    [sys.executable, "-m", "pip", "install",
     "--only-binary=:all:", "--quiet", "pyvinecopulib>=0.7"],
    capture_output=True, text=True,
)
if r.returncode == 0:
    print("pyvinecopulib installed OK")
else:
    print("WARNING: pyvinecopulib install failed -- vine tests will be skipped")
    print(r.stderr[-300:])

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

pkg_id = uuid.uuid4().hex[:8]
pkg_dir = f"/tmp/insurance_copula_{{pkg_id}}"
src_dir = f"{{pkg_dir}}/src/insurance_copula"
tests_dir = f"{{pkg_dir}}/tests"
os.makedirs(src_dir, exist_ok=True)
os.makedirs(tests_dir, exist_ok=True)

FILES_JSON = {files_json!r}
PYPROJECT_JSON = {pyproject_json!r}
SRC_FILE_NAMES = {src_file_names_json!r}

files_map = json.loads(FILES_JSON)
pyproject_src = json.loads(PYPROJECT_JSON)
src_names = set(json.loads(SRC_FILE_NAMES))

for name, content in files_map.items():
    if name in src_names:
        path = f"{{src_dir}}/{{name}}"
    else:
        path = f"{{tests_dir}}/{{name}}"
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as fh:
        fh.write(content)

with open(f"{{pkg_dir}}/pyproject.toml", "w") as fh:
    fh.write(pyproject_src)

print(f"Written {{len(files_map) + 1}} files to {{pkg_dir}}")

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

# Add source directory to path directly
sys.path.insert(0, f"{{pkg_dir}}/src")
print(f"Added {{pkg_dir}}/src to sys.path")

import insurance_copula
print(f"HAS_PYVINECOPULIB = {{insurance_copula.HAS_PYVINECOPULIB}}")

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

env = os.environ.copy()
env["PYTHONPATH"] = f"{{pkg_dir}}/src:" + env.get("PYTHONPATH", "")
r = subprocess.run(
    [sys.executable, "-m", "pytest", tests_dir, "-v", "--tb=short",
     "--no-header", "-p", "no:cacheprovider", "-q"],
    capture_output=True, text=True, cwd=pkg_dir, env=env
)

print(r.stdout[-10000:])
if r.stderr:
    print("STDERR:", r.stderr[-1000:])

if r.returncode == 0:
    print("\\n=== ALL TESTS PASSED ===")
    try:
        dbutils.notebook.exit("ALL TESTS PASSED")
    except NameError:
        pass
else:
    msg = f"TESTS FAILED (exit {{r.returncode}})"
    print(f"\\n=== {{msg}} ===")
    try:
        dbutils.notebook.exit(msg)
    except NameError:
        pass
"""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def api_call(method: str, endpoint: str, body: dict | None = None) -> dict:
    data = json.dumps(body).encode("utf-8") if body else None
    req = urllib.request.Request(
        f"{api_base}/{endpoint}",
        data=data,
        headers=headers,
        method=method,
    )
    try:
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        err_body = e.read().decode()
        raise RuntimeError(f"API {method} {endpoint} failed {e.code}: {err_body}")


# ---------------------------------------------------------------------------
# Ensure workspace folder exists
# ---------------------------------------------------------------------------

print(f"Creating workspace folder {WORKSPACE_FOLDER} ...")
try:
    api_call("POST", "api/2.0/workspace/mkdirs", {"path": WORKSPACE_FOLDER})
    print("Folder ready.")
except RuntimeError as exc:
    print(f"mkdirs note: {exc}")

# ---------------------------------------------------------------------------
# Upload notebook
# ---------------------------------------------------------------------------

print(f"Uploading notebook to {NOTEBOOK_PATH} ...")
notebook_b64 = base64.b64encode(NOTEBOOK_SOURCE.encode("utf-8")).decode("ascii")

api_call("POST", "api/2.0/workspace/import", {
    "path": NOTEBOOK_PATH,
    "format": "SOURCE",
    "language": "PYTHON",
    "content": notebook_b64,
    "overwrite": True,
})
print("Upload OK")

# ---------------------------------------------------------------------------
# Submit serverless run (no cluster spec = serverless)
# ---------------------------------------------------------------------------

print(f"Submitting serverless run {RUN_ID} ...")
submit_body = {
    "run_name": f"insurance-copula-pytest-{RUN_ID}",
    "tasks": [
        {
            "task_key": "pytest",
            "notebook_task": {
                "notebook_path": NOTEBOOK_PATH,
                "source": "WORKSPACE",
            },
        }
    ],
}
result = api_call("POST", "api/2.1/jobs/runs/submit", submit_body)
run_id = result["run_id"]
print(f"Run submitted: run_id={run_id}")

# ---------------------------------------------------------------------------
# Poll
# ---------------------------------------------------------------------------

print("Polling ...")
lc, rs = "PENDING", "-"
for i in range(120):
    time.sleep(15)
    run_state = api_call("GET", f"api/2.1/jobs/runs/get?run_id={run_id}")
    lc = run_state.get("state", {}).get("life_cycle_state", "UNKNOWN")
    rs = run_state.get("state", {}).get("result_state", "-")
    print(f"  [{i * 15}s] {lc} / {rs}")
    if lc in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
        break

print(f"\nFinal state: {lc} / {rs}")

# ---------------------------------------------------------------------------
# Fetch output — fetch per-task run output
# ---------------------------------------------------------------------------

def fetch_output(rid: int) -> None:
    try:
        output = api_call("GET", f"api/2.1/jobs/runs/get-output?run_id={rid}")
        nb_result = output.get("notebook_output", {}).get("result", "")
        error = output.get("error", "")
        error_trace = output.get("error_trace", "")
        logs = output.get("logs", "")
        if nb_result:
            print(f"Exit value: {nb_result}")
        if error:
            print(f"Error: {error}")
        if error_trace:
            print(f"Trace:\n{error_trace[:3000]}")
        if logs:
            print(f"Logs:\n{logs[-8000:]}")
    except Exception as e:
        print(f"Could not fetch output for run {rid}: {e}")


# Fetch per-task output (avoids "multiple tasks" error on parent run_id)
run_details = api_call("GET", f"api/2.1/jobs/runs/get?run_id={run_id}")
tasks = run_details.get("tasks", [])
if tasks:
    for task in tasks:
        task_run_id = task.get("run_id")
        if task_run_id:
            print(f"\n--- Output for task '{task.get('task_key')}' ---")
            fetch_output(task_run_id)
else:
    fetch_output(run_id)

# ---------------------------------------------------------------------------
# Clean up notebook
# ---------------------------------------------------------------------------
try:
    api_call("POST", "api/2.0/workspace/delete", {"path": NOTEBOOK_PATH})
except Exception:
    pass

# ---------------------------------------------------------------------------
# Exit
# ---------------------------------------------------------------------------

success = rs == "SUCCESS"
print(f"\nResult: {'PASSED' if success else 'FAILED'}")
sys.exit(0 if success else 1)
