from __future__ import annotations

import base64
import hashlib
import json
import os
import uuid
from datetime import UTC
from datetime import datetime
from datetime import timedelta
from pathlib import Path
from threading import Lock
from typing import Any


ALLOWED_TRAINING_ARTIFACTS = {
    "lstm_forecaster.json",
    "torch_retrain_guard.json",
    "torch_threshold_calibration.json",
}
RUNNING_TIMEOUT = timedelta(hours=12)
ONLINE_WINDOW = timedelta(minutes=3)


class TrainingCoordinator:
    def __init__(self, runtime_dir: Path) -> None:
        self.runtime_dir = runtime_dir
        self.state_path = runtime_dir / "training_coordination.json"
        self.upload_root = runtime_dir / ".training_uploads"
        self._lock = Lock()

    def status(self) -> dict[str, Any]:
        with self._lock:
            state = self._load_state()
            self._expire_stale_jobs(state)
            self._save_state(state)
            return self._public_status(state)

    def request_retrain(self, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        payload = payload or {}
        with self._lock:
            state = self._load_state()
            self._expire_stale_jobs(state)
            existing = self._active_job(state)
            if existing is not None:
                self._save_state(state)
                return {"queued": False, "reason": "active_job_exists", "job": existing, "status": self._public_status(state)}

            now = _now()
            job = {
                "id": str(uuid.uuid4()),
                "status": "pending",
                "requested_at": now,
                "requested_by": str(payload.get("source") or "api"),
                "parameters": _safe_parameters(payload.get("parameters")),
                "message": "",
                "artifacts": [],
            }
            state.setdefault("jobs", []).append(job)
            self._trim_jobs(state)
            self._save_state(state)
            return {"queued": True, "job": job, "status": self._public_status(state)}

    def heartbeat(self, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        payload = payload or {}
        with self._lock:
            state = self._load_state()
            worker = self._worker_from_payload(payload)
            state["worker"] = worker
            self._save_state(state)
            return {"ok": True, "worker": worker, "status": self._public_status(state)}

    def claim(self, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        payload = payload or {}
        with self._lock:
            state = self._load_state()
            self._expire_stale_jobs(state)
            worker = self._worker_from_payload(payload)
            state["worker"] = worker
            job = self._oldest_pending_job(state)
            if job is None:
                self._save_state(state)
                return {"claimed": False, "job": None, "status": self._public_status(state)}

            now = _now()
            job["status"] = "running"
            job["claimed_at"] = now
            job["claimed_by"] = worker["id"]
            job["worker"] = worker
            self._save_state(state)
            return {"claimed": True, "job": job, "status": self._public_status(state)}

    def save_artifact_chunk(self, job_id: str, payload: dict[str, Any]) -> dict[str, Any]:
        name = Path(str(payload.get("name") or "")).name
        if name not in ALLOWED_TRAINING_ARTIFACTS:
            raise ValueError(f"artifact is not allowed: {name}")
        index = int(payload.get("index", -1))
        total = int(payload.get("total", 0))
        sha256 = str(payload.get("sha256") or "").strip().lower()
        if index < 0 or total <= 0 or index >= total:
            raise ValueError("invalid artifact chunk index")
        if not sha256:
            raise ValueError("artifact sha256 is required")
        try:
            chunk = base64.b64decode(str(payload.get("data_base64") or ""), validate=True)
        except (ValueError, TypeError) as exc:
            raise ValueError("invalid artifact chunk payload") from exc

        chunk_dir = self.upload_root / job_id / name
        chunk_dir.mkdir(parents=True, exist_ok=True)
        (chunk_dir / f"{index:06d}.part").write_bytes(chunk)

        if not all((chunk_dir / f"{part:06d}.part").is_file() for part in range(total)):
            return {"complete": False, "received": index + 1, "total": total}

        target_tmp = self.runtime_dir / f".{name}.{job_id}.tmp"
        digest = hashlib.sha256()
        with target_tmp.open("wb") as output:
            for part in range(total):
                data = (chunk_dir / f"{part:06d}.part").read_bytes()
                digest.update(data)
                output.write(data)
        if digest.hexdigest().lower() != sha256:
            target_tmp.unlink(missing_ok=True)
            raise ValueError("artifact sha256 mismatch")

        self.runtime_dir.mkdir(parents=True, exist_ok=True)
        os.replace(target_tmp, self.runtime_dir / name)
        _remove_tree(chunk_dir)

        with self._lock:
            state = self._load_state()
            job = self._job_by_id(state, job_id)
            if job is not None:
                artifacts = job.setdefault("artifacts", [])
                artifacts = [item for item in artifacts if item.get("name") != name]
                artifacts.append({"name": name, "sha256": sha256, "uploaded_at": _now()})
                job["artifacts"] = artifacts
                self._save_state(state)
        return {"complete": True, "name": name, "sha256": sha256}

    def progress(self, job_id: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        payload = payload or {}
        with self._lock:
            state = self._load_state()
            job = self._job_by_id(state, job_id)
            if job is None:
                raise ValueError(f"training job not found: {job_id}")
            if isinstance(payload.get("worker"), dict):
                state["worker"] = self._worker_from_payload(payload["worker"])
            job["status"] = str(payload.get("status") or job.get("status") or "running")
            job["phase"] = str(payload.get("phase") or job.get("phase") or "running")
            job["message"] = str(payload.get("message") or job.get("message") or "")
            job["progress_percent"] = _coerce_percent(payload.get("progress_percent"), job.get("progress_percent", 0))
            job["updated_at"] = _now()
            if isinstance(payload.get("details"), dict):
                job["details"] = payload["details"]
            self._save_state(state)
            return {"ok": True, "job": job, "status": self._public_status(state)}

    def complete(self, job_id: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        payload = payload or {}
        with self._lock:
            state = self._load_state()
            job = self._job_by_id(state, job_id)
            if job is None:
                raise ValueError(f"training job not found: {job_id}")
            success = bool(payload.get("success", payload.get("status") == "completed"))
            job["status"] = "completed" if success else "failed"
            job["phase"] = "completed" if success else "failed"
            job["progress_percent"] = 100 if success else _coerce_percent(payload.get("progress_percent"), job.get("progress_percent", 0))
            job["completed_at"] = _now()
            job["message"] = str(payload.get("message") or "")
            if isinstance(payload.get("summary"), dict):
                job["summary"] = payload["summary"]
            self._save_state(state)
            return {"ok": True, "job": job, "status": self._public_status(state)}

    def _load_state(self) -> dict[str, Any]:
        try:
            data = json.loads(self.state_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            data = {}
        if not isinstance(data, dict):
            data = {}
        data.setdefault("jobs", [])
        return data

    def _save_state(self, state: dict[str, Any]) -> None:
        self.runtime_dir.mkdir(parents=True, exist_ok=True)
        tmp = self.state_path.with_suffix(".tmp")
        tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        os.replace(tmp, self.state_path)

    def _worker_from_payload(self, payload: dict[str, Any]) -> dict[str, Any]:
        return {
            "id": str(payload.get("worker_id") or payload.get("id") or "windows-training-host"),
            "name": str(payload.get("name") or "DESKTOP-TMFDL0H"),
            "path": str(payload.get("path") or "C:\\Repos\\TradeBot"),
            "version": str(payload.get("version") or "1"),
            "last_seen_at": _now(),
        }

    def _public_status(self, state: dict[str, Any]) -> dict[str, Any]:
        worker = state.get("worker") if isinstance(state.get("worker"), dict) else {}
        last_seen = _parse_time(str(worker.get("last_seen_at") or ""))
        active = self._active_job(state)
        latest = _latest_job(state)
        recently_seen = bool(last_seen and datetime.now(UTC) - last_seen <= ONLINE_WINDOW)
        agent_busy = bool(
            active
            and active.get("status") == "running"
            and worker
            and active.get("claimed_by") == worker.get("id")
        )
        return {
            "available": True,
            "agent_online": recently_seen or agent_busy,
            "agent_recently_seen": recently_seen,
            "agent_busy": agent_busy,
            "worker": worker,
            "active_job": active,
            "latest_job": latest,
            "pending_jobs": sum(1 for job in state.get("jobs", []) if job.get("status") == "pending"),
        }

    def _active_job(self, state: dict[str, Any]) -> dict[str, Any] | None:
        for job in reversed(state.get("jobs", [])):
            if job.get("status") in {"pending", "running"}:
                return job
        return None

    def _oldest_pending_job(self, state: dict[str, Any]) -> dict[str, Any] | None:
        for job in state.get("jobs", []):
            if job.get("status") == "pending":
                return job
        return None

    def _job_by_id(self, state: dict[str, Any], job_id: str) -> dict[str, Any] | None:
        for job in state.get("jobs", []):
            if job.get("id") == job_id:
                return job
        return None

    def _expire_stale_jobs(self, state: dict[str, Any]) -> None:
        now = datetime.now(UTC)
        for job in state.get("jobs", []):
            if job.get("status") != "running":
                continue
            claimed_at = _parse_time(str(job.get("claimed_at") or ""))
            if claimed_at and now - claimed_at > RUNNING_TIMEOUT:
                job["status"] = "failed"
                job["completed_at"] = _now()
                job["message"] = "training worker timeout"

    def _trim_jobs(self, state: dict[str, Any]) -> None:
        jobs = state.get("jobs", [])
        if isinstance(jobs, list) and len(jobs) > 30:
            state["jobs"] = jobs[-30:]


def _safe_parameters(value: Any) -> dict[str, Any]:
    if not isinstance(value, dict):
        return {}
    allowed = {"symbols", "limit", "lookbacks", "architectures", "hidden_sizes", "layers", "dropouts", "epochs"}
    return {key: value[key] for key in allowed if key in value}


def _latest_job(state: dict[str, Any]) -> dict[str, Any] | None:
    jobs = state.get("jobs", [])
    if not jobs:
        return None
    latest = jobs[-1]
    return latest if isinstance(latest, dict) else None


def _now() -> str:
    return datetime.now(UTC).isoformat(timespec="seconds")


def _parse_time(value: str) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return None


def _coerce_percent(value: Any, default: Any = 0) -> int:
    try:
        number = int(float(value))
    except (TypeError, ValueError):
        try:
            number = int(float(default))
        except (TypeError, ValueError):
            number = 0
    return max(0, min(number, 100))


def _remove_tree(path: Path) -> None:
    if not path.exists():
        return
    for child in path.iterdir():
        if child.is_dir():
            _remove_tree(child)
        else:
            child.unlink(missing_ok=True)
    path.rmdir()
