"""Pipeline DSL for composing module steps.

Provides a simple, declarative way to define multi-step analysis workflows
where each step is a module invocation.

Usage::

    from hla_compass.pipeline import Pipeline, Step

    pipeline = Pipeline(
        name="my-analysis",
        steps=[
            Step("preprocessing-module", params={"threshold": 0.5}),
            Step(
                "analysis-module",
                depends_on="preprocessing-module",
                param_mapping={"input_data": "results.processed"},
            ),
            Step("report-module", depends_on="analysis-module"),
        ],
    )

    results = pipeline.execute(api_client)
    for r in results:
        print(r.step.module_id, r.status)
"""

from __future__ import annotations

import logging
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

_COMPLETED_STATUSES = {"completed"}
_FAILED_TERMINAL_STATUSES = {"failed", "cancelled", "cancel_requested", "timed_out", "timeout"}


class PipelineError(Exception):
    """Error during pipeline definition or execution."""


@dataclass
class Step:
    """A single step in a pipeline.

    Args:
        module_id: Module to invoke.
        params: Static input parameters.
        depends_on: Module ID of the preceding step whose output to use.
        param_mapping: Maps parameter names to dot-paths into the
            previous step's result (e.g. ``{"data": "results.filtered"}``).
        version: Optional module version constraint.
        timeout: Maximum seconds to wait for this step.
    """

    module_id: str
    params: Dict[str, Any] = field(default_factory=dict)
    depends_on: Optional[str] = None
    param_mapping: Optional[Dict[str, str]] = None
    version: Optional[str] = None
    timeout: float = 3600


@dataclass
class StepResult:
    """The outcome of executing a single pipeline step."""

    step: Step
    run_id: str
    status: str
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None


class Pipeline:
    """A sequence of module steps with dependency resolution.

    Steps are executed in topological order.  A step whose ``depends_on``
    references another step will receive that step's output, with specific
    fields extracted via ``param_mapping``.
    """

    def __init__(self, name: str, steps: List[Step]):
        self.name = name
        self.steps = steps
        self._validate()

    def _validate(self) -> None:
        ids = {s.module_id for s in self.steps}
        # Check for duplicate module_ids
        if len(ids) != len(self.steps):
            seen: set[str] = set()
            for step in self.steps:
                if step.module_id in seen:
                    raise PipelineError(f"Duplicate module_id '{step.module_id}' in pipeline steps")
                seen.add(step.module_id)
        for step in self.steps:
            if step.depends_on and step.depends_on not in ids:
                raise PipelineError(f"Step '{step.module_id}' depends on unknown step '{step.depends_on}'")
        # Check for cycles
        self._topological_sort()

    def _topological_sort(self) -> List[Step]:
        by_id: Dict[str, Step] = {s.module_id: s for s in self.steps}
        in_degree: Dict[str, int] = defaultdict(int)
        dependents: Dict[str, List[str]] = defaultdict(list)

        for step in self.steps:
            in_degree.setdefault(step.module_id, 0)
            if step.depends_on:
                in_degree[step.module_id] += 1
                dependents[step.depends_on].append(step.module_id)

        queue = deque(mid for mid, deg in in_degree.items() if deg == 0)
        ordered: List[Step] = []

        while queue:
            mid = queue.popleft()
            ordered.append(by_id[mid])
            for dep in dependents[mid]:
                in_degree[dep] -= 1
                if in_degree[dep] == 0:
                    queue.append(dep)

        if len(ordered) != len(self.steps):
            raise PipelineError("Pipeline contains a dependency cycle")

        return ordered

    def execute(
        self,
        api_client,
        *,
        poll_interval: float = 5.0,
    ) -> List[StepResult]:
        """Execute the pipeline by running steps in dependency order.

        Args:
            api_client: An :class:`~hla_compass.client.APIClient` instance.
            poll_interval: Seconds between status polls per step.

        Returns:
            A list of :class:`StepResult` for each executed step.
            If a step fails, later steps are skipped.
        """
        ordered = self._topological_sort()
        results: Dict[str, StepResult] = {}

        logger.info("Starting pipeline '%s' with %d steps", self.name, len(ordered))

        for step in ordered:
            params = dict(step.params)

            # Inject mapped parameters from the dependency's result
            if step.param_mapping and step.depends_on:
                prev = results.get(step.depends_on)
                if prev and prev.result:
                    for param_name, path in step.param_mapping.items():
                        value = _extract_path(prev.result, path)
                        if value is not None:
                            params[param_name] = value

            # Start the module run
            logger.info("Pipeline '%s': starting step '%s'", self.name, step.module_id)
            try:
                run_response = api_client.start_module_run(
                    step.module_id,
                    parameters=params,
                    mode="workflow",
                    version=step.version,
                )
            except Exception as exc:
                sr = StepResult(step=step, run_id="", status="failed", error=str(exc))
                results[step.module_id] = sr
                logger.error("Pipeline '%s': step '%s' failed to start: %s", self.name, step.module_id, exc)
                break

            run_id = ""
            if isinstance(run_response, dict):
                run_id = run_response.get("runId") or run_response.get("id") or run_response.get("run_id") or ""
            if not run_id:
                sr = StepResult(
                    step=step,
                    run_id="",
                    status="failed",
                    error="start_module_run() did not return a run identifier",
                )
                results[step.module_id] = sr
                logger.error(
                    "Pipeline '%s': step '%s' returned no run identifier — stopping pipeline",
                    self.name,
                    step.module_id,
                )
                break

            # Wait for completion
            sr = _wait_for_step(api_client, step, run_id, poll_interval)
            results[step.module_id] = sr

            if sr.status != "completed":
                logger.error(
                    "Pipeline '%s': step '%s' %s — stopping pipeline",
                    self.name,
                    step.module_id,
                    sr.status,
                )
                break

            logger.info("Pipeline '%s': step '%s' completed", self.name, step.module_id)

        return [results[s.module_id] for s in ordered if s.module_id in results]


def _wait_for_step(
    api_client,
    step: Step,
    run_id: str,
    poll_interval: float,
) -> StepResult:
    """Poll until a step completes, fails, or times out."""
    deadline = time.monotonic() + step.timeout

    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return StepResult(
                step=step,
                run_id=run_id,
                status="timeout",
                error=f"Step timed out after {step.timeout}s",
            )

        try:
            status_resp = api_client.get_module_run(run_id)
        except Exception as exc:
            return StepResult(step=step, run_id=run_id, status="failed", error=str(exc))

        status = str(status_resp.get("status") or status_resp.get("state") or "").lower()

        if status in _COMPLETED_STATUSES:
            try:
                result = api_client.get_module_run_result(run_id)
            except Exception:
                result = None
            return StepResult(step=step, run_id=run_id, status="completed", result=result)

        if status in _FAILED_TERMINAL_STATUSES:
            error_info = status_resp.get("error") or {}
            error_msg = error_info.get("message") if isinstance(error_info, dict) else str(error_info)
            return StepResult(step=step, run_id=run_id, status=status, error=error_msg or None)

        time.sleep(min(poll_interval, remaining))


def _extract_path(data: Any, path: str) -> Any:
    """Extract a value from nested dicts using a dot-separated path.

    Example::

        >>> _extract_path({"results": {"filtered": [1, 2, 3]}}, "results.filtered")
        [1, 2, 3]
    """
    current = data
    for part in path.split("."):
        if isinstance(current, dict):
            current = current.get(part)
        else:
            return None
    return current
