from __future__ import annotations

import hashlib
import json
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol


@dataclass
class BackendResult:
    status: str
    checkpoint_reference: str = ""
    checkpoint_hash: str = ""
    metrics: dict[str, float] | None = None
    stdout: str = ""
    stderr: str = ""
    failure_category: str = ""
    failure_details: str = ""


class TrainingBackend(Protocol):
    def estimate_runtime(self, recipe: dict[str, Any]) -> int: ...
    def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult: ...
    def validate_checkpoint(self, reference: str) -> bool: ...


class FakeTrainingBackend:
    """Deterministic backend for workflow tests; never launches training."""

    def __init__(self, outcomes: list[dict[str, Any]] | None = None) -> None:
        self.outcomes = list(outcomes or [{"status": "SUCCEEDED", "metrics": {"primary": 0.75}}])

    def estimate_runtime(self, recipe: dict[str, Any]) -> int:
        return int(recipe.get("estimated_runtime_seconds", 60))

    def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult:
        outcome = self.outcomes.pop(0) if self.outcomes else {"status": "SUCCEEDED", "metrics": {"primary": 0.75}}
        status = str(outcome.get("status", "SUCCEEDED"))
        reference = str(outcome.get("checkpoint_reference", f"fake://{hashlib.sha256(json.dumps(outcome, sort_keys=True).encode()).hexdigest()[:16]}"))
        return BackendResult(status=status, checkpoint_reference=reference, checkpoint_hash=hashlib.sha256(reference.encode()).hexdigest(), metrics=outcome.get("metrics", {}), failure_category=str(outcome.get("failure_category", "")), failure_details=str(outcome.get("failure_details", "")))

    def validate_checkpoint(self, reference: str) -> bool:
        return reference.startswith("fake://")


class GuardSubprocessBackend:
    """Scoped wrapper for the discovered Guard trainer; commands come from profiles, never an LLM."""

    def estimate_runtime(self, recipe: dict[str, Any]) -> int:
        return int(recipe.get("estimated_runtime_seconds", 90 * 60))

    def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult:
        try:
            completed = subprocess.run(command, cwd=working_directory, capture_output=True, text=True, timeout=timeout_seconds, check=False)
        except subprocess.TimeoutExpired as exc:
            return BackendResult(status="TIMEOUT", stdout=exc.stdout or "", stderr=exc.stderr or "", failure_category="TIMEOUT", failure_details=f"Exceeded {timeout_seconds}s")
        stdout, stderr = completed.stdout or "", completed.stderr or ""
        if completed.returncode:
            category = "OOM" if "out of memory" in (stdout + stderr).lower() else "PROCESS_FAILURE"
            return BackendResult(status=category if category == "OOM" else "FAILED", stdout=stdout, stderr=stderr, failure_category=category, failure_details=f"Exit code {completed.returncode}")
        return BackendResult(status="SUCCEEDED", stdout=stdout, stderr=stderr)

    def validate_checkpoint(self, reference: str) -> bool:
        path = Path(reference)
        return path.is_dir() and any(path.glob("adapter_model.*"))


class SparkGuardBackend(GuardSubprocessBackend):
    """Runs only profile-generated Guard commands over the configured Spark SSH alias."""

    def __init__(self, ssh_alias: str = "spark") -> None:
        self.ssh_alias = ssh_alias

    def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult:
        if not command:
            return BackendResult(status="FAILED", failure_category="COMMAND_SCOPE", failure_details="Missing profile-generated command.")
        import shlex

        remote = "cd " + shlex.quote(working_directory) + " && " + " ".join(shlex.quote(part) for part in command)
        try:
            completed = subprocess.run(["ssh", self.ssh_alias, remote], capture_output=True, text=True, timeout=timeout_seconds, check=False)
        except subprocess.TimeoutExpired as exc:
            return BackendResult(status="TIMEOUT", stdout=exc.stdout or "", stderr=exc.stderr or "", failure_category="TIMEOUT", failure_details=f"Exceeded {timeout_seconds}s")
        if completed.returncode:
            combined = (completed.stdout or "") + (completed.stderr or "")
            category = "OOM" if "out of memory" in combined.lower() else "REMOTE_PROCESS_FAILURE"
            return BackendResult(status="OOM" if category == "OOM" else "FAILED", stdout=completed.stdout or "", stderr=completed.stderr or "", failure_category=category, failure_details=f"Exit code {completed.returncode}")
        return BackendResult(status="SUCCEEDED", stdout=completed.stdout or "", stderr=completed.stderr or "")