"""Logging, atomic writes, stage state, GPU discovery."""
import json
import logging
import os
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd

from config import CFG, LOGS, STATE, WORK


def get_logger(name):
    LOGS.mkdir(parents=True, exist_ok=True)
    log = logging.getLogger(name)
    if not log.handlers:
        log.setLevel(logging.INFO)
        fmt = logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s", "%m-%d %H:%M:%S")
        h = logging.StreamHandler(sys.stdout)
        h.setFormatter(fmt)
        log.addHandler(h)
        log.propagate = False
    return log


def _tmp(path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    return path.with_name(f".{path.name}.tmp{os.getpid()}")


def save_parquet(df, path):
    tmp = _tmp(path)
    df.to_parquet(tmp, index=False)
    os.replace(tmp, path)


def save_npz(path, **arrays):
    tmp = _tmp(path)
    with open(tmp, "wb") as f:
        np.savez(f, **arrays)
    os.replace(tmp, path)


def save_npy(path, arr):
    tmp = _tmp(path)
    with open(tmp, "wb") as f:
        np.save(f, arr)
    os.replace(tmp, path)


def save_json(obj, path):
    tmp = _tmp(path)
    tmp.write_text(json.dumps(obj, indent=2, default=str))
    os.replace(tmp, path)


def load_json(path):
    return json.loads(Path(path).read_text())


def stage_done(name):
    return (STATE / f"{name}.done").exists()


def mark_done(name, info=None):
    save_json({"finished": time.strftime("%Y-%m-%d %H:%M:%S"), **(info or {})}, STATE / f"{name}.done")


class Progress:
    def __init__(self, log, total, what, every=60):
        self.log, self.total, self.what, self.every = log, total, what, every
        self.t0 = self.last = time.time()
        self.done = 0

    def step(self, n=1):
        self.done += n
        now = time.time()
        if now - self.last >= self.every or self.done >= self.total:
            self.last = now
            rate = self.done / max(now - self.t0, 1e-9)
            eta = (self.total - self.done) / rate if rate > 0 else float("nan")
            self.log.info(f"{self.what}: {self.done}/{self.total} ({100 * self.done / max(self.total, 1):.1f}%) "
                          f"elapsed {(now - self.t0) / 60:.1f}m eta {eta / 60:.1f}m")


def gpu_status():
    """[(index, free_mb, util_pct)] for the GPUs this process may use."""
    try:
        out = subprocess.run(["nvidia-smi", "--query-gpu=index,memory.free,utilization.gpu",
                              "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=30).stdout
    except Exception:
        return []
    allowed = os.environ.get("M4_GPUS")
    allowed = {int(x) for x in allowed.split(",")} if allowed else None
    res = []
    for line in out.strip().splitlines():
        i, free, util = [int(x.strip()) for x in line.split(",")]
        if allowed is None or i in allowed:
            res.append((i, free, util))
    return res


def gpu_patient(fn, log, what, wait_s=60, max_wait_s=3 * 3600):
    """Run fn(); on CUDA OOM (other users filling the shared GPU) free cache, wait, retry."""
    import torch
    t0 = time.time()
    while True:
        try:
            return fn()
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            if time.time() - t0 > max_wait_s:
                raise
            log.warning(f"{what}: GPU out of memory (shared card busy), retrying in {wait_s}s")
            time.sleep(wait_s)


def n_workers():
    return max(1, min(CFG["n_workers"], (os.cpu_count() or 8) // 2))


def shard_path(kind, split, i, ext="parquet", sub=None):
    base = WORK / kind / split
    if sub:
        base = base / sub
    return base / f"shard_{i:05d}.{ext}"


def read_shards(kind, split, n, columns=None, sub=None, ext="parquet"):
    return pd.concat([pd.read_parquet(shard_path(kind, split, i, ext, sub), columns=columns) for i in range(n)],
                     ignore_index=True)