"""Run insurance-copula tests on Databricks.

Uploads the project, installs dependencies (including pyvinecopulib
via binary wheel), and runs pytest. Use this on x86_64 where
pyvinecopulib wheels are available.

Usage:
    python run_tests_databricks.py
"""

from __future__ import annotations

import os
import sys
import time
import base64
import subprocess
import textwrap
from pathlib import Path

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

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import (
    NotebookTask,
    Task,
    JobSettings,
    Source,
)

WORKSPACE_PATH = "/Workspace/insurance-copula-tests"
NOTEBOOK_PATH = f"{WORKSPACE_PATH}/run_tests"

NOTEBOOK_CONTENT = textwrap.dedent('''
# Databricks notebook source
# MAGIC %md
# MAGIC # insurance-copula test runner

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

import subprocess
import sys

# Install pyvinecopulib via binary wheel (source build not available on serverless)
result = subprocess.run(
    [sys.executable, "-m", "pip", "install",
     "--only-binary=:all:",
     "--quiet",
     "pyvinecopulib>=0.7"],
    capture_output=True, text=True
)
print("pyvinecopulib install stdout:", result.stdout[-500:] if result.stdout else "")
print("pyvinecopulib install stderr:", result.stderr[-500:] if result.stderr else "")
print("Return code:", result.returncode)

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

# Install other dependencies
%pip install scipy numpy pandas matplotlib pytest pytest-cov --quiet

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

import sys
sys.path.insert(0, "/Workspace/insurance-copula-tests/src")

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

# Verify pyvinecopulib import
try:
    import pyvinecopulib as pv
    print(f"pyvinecopulib available: {pv.__version__ if hasattr(pv, '__version__') else 'yes'}")
except ImportError as e:
    print(f"pyvinecopulib NOT available: {e}")

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

import subprocess, sys

result = subprocess.run(
    [sys.executable, "-m", "pytest",
     "/Workspace/insurance-copula-tests/tests/",
     "-v", "--tb=short", "-q",
     "--no-header"],
    capture_output=True, text=True,
    cwd="/Workspace/insurance-copula-tests"
)
print(result.stdout[-5000:])
if result.stderr:
    print("STDERR:", result.stderr[-2000:])
print("Return code:", result.returncode)
''')


def upload_directory(w: WorkspaceClient, local_dir: Path, workspace_dir: str) -> None:
    """Upload all Python files to workspace."""
    from databricks.sdk.service.workspace import ImportFormat, Language

    # Create directory
    try:
        w.workspace.mkdirs(workspace_dir)
    except Exception:
        pass

    for path in local_dir.rglob("*"):
        if path.is_dir():
            rel = path.relative_to(local_dir)
            ws_path = f"{workspace_dir}/{rel}".replace("\\", "/")
            try:
                w.workspace.mkdirs(ws_path)
            except Exception:
                pass
        elif path.suffix == ".py" and ".venv" not in str(path):
            rel = path.relative_to(local_dir)
            ws_path = f"{workspace_dir}/{rel}".replace("\\", "/")
            content = path.read_bytes()
            w.workspace.upload(
                ws_path,
                content,
                overwrite=True,
                format=ImportFormat.AUTO,
            )
            print(f"  Uploaded: {rel}")


def upload_notebook(w: WorkspaceClient, path: str, content: str) -> None:
    from databricks.sdk.service.workspace import ImportFormat, Language, ObjectType

    encoded = base64.b64encode(content.encode()).decode()
    w.workspace.import_(
        path=path,
        content=encoded,
        overwrite=True,
        format=ImportFormat.SOURCE,
        language=Language.PYTHON,
    )
    print(f"Uploaded notebook: {path}")


def run_job(w: WorkspaceClient, notebook_path: str) -> int:
    """Create and run a one-shot job, return exit code."""
    from databricks.sdk.service.jobs import (
        RunNow,
        NotebookTask,
        Task,
        Source,
    )
    from databricks.sdk.service import jobs as j

    job = w.jobs.create(
        name="insurance-copula-test-run",
        tasks=[
            j.Task(
                task_key="run_tests",
                notebook_task=j.NotebookTask(
                    notebook_path=notebook_path,
                    source=j.Source.WORKSPACE,
                ),
                new_cluster=j.ClusterSpec(
                    spark_version="15.4.x-scala2.12",
                    node_type_id="m5d.large",
                    num_workers=0,
                    spark_conf={"spark.master": "local[*, 4]"},
                ),
            )
        ],
    )
    print(f"Created job {job.job_id}")

    run = w.jobs.run_now(job_id=job.job_id)
    run_id = run.run_id
    print(f"Started run {run_id}")

    while True:
        state = w.jobs.get_run(run_id=run_id)
        lc = state.state.life_cycle_state.value if state.state.life_cycle_state else "UNKNOWN"
        rs = state.state.result_state.value if state.state.result_state else ""
        print(f"  Status: {lc} {rs}")
        if lc in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
            break
        time.sleep(15)

    # Get output
    output = w.jobs.get_run_output(run_id=run_id)
    if output.notebook_output:
        print("\n--- Notebook output ---")
        print(output.notebook_output.result)

    # Cleanup
    w.jobs.delete(job_id=job.job_id)

    result_state = state.state.result_state.value if state.state.result_state else ""
    return 0 if result_state == "SUCCESS" else 1


def main() -> None:
    project_root = Path(__file__).parent
    w = WorkspaceClient()

    print("Uploading project files...")
    upload_directory(w, project_root, WORKSPACE_PATH)

    print("Uploading test notebook...")
    upload_notebook(w, NOTEBOOK_PATH, NOTEBOOK_CONTENT)

    print("Running tests on Databricks...")
    exit_code = run_job(w, NOTEBOOK_PATH)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
