from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

from app.db.platform_store import get_platform_store
from app.core.config import get_settings
from app.modules.compute_gateway.client import ComputeNodeClient
from app.modules.storage.minio_store import get_object_storage

# starting 状态允许的最大轮询次数（约 40 * 3s ≈ 2 分钟），超过即判定节点不可达
MAX_STARTING_ATTEMPTS = 40


def _extract_job_failure_reason(log_text: str, limit: int = 2000) -> str:
    """Return a concise actionable reason from a failed Compute job log."""
    lines = [line.strip() for line in str(log_text or "").splitlines() if line.strip()]
    if not lines:
        return ""
    markers = ("[eval] FAILED", "Traceback", "RuntimeError", "Error:", "ERROR")
    for index in range(len(lines) - 1, -1, -1):
        if any(marker in lines[index] for marker in markers):
            return "\n".join(lines[index : index + 8])[-limit:]
    return "\n".join(lines[-8:])[-limit:]


async def _archive_node_directory(
    store: Any,
    client: ComputeNodeClient,
    node: dict[str, Any],
    source_path: str,
    resource_type: str,
    resource_id: str,
    version_id: str,
    object_prefix: str,
) -> list[dict[str, Any]]:
    """Archive a completed node directory to MinIO, preserving subdirectories."""
    data_root = Path(str(node.get("data_root") or "/data/yg-ft")).resolve()
    source = Path(source_path).resolve()
    if source == data_root:
        raise RuntimeError("refuse to archive compute data root; output_dir must be a task subdirectory")
    try:
        relative_root = source.relative_to(data_root).as_posix()
    except ValueError as exc:
        raise RuntimeError(f"artifact path is outside compute data root: {source_path}") from exc
    queue = [relative_root]
    archived: list[dict[str, Any]] = []
    max_files = 10000
    while queue:
        relative = queue.pop(0)
        listing = await client.list_files(root="data", relative_path=relative)
        for item in listing.get("items") or []:
            item_relative = str(item.get("relative_path") or "")
            if item.get("type") == "directory":
                queue.append(item_relative)
                continue
            path = str(item.get("path") or "")
            if not path:
                continue
            try:
                relative_file = Path(item_relative).relative_to(Path(relative_root)).as_posix()
            except ValueError:
                relative_file = Path(str(item.get("name") or Path(path).name)).name
            object_key = f"{object_prefix}/{version_id}/{relative_file}"
            if len(archived) >= max_files:
                raise RuntimeError(f"archive file count exceeds limit {max_files}")
            upload_url = get_object_storage().presigned_put(object_key)
            result = await client.upload_file_to_url(path, upload_url, object_key)
            metadata = get_object_storage().stat(object_key)
            archived.append(
                store.create_storage_object(
                    {
                        "resource_type": resource_type,
                        "resource_id": resource_id,
                        "version_id": version_id,
                        "bucket": get_object_storage().bucket,
                        "object_key": object_key,
                        "file_name": relative_file,
                        "content_type": "application/octet-stream",
                        "byte_size": metadata.get("byte_size") or result.get("byte_size") or 0,
                        "checksum_sha256": result.get("checksum_sha256") or "",
                        "status": "available",
                    }
                )
            )
    return archived


def _node_for_task(task: dict[str, Any]) -> dict[str, Any] | None:
    return next((node for node in get_platform_store().compute_nodes() if node["id"] == task.get("compute_node_id")), None)


def _parse_inference_load_status(task: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    load_status = task.get("load_status") or {}
    if isinstance(load_status, str):
        try:
            load_status = json.loads(load_status)
        except (json.JSONDecodeError, TypeError):
            load_status = {}
    return load_status.get("loaded_models") or [], load_status


async def reconcile_inference_loads(store: Any) -> list[dict[str, Any]]:
    """推进处于 starting 状态的推理加载。

    模型加载已改为异步派发：/model-compare/{id}/load 立即返回，这里在每次
    轮询时查询对应计算节点的 /inference/status，把任务从 starting 推进到
    ready/error。使用短超时，单节点不可达不会阻塞整轮轮询。
    """
    reconciled: list[dict[str, Any]] = []
    now = time.time()
    for task in store.compare_tasks():
        items, _ = _parse_inference_load_status(task)
        if not any(item.get("status") == "starting" for item in items):
            continue
        # dirty 只要处理过任一 starting 项就置位：load_attempts / last_polled_at
        # 必须落库，否则节点不可达时计数不会累积，封顶逻辑永远触发不了
        dirty = False
        for item in items:
            if item.get("status") != "starting":
                continue
            # 节流：同一 item 每 3s 只查询一次
            if now - float(item.get("last_polled_at") or 0) < 3:
                continue
            item["last_polled_at"] = now
            item["load_attempts"] = int(item.get("load_attempts") or 0) + 1
            dirty = True
            node = next((n for n in store.compute_nodes() if n["id"] == item.get("node_id")), None)
            if not node:
                item["status"] = "error"
                item["error"] = "compute node deleted"
                store.mark_inference_unloaded(item.get("node_id") or "")
                store.release_external_gpus("inference", str(task["id"]), item.get("node_id"))
                continue
            if not node.get("enabled") or node.get("scheduler_status") != "online":
                item["status"] = "error"
                item["error"] = "compute node offline"
                store.mark_inference_unloaded(node["id"])
                store.release_external_gpus("inference", str(task["id"]), node["id"])
                continue
            try:
                status = await ComputeNodeClient(node["api_base_url"]).inference_status()
            except Exception as exc:  # noqa: BLE001 - node unreachable; keep retrying until cap
                if int(item.get("load_attempts") or 0) >= MAX_STARTING_ATTEMPTS:
                    item["status"] = "error"
                    item["error"] = f"compute node unreachable: {exc}"
                    store.mark_inference_unloaded(node["id"])
                    store.release_external_gpus("inference", str(task["id"]), node["id"])
                continue
            node_status = status.get("status")
            if node_status == "ready":
                item["status"] = "ready"
                item.pop("error", None)
                selected_gpus = item.get("gpu_indices") or item.get("gpus")
                store.mark_inference_loaded(node["id"], selected_gpus)
            elif node_status == "error":
                item["status"] = "error"
                item["error"] = status.get("error") or "model load failed on compute node"
                store.mark_inference_unloaded(node["id"])
                store.release_external_gpus("inference", str(task["id"]), node["id"])
            elif node_status == "idle":
                # 节点重启导致已加载模型丢失
                item["status"] = "error"
                item["error"] = "model disappeared from compute node (node may have restarted)"
                store.mark_inference_unloaded(node["id"])
                store.release_external_gpus("inference", str(task["id"]), node["id"])
            # node_status == "loading" -> 保持 starting，下轮再查
        if dirty:
            if any(i.get("status") in {"ready", "running"} for i in items):
                new_status = "loaded"
            elif any(i.get("status") == "starting" for i in items):
                new_status = "starting"  # 仍在加载中，保持 starting
            else:
                new_status = "failed"
            store.update_compare_task(task["id"], {"status": new_status, "load_status": {"loaded_models": items}})
            reconciled.append({"task_id": task["id"], "status": new_status})
    return reconciled


async def fetch_eval_result_content(
    client: ComputeNodeClient,
    node: dict[str, Any],
    job: dict[str, Any],
    file_name: str = "eval_results.json",
) -> dict[str, Any] | None:
    output_dir = job.get("output_dir")
    if not output_dir:
        return None
    full_path = f"{str(output_dir).rstrip('/')}/{file_name}"
    data_root = "/data/yg-ft/"
    if full_path.startswith(data_root):
        full_path = full_path[len(data_root):]
    rel_path = full_path.lstrip("/")
    import httpx
    url = f"{node['api_base_url'].rstrip('/')}/modelTF/compute/files/read"
    async with httpx.AsyncClient(timeout=30, headers=client.headers()) as http:
        response = await http.get(url, params={"path": rel_path})
        response.raise_for_status()
        payload = response.json()
        return payload if isinstance(payload, dict) else None


async def fetch_eval_progress_content(
    client: ComputeNodeClient,
    node: dict[str, Any],
    job: dict[str, Any],
) -> dict[str, Any] | None:
    return await fetch_eval_result_content(client, node, job, "eval_progress.json")


async def poll_compute_jobs_once(store: Any | None = None) -> dict[str, Any]:
    store = store or get_platform_store()
    synced: list[dict[str, Any]] = []
    failed: list[dict[str, str]] = []
    online_nodes = {
        str(node.get("id"))
        for node in store.compute_nodes()
        if node.get("enabled") and node.get("scheduler_status") in {"online", "draining"}
    }
    training_tasks = {str(task["id"]): task for task in store.running_compute_tasks()}
    # Completed tasks whose MinIO archive was interrupted remain eligible for
    # reconciliation after a Backend restart or a transient node failure.
    if get_settings().minio_enabled:
        for task in store.tasks():
            if task.get("status") != "completed" or not task.get("compute_job_id"):
                continue
            if (
                str(task.get("archive_status") or "") != "completed"
                and str(task.get("compute_node_id")) in online_nodes
            ):
                training_tasks.setdefault(str(task["id"]), task)
    for task in training_tasks.values():
        node = _node_for_task(task)
        if not node:
            failed.append({"task_id": task["id"], "error": "compute node not found"})
            continue
        try:
            client = ComputeNodeClient(node["api_base_url"])
            job = await client.get_job(task["compute_job_id"])
            try:
                logs = await client.job_logs(task["compute_job_id"], tail_lines=5000)
                store.record_training_log_metrics(task["id"], str(logs.get("content") or ""))
            except Exception:
                pass
            # P0-4: Force-fetch last log snippet when job reaches terminal state
            if job.get("status") in {"failed", "stopped"}:
                try:
                    last_logs = await client.job_logs(task["compute_job_id"], tail_lines=200)
                    job["log_snippet"] = str(last_logs.get("content") or "")[:8192]
                except Exception:
                    pass
            updated_task = store.apply_compute_job(task["id"], job)
            if (
                get_settings().minio_enabled
                and job.get("status") == "completed"
                and job.get("output_dir")
            ):
                trained_model = next(
                    (
                        item
                        for item in store.trained_models()
                        if item.get("name")
                        == (task.get("output_model_name") or f"{task.get('name')}-lora")
                    ),
                    None,
                )
                if trained_model:
                    try:
                        archived = await _archive_node_directory(
                            store,
                            client,
                            node,
                            str(job["output_dir"]),
                            "trained_model",
                            str(trained_model["id"]),
                            str(job.get("id") or task.get("compute_job_id") or task["id"]),
                            f"trained_models/{trained_model['id']}",
                        )
                        artifacts = store.model_artifacts(str(trained_model["id"]))
                        if archived and artifacts:
                            store.link_model_artifact_storage_object(
                                str(artifacts[0]["id"]), str(archived[0]["id"])
                            )
                        store.update_task(task["id"], {
                            "archive_status": "completed",
                            "archive_object_ids": [str(item["id"]) for item in archived],
                            "archive_error": "",
                        })
                    except Exception as archive_exc:
                        store.update_task(task["id"], {
                            "archive_status": "pending",
                            "archive_error": str(archive_exc)[:2000],
                        })
                        raise
            synced.append(updated_task)
        except Exception as exc:  # noqa: BLE001 - keep polling other jobs
            failed.append({"task_id": task["id"], "error": str(exc)})
    standalone_synced: list[dict[str, Any]] = []
    standalone_jobs = {
        str(record["id"]): record
        for record in store.active_standalone_compute_jobs()
        if str(record.get("node_id")) in online_nodes
    }
    if get_settings().minio_enabled:
        for record in store.standalone_compute_jobs_pending_archive():
            if str(record.get("node_id")) in online_nodes:
                standalone_jobs.setdefault(str(record["id"]), record)
    for record in standalone_jobs.values():
        node = next((item for item in store.compute_nodes() if item["id"] == record.get("node_id")), None)
        if not node:
            failed.append({"job_id": record["id"], "error": "compute node not found"})
            continue
        try:
            job = await ComputeNodeClient(node["api_base_url"]).get_job(record["id"])
            standalone_synced.append(store.sync_model_merge_job(record["id"], job))
            if get_settings().minio_enabled and job.get("status") == "completed" and job.get("output_dir"):
                payload = (store.compute_job(record["id"]).get("payload") or {})
                trained_model_id = str(payload.get("trained_model_id") or payload.get("model_name") or "")
                if trained_model_id:
                    trained_model = next(
                        (item for item in store.trained_models() if item.get("id") == trained_model_id or item.get("name") == trained_model_id),
                        None,
                    )
                    if trained_model:
                        archived = await _archive_node_directory(
                            store,
                            ComputeNodeClient(node["api_base_url"], timeout=900),
                            node,
                            str(job["output_dir"]),
                            "trained_model",
                            str(trained_model["id"]),
                            str(job.get("id") or record["id"]),
                            f"trained_models/{trained_model['id']}",
                        )
                        artifacts = store.model_artifacts(str(trained_model["id"]))
                        if archived and artifacts:
                            store.link_model_artifact_storage_object(
                                str(artifacts[0]["id"]), str(archived[0]["id"])
                            )
                        store.update_compute_job_archive(record["id"], "completed", [str(item["id"]) for item in archived])
        except Exception as exc:  # noqa: BLE001 - keep polling other jobs
            try:
                store.update_compute_job_archive(record["id"], "pending", [], str(exc)[:2000])
            except Exception:
                pass
            failed.append({"job_id": record["id"], "error": str(exc)})

    # ── Eval job sync ────────────────────────────────────────────────
    eval_synced = 0
    eval_tasks = {str(task["id"]): task for task in store.running_eval_tasks()}
    if get_settings().minio_enabled:
        # A completed evaluation can win the race with the poller: its status
        # is persisted before the report archive finishes. Keep such tasks in
        # the reconciliation set until the report object is available.
        for task in store.eval_tasks():
            if (
                task.get("status") == "completed"
                and task.get("compute_job_id")
                and str(task.get("archive_status") or "") != "completed"
                and str(task.get("compute_node_id")) in online_nodes
            ):
                eval_tasks.setdefault(str(task["id"]), task)
    for eval_task in eval_tasks.values():
        if str(eval_task.get("compute_node_id")) not in online_nodes:
            continue
        node = next(
            (item for item in store.compute_nodes() if item["id"] == eval_task.get("compute_node_id")),
            None,
        )
        if not node:
            failed.append({"eval_task_id": eval_task["id"], "error": "compute node not found"})
            continue
        try:
            client = ComputeNodeClient(node["api_base_url"])
            job = await client.get_job(eval_task["compute_job_id"])
            result_content = None
            # Read live progress and partial results while the evaluator is running.
            if job.get("status") in {"queued", "running"} and job.get("output_dir"):
                try:
                    progress_content = await fetch_eval_progress_content(client, node, job)
                    if progress_content:
                        store.update_eval_task(
                            eval_task["id"],
                            {
                                "progress_detail": progress_content,
                                "progress": progress_content.get("percentage", eval_task.get("progress", 0)),
                            },
                        )
                except Exception:
                    pass
                try:
                    result_content = await fetch_eval_result_content(client, node, job)
                except Exception:
                    result_content = None
            # Try to read eval_results.json from the job output directory on completion.
            if job.get("status") == "completed" and job.get("output_dir"):
                try:
                    result_content = await fetch_eval_result_content(client, node, job)
                except Exception:
                    pass
            if job.get("status") in {"failed", "stopped"} and not job.get("error"):
                try:
                    failure_logs = await client.job_logs(eval_task["compute_job_id"], tail_lines=120)
                    job["error"] = _extract_job_failure_reason(str(failure_logs.get("content") or ""))
                except Exception:
                    pass
            store.apply_eval_job_result(eval_task["id"], job, result_content)
            if get_settings().minio_enabled and job.get("status") == "completed" and job.get("output_dir"):
                archived = await _archive_node_directory(
                    store,
                    client,
                    node,
                    str(job["output_dir"]),
                    "eval",
                    str(eval_task["id"]),
                    str(job.get("id") or eval_task.get("compute_job_id") or eval_task["id"]),
                    f"evaluations/{eval_task['id']}",
                )
                report_object = next(
                    (item for item in archived if Path(str(item.get("file_name") or "")).name == "eval_results.json"),
                    archived[0] if archived else None,
                )
                store.update_eval_task(eval_task["id"], {
                    "report_storage_object_id": str(report_object["id"]) if report_object else "",
                    "archive_status": "completed",
                    "archive_object_ids": [str(item["id"]) for item in archived],
                    "archive_error": "",
                })
            # 评测 GPU 占用由 eval_tasks 状态派生，无需维护推理内存标记
            eval_synced += 1
        except Exception as exc:  # noqa: BLE001
            try:
                current_eval = store.eval_task(eval_task["id"])
            except Exception:
                current_eval = eval_task
            if current_eval.get("status") == "completed":
                try:
                    store.update_eval_task(eval_task["id"], {"archive_status": "pending", "archive_error": str(exc)[:2000]})
                except Exception:
                    pass
            failed.append({"eval_task_id": eval_task["id"], "error": str(exc)})

    # ── Inference load reconciliation ─────────────────────────────────────
    try:
        inference_reconciled = await reconcile_inference_loads(store)
    except Exception as exc:  # noqa: BLE001 - keep polling alive
        failed.append({"inference_reconcile": str(exc)})
        inference_reconciled = []

    return {"synced": len(synced) + len(standalone_synced) + eval_synced, "failed": failed,
            "items": synced, "standalone": standalone_synced, "eval_synced": eval_synced,
            "inference_reconciled": inference_reconciled}