"""
Run insurance-monitoring test suite on Databricks serverless compute.
Uploads current state of the repo and runs pytest.
"""

import os
import sys
import time
import base64
import pathlib

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] = v

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

w = WorkspaceClient()

WORKSPACE_PATH = "/Workspace/insurance-monitoring-bench"
PROJECT_ROOT = pathlib.Path("/home/ralph/burning-cost/repos/insurance-monitoring")


def upload_file(local_path: pathlib.Path, remote_path: str) -> None:
    content = local_path.read_bytes()
    encoded = base64.b64encode(content).decode()
    parent = "/".join(remote_path.split("/")[:-1])
    try:
        w.workspace.mkdirs(path=parent)
    except Exception:
        pass
    w.workspace.import_(
        path=remote_path,
        content=encoded,
        overwrite=True,
        format=ImportFormat.AUTO,
    )


def upload_notebook(local_path: pathlib.Path, remote_path: str) -> None:
    content = local_path.read_bytes()
    encoded = base64.b64encode(content).decode()
    parent = "/".join(remote_path.split("/")[:-1])
    try:
        w.workspace.mkdirs(path=parent)
    except Exception:
        pass
    w.workspace.import_(
        path=remote_path,
        content=encoded,
        overwrite=True,
        format=ImportFormat.SOURCE,
        language=Language.PYTHON,
    )


# Upload src/ and tests/
upload_files = []
for rel_dir in ["src", "tests"]:
    for p in (PROJECT_ROOT / rel_dir).rglob("*.py"):
        if "__pycache__" in p.parts:
            continue
        upload_files.append(p)
upload_files.append(PROJECT_ROOT / "pyproject.toml")

print(f"Uploading {len(upload_files)} source/test files...")
for local in sorted(upload_files):
    rel = local.relative_to(PROJECT_ROOT)
    remote = f"{WORKSPACE_PATH}/{rel}"
    upload_file(local, remote)
    print(f"  {rel}")

# Upload test runner notebook
nb_local = PROJECT_ROOT / "notebooks" / "run_tests_v070.py"
nb_remote_name = "run_tests_bench"
nb_remote = f"{WORKSPACE_PATH}/notebooks/{nb_remote_name}"

# We need a slightly adapted notebook that references our workspace path
nb_content = f"""# Databricks notebook source
# COMMAND ----------
# MAGIC %pip install polars scipy matplotlib pytest sortedcontainers

# COMMAND ----------
dbutils.library.restartPython()

# COMMAND ----------
import subprocess
import sys
import os

workspace_path = "{WORKSPACE_PATH}"
sys.path.insert(0, f"{{workspace_path}}/src")

import insurance_monitoring
print(f"insurance_monitoring version: {{insurance_monitoring.__version__}}")

from insurance_monitoring import PITMonitor, InterpretableDriftDetector
print("PITMonitor and InterpretableDriftDetector available (v0.7.0)")

# COMMAND ----------
env = {{
    **os.environ,
    "PYTHONPATH": f"{{workspace_path}}/src",
    "PYTHONDONTWRITEBYTECODE": "1",
}}

result = subprocess.run(
    [
        sys.executable, "-m", "pytest",
        f"{{workspace_path}}/tests",
        "-v", "--tb=long",
        "--import-mode=importlib",
        "-p", "no:cacheprovider",
        f"--rootdir={{workspace_path}}",
    ],
    capture_output=True,
    text=True,
    env=env,
)
output = result.stdout + "\\n" + result.stderr
print(output)
dbutils.notebook.exit(output[-8000:] if len(output) > 8000 else output)
"""

encoded_nb = base64.b64encode(nb_content.encode()).decode()
try:
    w.workspace.mkdirs(path=f"{WORKSPACE_PATH}/notebooks")
except Exception:
    pass
w.workspace.import_(
    path=nb_remote,
    content=encoded_nb,
    overwrite=True,
    format=ImportFormat.SOURCE,
    language=Language.PYTHON,
)
print(f"Uploaded notebook: {nb_remote}")

print("\nSubmitting test job...")

run = w.jobs.submit(
    run_name="insurance-monitoring-bench-tests",
    tasks=[
        jobs.SubmitTask(
            task_key="run_tests",
            environment_key="default",
            notebook_task=jobs.NotebookTask(
                notebook_path=nb_remote,
            ),
        )
    ],
    environments=[
        jobs.JobEnvironment(
            environment_key="default",
            spec=compute.Environment(),
        )
    ],
)

run_id = run.run_id
print(f"Run ID: {run_id}")
print(f"Monitor: {os.environ['DATABRICKS_HOST']}#job/runs/{run_id}")

print("\nPolling (max 15 min)...")
deadline = time.time() + 900
while time.time() < deadline:
    state = w.jobs.get_run(run_id=run_id)
    life = state.state.life_cycle_state.value
    print(f"  {life}")
    if life in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
        break
    time.sleep(20)

result_state = state.state.result_state
print(f"\nResult: {result_state.value if result_state else 'UNKNOWN'}")

try:
    task_run = state.tasks[0] if state.tasks else None
    if task_run:
        output = w.jobs.get_run_output(run_id=task_run.run_id)
        if output.notebook_output and output.notebook_output.result:
            print("\n--- Output ---")
            print(output.notebook_output.result)
        if output.error:
            print(f"\nError: {output.error}")
        if output.error_trace:
            print(f"\nTrace:\n{output.error_trace}")
except Exception as e:
    print(f"Could not fetch output: {e}")

if result_state and result_state.value == "SUCCESS":
    print("\nAll tests PASSED.")
    sys.exit(0)
else:
    print("\nTests FAILED or incomplete.")
    sys.exit(1)
