"""
Re-run insurance-optimise tests using pip install -e . (editable) instead of PYTHONPATH.
"""

import os, base64, time, json
from pathlib import Path

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

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

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

ENV_KEY = "optimise-test-env-v4"
WS_BASE = "/Workspace/Shared/test-sweep-batch2"

NB_CONTENT = '''# Databricks notebook source
import subprocess, sys, os, shutil, json

name       = "insurance-optimise"
github_url = "https://github.com/burning-cost/insurance-optimise"

work_dir = f"/tmp/work-{name}"
shutil.rmtree(work_dir, ignore_errors=True)

# Clone
print(f"Cloning {github_url}...")
r_clone = subprocess.run(
    ["git", "clone", "--depth=1", github_url, work_dir],
    capture_output=True, text=True, timeout=120
)
if r_clone.returncode != 0:
    dbutils.notebook.exit(json.dumps({"status": "CLONE_FAILED", "output": r_clone.stderr[-2000:]}))
print("Clone OK")

# Show head of demand/__init__.py to confirm what was cloned
demand_init = os.path.join(work_dir, "src/insurance_optimise/demand/__init__.py")
with open(demand_init) as f:
    lines = f.readlines()
print(f"demand/__init__.py first 70 lines:")
for i, line in enumerate(lines[:70], 1):
    print(f"  {i}: {line}", end="")

# Show what pip says about insurance-optimise pre-installed
r_pip_show = subprocess.run(
    [sys.executable, "-m", "pip", "show", "insurance-optimise"],
    capture_output=True, text=True
)
print(f"\\nPre-existing install: {r_pip_show.stdout[:300] if r_pip_show.returncode == 0 else 'not installed'}")

# Uninstall any pre-existing version first
subprocess.run(
    [sys.executable, "-m", "pip", "uninstall", "insurance-optimise", "-y"],
    capture_output=True
)
subprocess.run(
    [sys.executable, "-m", "pip", "uninstall", "insurance-demand", "-y"],
    capture_output=True
)

# Install deps first
base_deps = [
    "numpy>=1.24", "polars>=1.0", "pandas>=2.0", "scikit-learn>=1.3",
    "scipy>=1.10", "statsmodels>=0.14.5", "pyarrow>=12",
    "pytest>=7", "pytest-cov>=4", "matplotlib"
]
print(f"\\nInstalling base deps...")
r_deps = subprocess.run(
    [sys.executable, "-m", "pip", "install"] + base_deps + ["--quiet"],
    capture_output=True, text=True, timeout=300
)
if r_deps.returncode != 0:
    dbutils.notebook.exit(json.dumps({"status": "DEPS_FAILED", "output": (r_deps.stdout + r_deps.stderr)[-1000:]}))
print("Deps OK")

# Install package in editable mode
print("Running pip install -e ...")
r_install = subprocess.run(
    [sys.executable, "-m", "pip", "install", "-e", work_dir, "--no-deps"],
    capture_output=True, text=True, timeout=120, cwd=work_dir
)
install_out = (r_install.stdout + r_install.stderr).strip()
print(f"Install output:\\n{install_out[-500:]}")
if r_install.returncode != 0:
    dbutils.notebook.exit(json.dumps({"status": "INSTALL_FAILED", "output": install_out[-1000:]}))
print("Install OK")

# Run pytest
log_file = f"/tmp/pytest-{name}.txt"
print("Running pytest...")
with open(log_file, "w") as f:
    r_pytest = subprocess.run(
        [sys.executable, "-m", "pytest", "tests/", "-v", "--tb=short", "-p", "no:cacheprovider"],
        stdout=f, stderr=subprocess.STDOUT,
        timeout=360, cwd=work_dir
    )

with open(log_file) as f:
    output = f.read()

all_lines = output.strip().split("\\n")
tail = "\\n".join(all_lines[-80:])
print(f"--- pytest tail ---\\n{tail}")

summary = ""
for line in reversed(all_lines):
    s = line.strip()
    if s and ("passed" in s or "failed" in s or "error" in s):
        summary = s
        break

status = "PASSED" if r_pytest.returncode == 0 else "FAILED"
dbutils.notebook.exit(json.dumps({
    "status": status,
    "returncode": r_pytest.returncode,
    "summary": summary,
    "tail": tail,
}))
'''

nb_path = f"{WS_BASE}/test_insurance_optimise_v4"
encoded = base64.b64encode(NB_CONTENT.encode()).decode()

try:
    w.workspace.delete(nb_path)
except Exception:
    pass
w.workspace.import_(
    path=nb_path,
    content=encoded,
    overwrite=True,
    format=ImportFormat.SOURCE,
    language=Language.PYTHON,
)
print(f"Notebook uploaded: {nb_path}")

run_waiter = w.jobs.submit(
    run_name="test-optimise-v4",
    environments=[
        jobs.JobEnvironment(
            environment_key=ENV_KEY,
            spec=compute.Environment(client="2"),
        )
    ],
    tasks=[
        jobs.SubmitTask(
            task_key="run-tests",
            environment_key=ENV_KEY,
            notebook_task=jobs.NotebookTask(notebook_path=nb_path),
        )
    ],
)
run_id = run_waiter.run_id
print(f"Run ID: {run_id}")

while True:
    state = w.jobs.get_run(run_id=run_id)
    life_str = state.state.life_cycle_state.value if hasattr(state.state.life_cycle_state, "value") else str(state.state.life_cycle_state)
    result_str = state.state.result_state.value if state.state.result_state and hasattr(state.state.result_state, "value") else str(state.state.result_state or "")
    print(f"  [{time.strftime('%H:%M:%S')}] {life_str}  {result_str}")
    if any(s in life_str for s in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR")):
        break
    time.sleep(20)

nb_result = None
for task in (state.tasks or []):
    try:
        out = w.jobs.get_run_output(run_id=task.run_id)
        if out.notebook_output and out.notebook_output.result:
            nb_result = out.notebook_output.result
        if out.error:
            print(f"Error: {out.error}")
        if out.error_trace:
            print(f"Trace: {out.error_trace[-3000:]}")
    except Exception as e:
        print(f"Could not fetch: {e}")

if nb_result:
    print(f"\n=== NOTEBOOK OUTPUT ===\n{nb_result}")
    try:
        res = json.loads(nb_result)
        print(f"\n[insurance-optimise] Status: {res.get('status')} rc={res.get('returncode')}")
        print(f"Summary: {res.get('summary','')}")
    except Exception:
        pass
