"""Submit a one-time Databricks job to run pytest on insurance-transfer using serverless compute."""

import os
import time
import base64
import requests
import re

os.environ["DATABRICKS_HOST"] = "https://dbc-150a27f5-e1e7.cloud.databricks.com/"
os.environ["DATABRICKS_TOKEN"] = "dapi895b482f1aaae1b0cfaab97cbfe4546b"

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

# Find a writable temp directory; use pytest --import-mode=importlib to avoid __pycache__
test_notebook_content = (
    "# Databricks notebook source\n"
    "# MAGIC %pip install numpy scipy scikit-learn statsmodels pytest -q\n"
    "\n"
    "# COMMAND ----------\n"
    "import subprocess, sys, tempfile, os, shutil\n"
    "\n"
    "# Install package\n"
    "r = subprocess.run(\n"
    "    [sys.executable, '-m', 'pip', 'install', '-e', '/Workspace/insurance-transfer/', '-q'],\n"
    "    capture_output=True, text=True\n"
    ")\n"
    "print('pip returncode:', r.returncode)\n"
    "if r.returncode != 0:\n"
    "    print(r.stderr[-500:])\n"
    "\n"
    "# Find writable directory\n"
    "for candidate in ['/tmp', '/var/tmp', tempfile.gettempdir()]:\n"
    "    try:\n"
    "        td = os.path.join(candidate, 'it_tests')\n"
    "        os.makedirs(td, exist_ok=True)\n"
    "        print('Writable dir:', td)\n"
    "        break\n"
    "    except Exception as e:\n"
    "        print(f'{candidate} not writable: {e}')\n"
    "        td = None\n"
    "\n"
    "# Copy test files\n"
    "for f in os.listdir('/Workspace/insurance-transfer/tests'):\n"
    "    if f.endswith('.py'):\n"
    "        shutil.copy(\n"
    "            f'/Workspace/insurance-transfer/tests/{f}',\n"
    "            os.path.join(td, f)\n"
    "        )\n"
    "print('Tests in', td, ':', os.listdir(td))\n"
    "\n"
    "# COMMAND ----------\n"
    "import subprocess, sys, base64, json, os\n"
    "from urllib.request import Request, urlopen\n"
    "\n"
    "td = '/tmp/it_tests'\n"
    "\n"
    "r = subprocess.run(\n"
    "    [sys.executable, '-m', 'pytest', td,\n"
    "     '--tb=short', '-v',\n"
    "     '--import-mode=importlib',\n"
    "     '--ignore=' + os.path.join(td, 'test_gbm_transfer.py'),\n"
    "     '--ignore=' + os.path.join(td, 'test_cann_transfer.py'),\n"
    "    ],\n"
    "    capture_output=True, text=True,\n"
    "    cwd=td,\n"
    ")\n"
    "run_out = r.stdout + '\\nSTDERR:\\n' + r.stderr\n"
    "\n"
    "# Save output to workspace\n"
    "token = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()\n"
    "hdrs = {'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'}\n"
    "payload = {\n"
    "    'path': '/Workspace/insurance-transfer-tests/run_output.txt',\n"
    "    'format': 'AUTO',\n"
    "    'content': base64.b64encode(run_out.encode()).decode(),\n"
    "    'overwrite': True,\n"
    "    'language': 'PYTHON',\n"
    "}\n"
    "urlopen(Request(\n"
    "    'https://dbc-150a27f5-e1e7.cloud.databricks.com/api/2.0/workspace/import',\n"
    "    data=json.dumps(payload).encode(),\n"
    "    headers=hdrs,\n"
    "))\n"
    "\n"
    "print('returncode:', r.returncode)\n"
    "lines = run_out.splitlines()\n"
    "print('\\n'.join(lines[-30:]))\n"
    "assert r.returncode == 0, f'TESTS FAILED code={r.returncode}'\n"
    "print('ALL TESTS PASSED')\n"
)

payload = {
    "path": "/Workspace/insurance-transfer-tests/run_tests",
    "format": "SOURCE",
    "language": "PYTHON",
    "content": base64.b64encode(test_notebook_content.encode()).decode(),
    "overwrite": True,
}
resp = requests.post(f"{host}/api/2.0/workspace/import", json=payload, headers=headers)
if resp.status_code != 200:
    print(f"Upload failed: {resp.status_code} {resp.text}")
    raise SystemExit(1)
print("Uploaded test notebook")

run_payload = {
    "run_name": "insurance-transfer-pytest-v7",
    "tasks": [
        {
            "task_key": "pytest",
            "notebook_task": {
                "notebook_path": "/Workspace/insurance-transfer-tests/run_tests",
                "source": "WORKSPACE",
            },
            "environment_key": "default",
        }
    ],
    "environments": [
        {
            "environment_key": "default",
            "spec": {
                "client": "2",
                "dependencies": [
                    "numpy",
                    "scipy",
                    "scikit-learn",
                    "statsmodels",
                    "pytest",
                ],
            },
        }
    ],
}

resp = requests.post(f"{host}/api/2.1/jobs/runs/submit", json=run_payload, headers=headers)
if resp.status_code != 200:
    print(f"Submit failed: {resp.status_code} {resp.text}")
    raise SystemExit(1)
run_id = resp.json()["run_id"]
print(f"Run ID: {run_id}")


def read_workspace_file(path):
    resp = requests.get(
        f"{host}/api/2.0/workspace/export",
        params={"path": path, "format": "SOURCE"},
        headers=headers,
    )
    if resp.status_code == 200:
        data = resp.json()
        return base64.b64decode(data.get("content", "")).decode(errors="replace")
    return None


for i in range(60):
    time.sleep(20)
    status_resp = requests.get(f"{host}/api/2.1/jobs/runs/get", params={"run_id": run_id}, headers=headers)
    status = status_resp.json()
    state = status["state"]["life_cycle_state"]
    print(f"[{(i+1)*20}s] {state}")
    if state in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
        result_state = status["state"].get("result_state", "UNKNOWN")
        print(f"Result: {result_state}")

        content = read_workspace_file("/Workspace/insurance-transfer-tests/run_output.txt")
        if content:
            print("\n=== PYTEST OUTPUT ===")
            clean = re.sub(r'\x1b\[[0-9;]*m', '', content)
            print(clean[:10000])

        for task in status.get("tasks", []):
            task_run_id = task["run_id"]
            out_resp = requests.get(
                f"{host}/api/2.1/jobs/runs/get-output",
                params={"run_id": task_run_id},
                headers=headers,
            )
            out = out_resp.json()
            nb_out = out.get("notebook_output", {})
            if nb_out.get("result"):
                print("\n=== NOTEBOOK RESULT ===")
                print(re.sub(r'\x1b\[[0-9;]*m', '', nb_out["result"]))
            if out.get("error"):
                print("\n=== ERROR ===")
                print(out["error"])
            if out.get("error_trace"):
                print("\n=== TRACE ===")
                trace = re.sub(r'\x1b\[[0-9;]*m', '', out["error_trace"])
                print(trace[-2000:])
        break
else:
    print("Timed out")
