"""Templates resource — sync and async implementations."""
from __future__ import annotations
import time
from pathlib import Path
from typing import Any, BinaryIO, Dict, List, Optional, Union

from nanorix._constants import DEFAULT_UPLOAD_TIMEOUT
from nanorix._exceptions import NanorixError
from nanorix._http import AsyncHttpTransport, HttpTransport
from nanorix._models import Template, TemplateList


class Templates:
    """Synchronous template operations."""

    def __init__(self, http: HttpTransport) -> None:
        self._http = http

    def create(
        self,
        *,
        name: str,
        runtime: str,
        entrypoint: str,
        files: Union[List[str], Dict[str, bytes], Dict[str, BinaryIO]],
        description: Optional[str] = None,
        labels: Optional[Dict[str, str]] = None,
        env_vars: Optional[Dict[str, str]] = None,
        capsule_memory_mb: Optional[int] = None,
        input_path: Optional[str] = None,
        output_path: Optional[str] = None,
    ) -> Template:
        """
        Create a template with files. Handles the full three-step flow:
        1. POST /v1/templates — create template metadata
        2. POST /v1/templates/:id/files?filename= — upload each file (raw body)
        3. POST /v1/templates/:id/finalize — finalize template

        ``files`` accepts:
        - List[str]: file paths on disk (reads and uploads each)
        - Dict[str, bytes]: filename -> content bytes
        - Dict[str, BinaryIO]: filename -> file-like object
        """
        body: Dict[str, Any] = {"name": name, "runtime": runtime, "entrypoint": entrypoint}
        if description is not None:
            body["description"] = description
        if labels is not None:
            body["labels"] = labels
        if env_vars is not None:
            body["env_vars"] = env_vars
        if capsule_memory_mb is not None:
            body["capsule_memory_mb"] = capsule_memory_mb
        if input_path is not None:
            body["input_path"] = input_path
        if output_path is not None:
            body["output_path"] = output_path

        result = self._http.request("POST", "/v1/templates", json=body)
        template_id = result["id"]

        # Step 2: Upload each file (raw body, not multipart)
        file_map = self._resolve_files(files)
        for filename, content in file_map.items():
            self._http.request(
                "POST",
                f"/v1/templates/{template_id}/files",
                data=content,
                headers={"Content-Type": "application/octet-stream"},
                params={"filename": filename},
                timeout=DEFAULT_UPLOAD_TIMEOUT,
            )

        # Step 3: Finalize
        result = self._http.request("POST", f"/v1/templates/{template_id}/finalize")
        return Template.model_validate(result)

    def list(self, **kwargs: Any) -> TemplateList:
        """List templates. Supports ?status=ready, ?cursor=, ?limit= query params."""
        result = self._http.request("GET", "/v1/templates", params=kwargs or None)
        return TemplateList.model_validate(result)

    def get(self, template_id: str) -> Template:
        """Get a template by ID."""
        result = self._http.request("GET", f"/v1/templates/{template_id}")
        return Template.model_validate(result)

    def delete(self, template_id: str) -> None:
        """Delete a template (crypto-shreds encryption key, deletes files)."""
        self._http.request("DELETE", f"/v1/templates/{template_id}")

    def update(self, template_id: str, **kwargs: Any) -> Template:
        """Update a template. Not yet implemented."""
        raise NotImplementedError(
            "PUT /v1/templates/:id not yet implemented. Coming in v1.1."
        )

    def wait_until_ready(
        self,
        template_id: str,
        *,
        timeout: float = 120.0,
        poll_interval: float = 2.0,
    ) -> Template:
        """Poll template status until 'ready' or timeout."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            tmpl = self.get(template_id)
            if tmpl.status == "ready":
                return tmpl
            if tmpl.status == "failed":
                raise NanorixError(
                    f"Template build failed: {tmpl.error_message}",
                    status_code=None,
                )
            time.sleep(poll_interval)
        raise NanorixError(
            f"Template {template_id} not ready after {timeout}s",
            status_code=None,
        )

    @staticmethod
    def _resolve_files(
        files: Union[List[str], Dict[str, bytes], Dict[str, BinaryIO]],
    ) -> Dict[str, bytes]:
        if isinstance(files, list):
            result: Dict[str, bytes] = {}
            for path_str in files:
                p = Path(path_str)
                result[p.name] = p.read_bytes()
            return result
        elif isinstance(files, dict):
            result = {}
            for name, content in files.items():
                if isinstance(content, bytes):
                    result[name] = content
                else:
                    result[name] = content.read()
            return result
        raise TypeError(
            f"files must be list, dict[str,bytes], or dict[str,BinaryIO], got {type(files)}"
        )


class AsyncTemplates:
    """Asynchronous template operations."""

    def __init__(self, http: AsyncHttpTransport) -> None:
        self._http = http

    async def create(
        self,
        *,
        name: str,
        runtime: str,
        entrypoint: str,
        files: Union[List[str], Dict[str, bytes], Dict[str, BinaryIO]],
        description: Optional[str] = None,
        labels: Optional[Dict[str, str]] = None,
        env_vars: Optional[Dict[str, str]] = None,
        capsule_memory_mb: Optional[int] = None,
        input_path: Optional[str] = None,
        output_path: Optional[str] = None,
    ) -> Template:
        """Create a template with files (async). Same three-step flow as sync."""
        body: Dict[str, Any] = {"name": name, "runtime": runtime, "entrypoint": entrypoint}
        if description is not None:
            body["description"] = description
        if labels is not None:
            body["labels"] = labels
        if env_vars is not None:
            body["env_vars"] = env_vars
        if capsule_memory_mb is not None:
            body["capsule_memory_mb"] = capsule_memory_mb
        if input_path is not None:
            body["input_path"] = input_path
        if output_path is not None:
            body["output_path"] = output_path

        result = await self._http.request("POST", "/v1/templates", json=body)
        template_id = result["id"]

        file_map = Templates._resolve_files(files)
        for filename, content in file_map.items():
            await self._http.request(
                "POST",
                f"/v1/templates/{template_id}/files",
                data=content,
                headers={"Content-Type": "application/octet-stream"},
                params={"filename": filename},
                timeout=DEFAULT_UPLOAD_TIMEOUT,
            )

        result = await self._http.request("POST", f"/v1/templates/{template_id}/finalize")
        return Template.model_validate(result)

    async def list(self, **kwargs: Any) -> TemplateList:
        """List templates (async)."""
        result = await self._http.request("GET", "/v1/templates", params=kwargs or None)
        return TemplateList.model_validate(result)

    async def get(self, template_id: str) -> Template:
        """Get a template by ID (async)."""
        result = await self._http.request("GET", f"/v1/templates/{template_id}")
        return Template.model_validate(result)

    async def delete(self, template_id: str) -> None:
        """Delete a template (async)."""
        await self._http.request("DELETE", f"/v1/templates/{template_id}")

    async def update(self, template_id: str, **kwargs: Any) -> Template:
        """Update a template. Not yet implemented."""
        raise NotImplementedError(
            "PUT /v1/templates/:id not yet implemented. Coming in v1.1."
        )

    async def wait_until_ready(
        self,
        template_id: str,
        *,
        timeout: float = 120.0,
        poll_interval: float = 2.0,
    ) -> Template:
        """Poll template status until 'ready' or timeout (async)."""
        import asyncio
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            tmpl = await self.get(template_id)
            if tmpl.status == "ready":
                return tmpl
            if tmpl.status == "failed":
                raise NanorixError(
                    f"Template build failed: {tmpl.error_message}",
                    status_code=None,
                )
            await asyncio.sleep(poll_interval)
        raise NanorixError(
            f"Template {template_id} not ready after {timeout}s",
            status_code=None,
        )
