"""Batch orchestration for submitting and polling many energy models.

``BatchRunner`` is a client-side coordination helper — it does not map to an
API endpoint.  It wraps ``client.energy_models.create()`` for submission of
accepted envelopes and ``client.tasks.get()`` for lightweight status polling,
then optionally fetches full results via ``client.energy_models.get()``.

Usage::

    from dalysdk import DalyClient, BatchRunner

    with DalyClient(workspace_api_key="wk_...", user_api_key="uk_...") as client:
        batch = BatchRunner(client)

        batch.submit("Site A - Tracker", epm_input_a)
        batch.submit("Site B - Fixed", epm_input_b)

        batch.wait(timeout=3600, poll_interval=15)

        print(batch.summary())

        for result in batch.results():
            print(result["name"], result["data"].get("annualProductionMWh"))
"""

from __future__ import annotations

import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Callable

    from dalysdk.client import DalyClient

log = logging.getLogger(__name__)

_TERMINAL_STATUSES = {"COMPLETED", "FAILED", "CANCELLED"}


class BatchRunner:
    """Submit and poll many async energy model runs against a :class:`DalyClient`.

    Args:
        client: An initialised :class:`DalyClient` instance.
        max_workers: Maximum threads when using ``parallel=True`` in
            :meth:`submit_many`.  Defaults to ``4``.
    """

    def __init__(self, client: DalyClient, *, max_workers: int = 4) -> None:
        self._client = client
        self._max_workers = max_workers
        self._runs: list[dict[str, Any]] = []
        self._lock = threading.Lock()

    # -- Submission ---------------------------------------------------------

    def submit(
        self,
        name: str,
        data: dict[str, Any],
        *,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        """Submit a single energy model run and track its accepted envelope.

        Args:
            name: Human-readable label for this run.
            data: Full energy model input dict.
            timeout: Optional request timeout override (seconds).

        Returns:
            The raw API response dict containing task and model identifiers.
        """
        kwargs: dict[str, Any] = {}
        if timeout is not None:
            kwargs["timeout"] = timeout

        response = self._client.energy_models.create(data, **kwargs)
        task_id, energy_model_id, run_index = _extract_submission_ids(response)

        run: dict[str, Any] = {
            "name": name,
            "task_id": task_id,
            "energy_model_id": energy_model_id,
            "run_index": run_index,
            "status": "PENDING",
            "response": response,
            "error": None,
        }

        with self._lock:
            self._runs.append(run)

        log.info("Submitted '%s' — task_id=%s", name, run["task_id"])
        return response

    def submit_many(
        self,
        configs: list[tuple[str, dict[str, Any]]],
        *,
        parallel: bool = False,
        timeout: float | None = None,
    ) -> int:
        """Submit multiple energy model runs.

        Args:
            configs: List of ``(name, epm_input_dict)`` tuples.
            parallel: If ``True``, submit concurrently using a thread pool.
            timeout: Optional per-request timeout override (seconds).

        Returns:
            Number of successfully submitted runs.
        """
        if not parallel:
            success = 0
            for name, data in configs:
                try:
                    self.submit(name, data, timeout=timeout)
                    success += 1
                except Exception:
                    log.exception("Failed to submit '%s'", name)
            return success

        success = 0
        with ThreadPoolExecutor(max_workers=self._max_workers) as pool:
            futures = {pool.submit(self.submit, name, data, timeout=timeout): name for name, data in configs}
            for fut in as_completed(futures):
                name = futures[fut]
                try:
                    fut.result()
                    success += 1
                except Exception:
                    log.exception("Failed to submit '%s'", name)
        return success

    # -- Polling ------------------------------------------------------------

    def poll(self) -> tuple[int, int, int]:
        """Poll task statuses for all non-terminal runs.

        Returns:
            Tuple of ``(completed, failed, pending)`` counts across *all*
            tracked runs (not just those polled this cycle).
        """
        for run in self._runs:
            if run["status"] in _TERMINAL_STATUSES:
                continue
            task_id = run["task_id"]
            if task_id is None:
                continue
            try:
                task = self._client.tasks.get(task_id)
                new_status = task.get("status", run["status"])
                if isinstance(new_status, str):
                    new_status = new_status.upper()
                run["status"] = new_status
                if new_status == "FAILED":
                    task_details = task.get("taskDetails")
                    if not isinstance(task_details, dict):
                        task_details = {}
                    run["error"] = (
                        task.get("error")
                        or task.get("errorMessage")
                        or task.get("detail")
                        or task_details.get("error")
                        or task_details.get("errorMessage")
                    )
            except Exception:
                log.exception("Error polling task %s for '%s'", task_id, run["name"])

        return (
            len(self.completed),
            len(self.failed),
            len(self.pending),
        )

    def wait(
        self,
        *,
        timeout: float | None = None,
        poll_interval: float = 15.0,
        on_poll: Callable[[int, int, int], Any] | None = None,
    ) -> bool:
        """Block until all runs reach a terminal status or *timeout* elapses.

        Args:
            timeout: Maximum seconds to wait.  ``None`` means wait forever.
            poll_interval: Seconds between poll cycles.
            on_poll: Optional callback invoked after each poll with
                ``(completed, failed, pending)``.

        Returns:
            ``True`` if all runs are terminal, ``False`` if timed out.
        """
        deadline = (time.monotonic() + timeout) if timeout is not None else None

        while True:
            completed, failed, pending = self.poll()
            log.info(
                "Poll: %d completed, %d failed, %d pending",
                completed,
                failed,
                pending,
            )
            if on_poll is not None:
                on_poll(completed, failed, pending)

            if pending == 0:
                return True

            if deadline is not None and time.monotonic() >= deadline:
                log.warning("Timeout reached with %d runs still pending", pending)
                return False

            time.sleep(poll_interval)

    # -- Results / summary --------------------------------------------------

    def summary(self) -> dict[str, Any]:
        """Return a summary dict of current batch state.

        Returns:
            Dict with keys ``total``, ``completed``, ``failed``, ``pending``,
            ``results`` (completed run names/IDs), and ``failures`` (failed
            run names/errors).
        """
        return {
            "total": len(self._runs),
            "completed": len(self.completed),
            "failed": len(self.failed),
            "pending": len(self.pending),
            "results": [
                {
                    "name": r["name"],
                    "energy_model_id": r["energy_model_id"],
                    "task_id": r["task_id"],
                }
                for r in self.completed
            ],
            "failures": [
                {
                    "name": r["name"],
                    "task_id": r["task_id"],
                    "error": r["error"],
                }
                for r in self.failed
            ],
        }

    def results(self) -> list[dict[str, Any]]:
        """Fetch full energy model data for every completed run.

        Calls ``client.energy_models.get()`` for each completed run and
        returns a list of dicts with ``name`` and ``data`` keys.

        Returns:
            List of ``{"name": ..., "energy_model_id": ..., "data": ...}``
            dicts.
        """
        out: list[dict[str, Any]] = []
        for run in self.completed:
            em_id = run["energy_model_id"]
            if em_id is None:
                continue
            try:
                data = self._client.energy_models.get(em_id)
                out.append(
                    {
                        "name": run["name"],
                        "energy_model_id": em_id,
                        "data": data,
                    }
                )
            except Exception:
                log.exception("Failed to fetch results for '%s' (id=%s)", run["name"], em_id)
        return out

    # -- Filtered views -----------------------------------------------------

    @property
    def runs(self) -> list[dict[str, Any]]:
        """All tracked runs."""
        return list(self._runs)

    @property
    def completed(self) -> list[dict[str, Any]]:
        """Runs with status ``COMPLETED``."""
        return [r for r in self._runs if r["status"] == "COMPLETED"]

    @property
    def failed(self) -> list[dict[str, Any]]:
        """Runs with status ``FAILED``."""
        return [r for r in self._runs if r["status"] == "FAILED"]

    @property
    def pending(self) -> list[dict[str, Any]]:
        """Runs that have not yet reached a terminal status."""
        return [r for r in self._runs if r["status"] not in _TERMINAL_STATUSES]


def _extract_submission_ids(response: dict[str, Any]) -> tuple[Any, Any, Any]:
    task_details = response.get("taskDetails")
    if not isinstance(task_details, dict):
        task_details = {}

    task_id = response.get("taskId")
    if task_id is None:
        task_id = response.get("task_id")
    if task_id is None:
        task_id = task_details.get("taskId")
    if task_id is None:
        task_id = task_details.get("task_id")

    energy_model_id = response.get("energyModelId")
    if energy_model_id is None:
        energy_model_id = response.get("energy_model_id")
    if energy_model_id is None:
        energy_model_id = task_details.get("energyModelId")
    if energy_model_id is None:
        energy_model_id = task_details.get("energy_model_id")

    run_index = response.get("runIndex")
    if run_index is None:
        run_index = response.get("run_index")
    if run_index is None:
        run_index = task_details.get("runIndex")
    if run_index is None:
        run_index = task_details.get("run_index")

    return task_id, energy_model_id, run_index
