#!/usr/bin/env python3
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from threading import Thread


# =============================================================================
# EARLY CUDA FABRIC MANAGER KICK (before ANY CUDA-touching imports)
# =============================================================================
# On H200 hosts, cudaGetDeviceCount can return Error 802 "system not yet
# initialized" on first use, because nvidia-fabricmanager on the host
# synchronizes with the container's first driver call. Once any NVML/CUDA
# call succeeds once (even just nvidia-smi), the fabric is up for the rest
# of the container lifetime.
#
# Our previous approach (wait in a subprocess before training) didn't work
# because the "initialization failed" state persisted across calls in the
# same container. The real fix: kick the driver exactly once with
# nvidia-smi, which is what successfully-working baseline containers do
# implicitly via their first torch.cuda call.
#
# Must happen BEFORE `import torch` (because any import that eagerly calls
# cudaGetDeviceCount will cache the Error 802 state).
def _early_cuda_kick() -> None:
    deadline = time.time() + 120.0
    attempt = 0
    while time.time() < deadline:
        attempt += 1
        r = subprocess.run(['nvidia-smi'], capture_output=True, text=True, timeout=30)
        if r.returncode == 0 and 'H200' in (r.stdout or '') or 'H100' in (r.stdout or '') \
                or 'A100' in (r.stdout or '') or r.returncode == 0:
            print(f'[boot] nvidia-smi OK on attempt {attempt}', flush=True)
            break
        print(f'[boot] nvidia-smi attempt {attempt} rc={r.returncode} stderr={(r.stderr or "")[:120]}',
              flush=True)
        time.sleep(2)
    # After nvidia-smi, probe torch in a subprocess so any latent error state
    # doesn't leak into the main process's CUDA context.
    probe = 'import torch; import sys; sys.exit(0 if torch.cuda.is_available() else 1)'
    torch_deadline = time.time() + 120.0
    t_attempt = 0
    while time.time() < torch_deadline:
        t_attempt += 1
        r = subprocess.run([sys.executable, '-c', probe], capture_output=True, text=True, timeout=60)
        if r.returncode == 0:
            print(f'[boot] torch.cuda.is_available() = True after {t_attempt} probe(s)', flush=True)
            return
        if t_attempt == 1:
            print(f'[boot] torch cuda probe {t_attempt}: {(r.stderr or "")[:200]}', flush=True)
        time.sleep(2)
    print('[boot] WARNING: torch.cuda never became ready — training will likely fail', flush=True)


_early_cuda_kick()

# Hydrate triton compilation cache from HF Hub before any triton/mamba_ssm import.
# triton_cache_setup.py is copied next to this file by the job bash command.
try:
    import triton_cache_setup as _tcs
    _tcs.setup()
except ImportError:
    print('[boot] triton_cache_setup not found; skipping cache hydrate', flush=True)

from huggingface_hub import HfApi  # noqa: E402  (import after cuda kick)

REPO_ROOT = Path('/workspace/feather')
CACHE_ROOT = Path.home() / '.cache' / 'autoresearch'
LOG_FILE = REPO_ROOT / 'run_domain_expanded.log'
JOB_ID = os.environ.get('JOB_ID', 'local-job')
OUTPUT_REPO = os.environ.get('HF_REPO_ID', 'icarus112/feather-pretrain-checkpoints')
TOKEN = os.environ.get('HF_TOKEN')
RUNTIME_MODE = os.environ.get('FEATHER_RUNTIME_MODE', 'space')
APP_PORT = int(os.environ.get('PORT', '7860'))


class _HealthHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path in ('/', '/health', '/healthz', '/ready'):
            payload = {
                'status': 'ok',
                'mode': RUNTIME_MODE,
                'job_id': JOB_ID,
            }
            body = json.dumps(payload).encode('utf-8')
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        self.send_response(404)
        self.end_headers()

    def log_message(self, format, *args):
        return


def _start_health_server() -> HTTPServer:
    server = HTTPServer(('0.0.0.0', APP_PORT), _HealthHandler)
    thread = Thread(target=server.serve_forever, daemon=True)
    thread.start()
    print(f'[space] health server listening on 0.0.0.0:{APP_PORT}', flush=True)
    return server


def upload_artifact(api: HfApi, path: Path, dest: str) -> None:
    if not path.exists():
        print(f'[upload] skip missing {path}', flush=True)
        return
    api.upload_file(
        path_or_fileobj=str(path),
        path_in_repo=dest,
        repo_id=OUTPUT_REPO,
        repo_type='model',
    )
    print(f'[upload] uploaded {path} -> {OUTPUT_REPO}/{dest}', flush=True)


def _wait_for_cuda_ready(timeout_s: int = 120) -> None:
    """Block until CUDA is fully initialized or timeout.

    On H200 hosts with NVSwitch/fabric manager, nvidia driver setup can race
    with container start. cudaGetDeviceCount can return CUDA_ERROR_SYSTEM_NOT_READY
    (error 802) for the first few seconds, and any import that triggers
    @triton.autotune (e.g. mamba_ssm, torch amp utilities) blows up with
    "0 active drivers" if it happens during that window.

    We pre-init CUDA in a throwaway Python subprocess (so any error state does
    not leak into the main training process) and retry until torch.cuda
    reports ready.
    """
    import time as _t
    probe = (
        "import torch; "
        "import sys; "
        "avail = torch.cuda.is_available(); "
        "count = torch.cuda.device_count() if avail else 0; "
        "sys.exit(0 if (avail and count > 0) else 1)"
    )
    deadline = _t.time() + timeout_s
    attempt = 0
    while _t.time() < deadline:
        attempt += 1
        r = subprocess.run(['python', '-c', probe], capture_output=True, text=True)
        if r.returncode == 0:
            print(f'[job] CUDA ready after {attempt} probe(s)', flush=True)
            return
        if attempt == 1:
            print(f'[job] CUDA not ready yet (will retry up to {timeout_s}s): {r.stderr.strip()[:200]}', flush=True)
        _t.sleep(2)
    print(f'[job] CUDA still not ready after {timeout_s}s — continuing anyway (training will likely fail)', flush=True)


def _truthy_env(name: str, default: str = '0') -> bool:
    return os.environ.get(name, default).strip().lower() in {'1', 'true', 'yes', 'on'}


def _check_training_artifacts_ready() -> tuple[bool, bool]:
    """Return whether metrics and final checkpoints are visible to the job wrapper."""
    metrics_seen = False
    if LOG_FILE.exists():
        try:
            tail = LOG_FILE.read_text(errors='replace')[-20000:]
            metrics_seen = '[METRICS_JSON]' in tail or '[METRICS] wrote' in tail
        except OSError:
            metrics_seen = False
    checkpoints_ready = (CACHE_ROOT / 'latest.pt').exists() and (CACHE_ROOT / 'pretrain_final.pt').exists()
    return metrics_seen, checkpoints_ready


def _run_training_subprocess(cmd: list[str]) -> int:
    """Run training, optionally stopping after metrics/checkpoints for clean upload.

    Full-corpus streaming can leave dataset downloader worker threads alive during
    Python finalization after useful metrics/checkpoints have already been written.
    On HF Jobs this may keep the job RUNNING or flip it to ERROR before the
    entrypoint uploads artifacts. The watcher preserves the completed canary by
    terminating the train subprocess once the metrics/checkpoint contract is met.
    """
    if not _truthy_env('FEATHER_HF_EXIT_AFTER_METRICS', '1'):
        return subprocess.run(cmd, check=False).returncode

    proc = subprocess.Popen(cmd)
    metrics_seen = False
    checkpoints_ready = False
    while proc.poll() is None:
        metrics_seen, checkpoints_ready = _check_training_artifacts_ready()
        if metrics_seen and checkpoints_ready:
            print('[job] metrics/checkpoints observed; terminating training subprocess for clean artifact upload', flush=True)
            proc.terminate()
            try:
                proc.wait(timeout=30)
            except subprocess.TimeoutExpired:
                print('[job] training subprocess did not terminate cleanly; killing it', flush=True)
                proc.kill()
                proc.wait(timeout=30)
            return 0
        time.sleep(5)

    metrics_seen, checkpoints_ready = _check_training_artifacts_ready()
    if proc.returncode != 0 and metrics_seen and checkpoints_ready:
        print(
            f'[job] training subprocess exited rc={proc.returncode} after writing metrics/checkpoints; treating canary as successful for upload',
            flush=True,
        )
        return 0
    return int(proc.returncode or 0)


def run_job_mode() -> int:
    os.chdir(REPO_ROOT)

    # Dynamic live patch from GitHub to bypass Space build errors
    GIT_REF = os.environ.get('FEATHER_GIT_REF')
    if GIT_REF:
        print(f'[bootstrap] dynamic sync to {GIT_REF}...', flush=True)
        subprocess.run(['git', 'fetch', 'origin'], cwd=REPO_ROOT, check=False)
        subprocess.run(['git', 'checkout', GIT_REF], cwd=REPO_ROOT, check=False)

    os.environ.setdefault('HYDRA_TIME_BUDGET', '43200')
    os.environ.setdefault('HYDRA_TARGET_SHARDS', '2048')
    os.environ.setdefault('HYDRA_DOWNLOAD_WORKERS', '16')
    os.environ.setdefault('HYDRA_CKPT_INTERVAL', '1000')
    os.environ.setdefault('HYDRA_RESUME_CKPT', str(CACHE_ROOT / 'latest.pt'))

    # CUDA readiness was kicked at module import via _early_cuda_kick. Keep
    # the wait as a second safety net — no-op if CUDA already ready.
    _wait_for_cuda_ready()

    cmd = [
        'bash',
        './scripts/run_domain_expanded_pretrain.sh',
        '--target-shards', os.environ['HYDRA_TARGET_SHARDS'],
        '--download-workers', os.environ['HYDRA_DOWNLOAD_WORKERS'],
    ]
    print('[job] starting Feather domain-expanded pretrain', flush=True)
    print(f'[job] command={cmd}', flush=True)
    proc_returncode = _run_training_subprocess(cmd)

    # Push triton compilation cache back to HF Hub for next run.
    try:
        import triton_cache_setup as _tcs
        _tcs.teardown()
    except Exception as _tcs_err:
        print(f'[triton_cache] teardown error (non-fatal): {_tcs_err}', flush=True)

    if TOKEN:
        api = HfApi(token=TOKEN)
        try:
            api.create_repo(repo_id=OUTPUT_REPO, repo_type='model', private=True, exist_ok=True)
        except Exception as e:
            print(f'[upload] create_repo warning: {type(e).__name__}: {e}', flush=True)
        prefix = f'jobs/{JOB_ID}'
        try:
            upload_artifact(api, LOG_FILE, f'{prefix}/run_domain_expanded.log')
            upload_artifact(api, CACHE_ROOT / 'latest.pt', f'{prefix}/latest.pt')
            upload_artifact(api, CACHE_ROOT / 'pretrain_final.pt', f'{prefix}/pretrain_final.pt')
        except Exception as e:
            print(f'[upload] upload warning: {type(e).__name__}: {e}', flush=True)
    else:
        print('[upload] HF_TOKEN not set; skipping artifact upload', flush=True)

    return proc_returncode


def run_space_mode() -> int:
    server = _start_health_server()
    print('[space] Feather runtime image ready', flush=True)
    try:
        while True:
            time.sleep(3600)
    finally:
        server.shutdown()
        server.server_close()


def main() -> int:
    if RUNTIME_MODE == 'job':
        return run_job_mode()
    return run_space_mode()


if __name__ == '__main__':
    raise SystemExit(main())