from __future__ import annotations

import time
from typing import Any
from urllib.parse import urljoin

import httpx

from app.core.config import get_settings


def _join_url(base_url: str, path: str) -> str:
    return urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))


def _unwrap_items(payload: Any) -> list[dict[str, Any]]:
    if isinstance(payload, list):
        return [item for item in payload if isinstance(item, dict)]
    if isinstance(payload, dict):
        data = payload.get("data")
        if isinstance(data, dict) and isinstance(data.get("items"), list):
            return [item for item in data["items"] if isinstance(item, dict)]
        if isinstance(payload.get("items"), list):
            return [item for item in payload["items"] if isinstance(item, dict)]
        if isinstance(data, list):
            return [item for item in data if isinstance(item, dict)]
    return []


def _unwrap_dict(payload: Any) -> dict[str, Any]:
    if isinstance(payload, dict) and isinstance(payload.get("data"), dict):
        return payload["data"]
    return payload if isinstance(payload, dict) else {}


class ComputeNodeClient:
    """Application-side client for one compute node.

    The client accepts both current YG Compute API responses and common
    wrapper shapes such as `{code,message,data}` to make future engine/node
    adapters less brittle.
    """

    def __init__(self, api_base_url: str, token: str | None = None, timeout: float | None = None) -> None:
        settings = get_settings()
        self.api_base_url = api_base_url.rstrip("/")
        self.token = token or settings.compute_service_token
        self.timeout = timeout or settings.compute_request_timeout_seconds
        self.route_prefix = settings.route_prefix.rstrip("/") or "/modelTF"

    def headers(self) -> dict[str, str]:
        if not self.token:
            return {}
        return {"X-Compute-Token": self.token}

    async def test_connection(self) -> dict[str, Any]:
        started = time.perf_counter()
        health = await self.health()
        gpus = await self.gpus()
        return {
            "success": True,
            "latency_ms": int((time.perf_counter() - started) * 1000),
            "health": health,
            "gpus": gpus,
        }

    async def health(self) -> dict[str, Any]:
        paths = [f"{self.route_prefix}/v1/compute/health", f"{self.route_prefix}/health", "/health"]
        last_error = ""
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            for path in paths:
                try:
                    response = await client.get(_join_url(self.api_base_url, path))
                    response.raise_for_status()
                    return _unwrap_dict(response.json())
                except Exception as exc:  # noqa: BLE001 - keep endpoint compatibility fallback broad
                    last_error = str(exc)
        raise RuntimeError(last_error or "compute health check failed")

    async def gpus(self) -> list[dict[str, Any]]:
        paths = [
            f"{self.route_prefix}/compute/resources/gpus",
            f"{self.route_prefix}/v1/compute/resources/gpus",
            "/compute/resources/gpus",
        ]
        last_error = ""
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            for path in paths:
                try:
                    response = await client.get(_join_url(self.api_base_url, path))
                    response.raise_for_status()
                    return _unwrap_items(response.json())
                except Exception as exc:  # noqa: BLE001
                    last_error = str(exc)
        raise RuntimeError(last_error or "compute gpu discovery failed")

    async def create_job(self, payload: dict[str, Any]) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs"), json=payload)
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def preview_job(self, payload: dict[str, Any]) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.post(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/preview"),
                json=payload,
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def validate_job(self, payload: dict[str, Any]) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.post(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/validate"),
                json=payload,
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def check_paths(self, paths: list[dict[str, Any]]) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.post(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/check-paths"),
                json={"paths": paths},
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def list_files(
        self,
        root: str = "data",
        relative_path: str = "",
        directories_only: bool = False,
    ) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.get(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/list"),
                params={"root": root, "relative_path": relative_path, "directories_only": directories_only},
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def get_job(self, job_id: str) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.get(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}"))
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def stop_job(self, job_id: str) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/stop"))
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def job_logs(
        self,
        job_id: str,
        tail_lines: int | None = None,
        offset: int | None = None,
        limit: int | None = None,
    ) -> dict[str, Any]:
        params = {
            key: value
            for key, value in {"tail_lines": tail_lines, "offset": offset, "limit": limit}.items()
            if value is not None
        }
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.get(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/logs"),
                params=params,
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def import_local_file(self, payload: dict[str, Any]) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
            response = await client.post(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/import-local"),
                json=payload,
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())

    async def upload_file(
        self,
        filename: str,
        content: bytes,
        target_relative_path: str,
        resource_type: str | None = None,
        resource_id: str | None = None,
    ) -> dict[str, Any]:
        data = {
            "target_relative_path": target_relative_path,
            "resource_type": resource_type or "",
            "resource_id": resource_id or "",
        }
        files = {"file": (filename, content)}
        async with httpx.AsyncClient(timeout=max(self.timeout, 60), headers=self.headers()) as client:
            response = await client.post(
                _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/upload"),
                data=data,
                files=files,
            )
            response.raise_for_status()
            return _unwrap_dict(response.json())
