"""Background training worker.

Runs *real* gradient-descent training of a small 3-layer MLP with NumPy
(no GPU required). Each epoch persists loss to the DB so clients can stream
metrics. To fine-tune real models, replace `_run` with a TRL/transformers
training call (see README "Fine-tuning real models").
"""
import datetime
import threading

import numpy as np

from .. import models as orm
from ..database import SessionLocal

_workers = {}
_cancel = set()


def _forward_backward(X, y, W1, b1, W2, b2):
    z = X @ W1 + b1
    h = np.tanh(z)
    logits = h @ W2 + b2
    ex = np.exp(logits - logits.max(axis=1, keepdims=True))
    p = ex / ex.sum(axis=1, keepdims=True)
    loss = -np.log(p[np.arange(len(y)), y]).mean()

    dp = p.copy()
    dp[np.arange(len(y)), y] -= 1.0
    dp /= len(y)

    dW2 = h.T @ dp
    db2 = dp.sum(0)
    dh = dp @ W2.T
    dz = dh * (1.0 - np.tanh(z) ** 2)
    dW1 = X.T @ dz
    db1 = dz.sum(0)

    return loss, dW1, db1, dW2, db2


def _run(job_id: str):
    db = SessionLocal()
    try:
        job = db.get(orm.TrainingJob, job_id)
        if job is None:
            return
        cfg = job.config or {}
        rng = np.random.default_rng(int(cfg.get("seed", 42)))
        n, d, k = 600, 8, 4
        X = rng.normal(size=(n, d))
        true_W = rng.normal(size=(d, d))
        y = (np.argmax(X @ true_W + 0.4 * rng.normal(size=(n, d)), axis=1) % k).astype(int)

        W1 = rng.normal(size=(d, 32)) * 0.1
        b1 = np.zeros(32)
        W2 = rng.normal(size=(32, k)) * 0.1
        b2 = np.zeros(k)

        lr = float(cfg.get("lr", 0.05))
        epochs = int(cfg.get("epochs", 20))
        bs = int(cfg.get("batch_size", 32))

        job.status = "running"
        job.started_at = datetime.datetime.utcnow()
        loss_history = (job.metrics or {}).get("loss") if isinstance(job.metrics, dict) else None
        loss_history = list(loss_history) if loss_history else []
        db.commit()

        for ep in range(epochs):
            if job_id in _cancel:
                job.status = "cancelled"
                job.finished_at = datetime.datetime.utcnow()
                db.commit()
                return

            perm = rng.permutation(n)
            tot, nb = 0.0, 0
            for i in range(0, n, bs):
                idx = perm[i:i + bs]
                loss, dW1, db1, dW2, db2 = _forward_backward(X[idx], y[idx], W1, b1, W2, b2)
                W2 -= lr * dW2
                b2 -= lr * db2
                W1 -= lr * dW1
                b1 -= lr * db1
                tot += loss
                nb += 1

            avg = float(tot / nb)
            loss_history.append(avg)
            job.current_epoch = ep + 1
            job.loss = avg
            job.metrics = {"loss": loss_history[-200:], "lr": lr, "batch_size": bs}
            db.commit()

        h = np.tanh(X @ W1 + b1)
        preds = np.argmax(h @ W2 + b2, axis=1)
        acc = float((preds == y).mean())
        job.status = "completed"
        job.finished_at = datetime.datetime.utcnow()
        job.metrics = {"loss": loss_history[-200:], "accuracy": acc}
        db.commit()
    except Exception as e:  # noqa: BLE001
        try:
            job.status = "failed"
            job.error = str(e)
            job.finished_at = datetime.datetime.utcnow()
            db.commit()
        except Exception:
            pass
    finally:
        db.close()


def launch(job_id: str) -> bool:
    """Start a training thread for the job; returns False if already running."""
    if job_id in _workers and _workers[job_id].is_alive():
        return False
    t = threading.Thread(target=_run, args=(job_id,), daemon=True)
    _workers[job_id] = t
    t.start()
    return True


def cancel(job_id: str) -> None:
    _cancel.add(job_id)


def is_running(job_id: str) -> bool:
    return job_id in _workers and _workers[job_id].is_alive()