"""Execution runner for dotpromptz prompts.

Provides reusable ``run_single`` and ``run_batch`` functions that handle
the render → generate pipeline without any CLI or I/O concerns.

These are the primary entry points for Python SDK users who want to
execute prompts programmatically, including batch execution with
concurrency control.

Example usage::

    import asyncio
    from pathlib import Path
    from dotpromptz import Dotprompt
    from dotpromptz.runner import run_single, run_batch
    from dotpromptz.adapters import get_adapter


    async def main():
        source = Path('my.prompt').read_text()
        adapter = get_adapter('openai')

        # Single execution
        response = await run_single(source, {'topic': 'AI'}, adapter)
        print(response.text)

        # Batch execution
        inputs = [{'topic': 'AI'}, {'topic': 'ML'}, {'topic': 'DL'}]
        results = await run_batch(source, inputs, adapter, max_workers=5)
        for r in results:
            if r.status == 'success':
                print(r.response.text)


    asyncio.run(main())
"""

import asyncio
import json
import random
import re
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import structlog
import yaml

from dotpromptz.adapters import get_adapter, list_adapters
from dotpromptz.adapters._base import Adapter, GenerateResponse
from dotpromptz.credentials import CredentialPool, load_credentials
from dotpromptz.dotprompt import Dotprompt
from dotpromptz.typing import (
    AdapterConfig,
    DataArgument,
    DataPart,
    ErrorAction,
    ErrorHandlerRule,
    MediaPart,
    PromptOutputConfig,
    RetryConfig,
    RuntimeConfig,
    TextPart,
    ToolRequestPart,
    ToolResponsePart,
)

logger = structlog.get_logger(__name__)


class BatchAbortError(Exception):
    """Raised when a batch item triggers the ``abort`` error-handler action.

    Attributes:
        index: The batch item index that triggered the abort.
        original: The original exception.
    """

    def __init__(self, index: int, original: Exception) -> None:
        """Initialise with the failing index and original error."""
        self.index = index
        self.original = original
        super().__init__(f'Batch aborted at index {index}: {original}')


@dataclass
class BatchResult:
    """Result of a single item within a batch execution.

    Attributes:
        index: Zero-based position in the original input list.
        input: The input variables dict for this item.
        status: Either ``'success'`` or ``'error'``.
        response: The ``GenerateResponse`` on success, ``None`` on error.
        error: Error message string on failure, ``None`` on success.
        attempts: Total number of attempts (1 = no retries).
        elapsed: Wall-clock seconds from first attempt to final result.
    """

    index: int
    input: dict[str, Any]
    status: str  # 'success' | 'error'
    response: GenerateResponse | None = field(default=None)
    error: str | None = field(default=None)
    attempts: int = field(default=1)
    elapsed: float = field(default=0.0)


def _compute_delay(attempt: int, retry_cfg: RetryConfig, *, backoff_base_override: float | None = None) -> float:
    """Compute the delay before the next retry attempt.

    Args:
        attempt: Zero-based retry attempt number (0 = first retry).
        retry_cfg: The retry configuration.
        backoff_base_override: Override ``retry_cfg.backoff_base`` (used by
            per-rule overrides in ``ErrorHandlerRule``).

    Returns:
        The delay in seconds, capped at ``retry_cfg.backoff_max`` and
        optionally jittered.
    """
    base = backoff_base_override if backoff_base_override is not None else retry_cfg.backoff_base
    strategy = retry_cfg.backoff_strategy

    if strategy == 'fixed':
        delay = base
    elif strategy == 'linear':
        delay = base * (attempt + 1)
    else:  # exponential (default)
        delay = base * (2**attempt)

    delay = min(delay, retry_cfg.backoff_max)

    if retry_cfg.jitter:
        delay = delay * random.uniform(0.5, 1.0)  # noqa: S311

    return delay


def _extract_status_code(exc: Exception) -> int | None:
    """Extract an HTTP status code from an SDK exception.

    Supports OpenAI, Anthropic, and Google GenAI SDK exceptions.
    Falls back to generic attribute checks and cause/context unwrapping.

    Args:
        exc: The exception to inspect.

    Returns:
        The HTTP status code, or ``None`` if not extractable.
    """
    # OpenAI: openai.APIStatusError has .status_code (int)
    # Anthropic: anthropic.APIStatusError has .status_code (int)
    status = getattr(exc, 'status_code', None)
    if isinstance(status, int):
        return status

    # Google GenAI: google.genai.errors.APIError has .code (int)
    code = getattr(exc, 'code', None)
    if isinstance(code, int):
        return code

    # Response object fallback (some SDK exceptions carry the response).
    resp = getattr(exc, 'response', None) or getattr(exc, 'http_response', None)
    if resp is not None:
        resp_status = getattr(resp, 'status_code', None) or getattr(resp, 'status', None)
        if isinstance(resp_status, int):
            return resp_status

    # Some Google exceptions expose grpc_status_code; map common ones.
    grpc_code = getattr(exc, 'grpc_status_code', None)
    if grpc_code is not None:
        _grpc_http_map = {
            14: 503,  # UNAVAILABLE
            4: 504,  # DEADLINE_EXCEEDED
            8: 429,  # RESOURCE_EXHAUSTED
            13: 500,  # INTERNAL
            2: 500,  # UNKNOWN
        }
        grpc_val = grpc_code.value[0] if hasattr(grpc_code, 'value') else grpc_code
        if isinstance(grpc_val, int):
            return _grpc_http_map.get(grpc_val)

    # Dig into wrapped exceptions (cause / context).
    for inner in (exc.__cause__, exc.__context__):
        if inner is not None:
            inner_code = _extract_status_code(inner)
            if inner_code is not None:
                return inner_code

    return None


def _classify_error(
    exc: Exception,
    retry_cfg: RetryConfig,
    error_handlers: list[ErrorHandlerRule] | None,
) -> tuple[ErrorAction, int, float | None]:
    """Classify an exception into an action, max retries, and backoff override.

    Evaluation order:
    1. Walk ``error_handlers`` in order; first matching rule wins.
    2. If no rule matches, fall back to the default retry policy:
       retry if the status code is in ``retry_cfg.retryable_status_codes``.
    3. If the exception has no status code and no handler matches,
       default to ``RETRY`` (the exception is assumed transient).

    Args:
        exc: The exception to classify.
        retry_cfg: The global retry configuration.
        error_handlers: Ordered list of handler rules (may be ``None``).

    Returns:
        A 3-tuple of ``(action, max_retries, backoff_base_override)``.
    """
    status_code = _extract_status_code(exc)
    error_msg = str(exc)

    # 1. Check explicit handler rules.
    if error_handlers:
        for rule in error_handlers:
            matched = False
            if rule.status_codes and status_code is not None:
                if status_code in rule.status_codes:
                    matched = True
            if not matched and rule.pattern:
                if re.search(rule.pattern, error_msg):
                    matched = True
            if matched:
                max_retries = rule.max_retries if rule.max_retries is not None else retry_cfg.max_retries
                return rule.action, max_retries, rule.backoff_base

    # 2. Default policy: retry if status code is retryable (or unknown).
    if status_code is not None:
        if status_code in retry_cfg.retryable_status_codes:
            return ErrorAction.RETRY, retry_cfg.max_retries, None
        # Known non-retryable status code -> skip.
        return ErrorAction.SKIP, 0, None

    # 3. No status code, no handler match -> assume transient, retry.
    return ErrorAction.RETRY, retry_cfg.max_retries, None


async def _execute_with_retry(
    coro_factory: Callable[[], Any],
    retry_cfg: RetryConfig | None,
    error_handlers: list[ErrorHandlerRule] | None,
    timeout: float | None,
) -> tuple[Any, int]:
    """Execute *coro_factory()* with optional retry and timeout.

    Args:
        coro_factory: A zero-arg callable that returns an awaitable
            (called fresh on each attempt).
        retry_cfg: Retry configuration.  ``None`` disables retrying.
        error_handlers: Per-error rules (may be ``None``).
        timeout: Per-attempt timeout in seconds (``None`` = no timeout).

    Returns:
        A 2-tuple ``(result, attempts)`` where *result* is the awaitable's
        return value and *attempts* is the total number of tries.

    Raises:
        The last exception if all retries are exhausted.
        asyncio.TimeoutError: If *timeout* is exceeded on the last attempt.
    """
    max_attempts = 1 + (retry_cfg.max_retries if retry_cfg else 0)
    last_exc: Exception | None = None

    for attempt in range(max_attempts):
        try:
            coro = coro_factory()
            if timeout is not None:
                result = await asyncio.wait_for(coro, timeout=timeout)
            else:
                result = await coro
            return result, attempt + 1
        except Exception as exc:
            last_exc = exc

            if retry_cfg is None:
                raise

            action, rule_max_retries, backoff_override = _classify_error(
                exc,
                retry_cfg,
                error_handlers,
            )

            if action == ErrorAction.ABORT:
                raise
            if action == ErrorAction.SKIP:
                raise

            # RETRY: check if we have attempts left (rule may override global).
            effective_max = min(rule_max_retries, max_attempts - 1)
            if attempt >= effective_max:
                raise

            delay = _compute_delay(attempt, retry_cfg, backoff_base_override=backoff_override)
            logger.warning(
                'Retrying after error.',
                attempt=attempt + 1,
                max_retries=effective_max,
                delay=f'{delay:.2f}s',
                error=str(exc),
            )
            await asyncio.sleep(delay)

    # Should not reach here, but satisfy type checker.
    raise last_exc  # type: ignore[misc]


async def run_single(
    source: str,
    input_data: dict[str, Any],
    adapter: Adapter,
    *,
    dp: 'Dotprompt | None' = None,
) -> GenerateResponse:
    """Render and generate a single prompt. Pure logic, no I/O.

    Args:
        source: The ``.prompt`` file content.
        input_data: Dict of input variables for template rendering.
        adapter: An LLM adapter instance (e.g. from ``get_adapter``).
        dp: Optional ``Dotprompt`` instance for reusing helpers/partials.
            A default instance is created if not provided.

    Returns:
        The ``GenerateResponse`` from the adapter.

    Raises:
        Exception: Any error from rendering or adapter generation is
            propagated directly.
    """
    dp = dp or Dotprompt()
    rendered = dp.render(source, data=DataArgument(input=input_data), variables=input_data)
    return await adapter.generate(rendered)


async def run_batch(
    source: str,
    input_list: list[dict[str, Any]],
    adapter: Adapter,
    *,
    max_workers: int = 5,
    dp: 'Dotprompt | None' = None,
    on_item_complete: Callable[[BatchResult], None] | None = None,
    runtime: RuntimeConfig | None = None,
) -> list[BatchResult]:
    """Execute batch prompt rendering and generation with concurrency control.

    Pure execution logic without file I/O, ``sys.exit``, or logging side
    effects. Results are returned as a list of ``BatchResult`` objects
    sorted by index.

    When *runtime* is provided, its ``retry``, ``timeout``, and
    ``error_handlers`` fields are used to control per-item behaviour.
    If an error handler returns ``ErrorAction.ABORT``, remaining in-flight
    tasks are cancelled and a ``BatchAbortError`` is raised.

    Args:
        source: The ``.prompt`` file content.
        input_list: List of input variable dicts (one per batch item).
        adapter: An LLM adapter instance.
        max_workers: Maximum number of concurrent tasks (default 5).
        dp: Optional ``Dotprompt`` instance for reusing helpers/partials.
            A default instance is created if not provided.
        on_item_complete: Optional callback invoked after each item
            completes (success or error). Useful for progress reporting.
        runtime: Optional ``RuntimeConfig`` with retry, timeout, and
            error-handler settings.

    Returns:
        A list of ``BatchResult`` objects sorted by ``index``.

    Raises:
        BatchAbortError: If an error handler triggers the ``abort`` action.
    """
    dp = dp or Dotprompt()
    retry_cfg = runtime.retry if runtime else None
    error_handlers = runtime.error_handlers if runtime else None
    timeout = runtime.timeout if runtime else None
    sem = asyncio.Semaphore(max_workers)
    abort_event = asyncio.Event()

    async def _process_one(index: int, variables: dict[str, Any]) -> BatchResult:
        async with sem:
            if abort_event.is_set():
                return BatchResult(
                    index=index,
                    input=variables,
                    status='error',
                    error='Batch aborted before execution.',
                )

            t0 = time.monotonic()
            attempts = 1
            try:
                rendered = dp.render(
                    source,
                    data=DataArgument(input=variables),
                    variables=variables,
                )

                def _make_coro() -> Any:
                    return adapter.generate(rendered)

                response, attempts = await _execute_with_retry(
                    _make_coro,
                    retry_cfg,
                    error_handlers,
                    timeout,
                )
                elapsed = time.monotonic() - t0
                result = BatchResult(
                    index=index,
                    input=variables,
                    status='success',
                    response=response,
                    error=None,
                    attempts=attempts,
                    elapsed=elapsed,
                )
            except Exception as exc:
                elapsed = time.monotonic() - t0
                # Determine if this was an abort action.
                if retry_cfg and error_handlers:
                    action, _, _ = _classify_error(exc, retry_cfg, error_handlers)
                    if action == ErrorAction.ABORT:
                        abort_event.set()
                        result = BatchResult(
                            index=index,
                            input=variables,
                            status='error',
                            error=str(exc),
                            attempts=attempts,
                            elapsed=elapsed,
                        )
                        if on_item_complete is not None:
                            on_item_complete(result)
                        raise BatchAbortError(index, exc) from exc
                result = BatchResult(
                    index=index,
                    input=variables,
                    status='error',
                    response=None,
                    error=str(exc),
                    attempts=attempts,
                    elapsed=elapsed,
                )
            if on_item_complete is not None:
                on_item_complete(result)
            return result

    tasks = [asyncio.create_task(_process_one(i, v)) for i, v in enumerate(input_list)]

    # Gather with return_exceptions so we can inspect BatchAbortError.
    raw_results = await asyncio.gather(*tasks, return_exceptions=True)

    # Check for abort — cancel remaining and collect partial results.
    aborted = False
    abort_error: BatchAbortError | None = None
    final: list[BatchResult] = []
    for i, r in enumerate(raw_results):
        if isinstance(r, BatchAbortError):
            aborted = True
            abort_error = r
            # The abort item already built its BatchResult before raising;
            # build one here for the gather result.
            final.append(
                BatchResult(
                    index=r.index,
                    input=input_list[r.index],
                    status='error',
                    error=str(r.original),
                )
            )
        elif isinstance(r, Exception):
            final.append(
                BatchResult(
                    index=i,
                    input=input_list[i],
                    status='error',
                    response=None,
                    error=str(r),
                )
            )
        else:
            final.append(r)
    final.sort(key=lambda x: x.index)

    if aborted and abort_error is not None:
        raise abort_error

    return final


def resolve_adapter(rendered: Any, *, pool: CredentialPool | None = None) -> Adapter:
    """Resolve an adapter from a RenderedPrompt's adapter field.

    Uses the ``CredentialPool`` to obtain credentials via round-robin.

    Args:
        rendered: The RenderedPrompt with adapter configuration.
        pool: Optional credential pool.  If ``None``, one is loaded from
            the ``API_CREDENTIALS`` environment variable.

    Returns:
        An ``Adapter`` instance ready to generate responses.

    Raises:
        ValueError: If no adapter is specified in *rendered*.
    """
    adapter_cfg = rendered.adapter  # str | AdapterConfig | None
    adapter_name: str | None = None
    group: str | None = None
    groups: list[str] | None = None

    if isinstance(adapter_cfg, AdapterConfig):
        adapter_name = adapter_cfg.name
        group = adapter_cfg.group
        groups = adapter_cfg.groups
    elif isinstance(adapter_cfg, str):
        adapter_name = adapter_cfg

    if not adapter_name:
        available = ', '.join(list_adapters())
        msg = f'No adapter specified. Add "adapter: <name>" to the .prompt frontmatter. Available adapters: {available}'
        raise ValueError(msg)

    # Extract model name from config for credential filtering.
    model_name: str | None = None
    config = rendered.config if isinstance(rendered.config, dict) else {}
    if isinstance(config, dict):
        model_name = config.get('model')

    if pool is None:
        pool = load_credentials()

    credential = pool.next(adapter_name, group=group, groups=groups, model=model_name)
    return get_adapter(adapter_name, credential=credential)


def serialize_result(data: Any, fmt: str) -> str:
    """Serialize *data* to a string according to *fmt*.

    Args:
        data: Arbitrary data (usually a dict or string).
        fmt: One of ``'json'``, ``'yaml'``, ``'txt'``.

    Returns:
        A formatted string.

    """
    if fmt == 'json':
        if isinstance(data, str):
            try:
                data = json.loads(data)
            except (json.JSONDecodeError, TypeError):
                pass
        return json.dumps(data, indent=2, ensure_ascii=False)
    if fmt == 'yaml':
        if isinstance(data, str):
            try:
                data = json.loads(data)
            except (json.JSONDecodeError, TypeError):
                pass
        return yaml.dump(data, allow_unicode=True, default_flow_style=False, sort_keys=False)
    # txt
    return str(data)


def part_to_dict(part: Any) -> dict[str, Any]:
    """Convert a Part to a simple dict for JSON output.

    Args:
        part: A dotprompt Part object.

    Returns:
        A simplified dict representation.
    """
    if isinstance(part, TextPart):
        return {'text': part.text}
    if isinstance(part, MediaPart):
        return {'media': part.media.model_dump(exclude_none=True, by_alias=True)}
    if isinstance(part, DataPart):
        return {'data': part.data}
    if isinstance(part, ToolRequestPart):
        return {'tool_request': part.tool_request.model_dump(exclude_none=True)}
    if isinstance(part, ToolResponsePart):
        return {'tool_response': part.tool_response.model_dump(exclude_none=True)}
    return {'type': type(part).__name__}


def write_output(
    content: str | bytes,
    output_config: PromptOutputConfig,
    output_dir: Path,
    file_name: str,
) -> Path:
    """Write a single result to disk.

    Determines the file extension from *output_config.format* (``'png'`` for
    image format, otherwise the format string itself) and writes *content*
    to ``output_dir / {file_name}.{ext}``.

    Args:
        content: Text or binary data to write.
        output_config: The output configuration (used for format/extension).
        output_dir: Resolved output directory (created if needed).
        file_name: File stem (no extension).

    Returns:
        The ``Path`` of the written file.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    ext = 'png' if output_config.format == 'image' else output_config.format
    filepath = output_dir / f'{file_name}.{ext}'
    if isinstance(content, bytes):
        filepath.write_bytes(content)
    else:
        filepath.write_text(content, encoding='utf-8')
    return filepath


def extract_response_content(
    response: GenerateResponse,
    fmt: str,
) -> str | bytes:
    """Extract the primary content from a ``GenerateResponse``.

    For image formats, returns the raw image bytes.  For text-based formats,
    the response text (or serialised tool-call list) is returned as a string
    via :func:`serialize_result`.

    Args:
        response: The adapter response.
        fmt: The output format (``'json'``, ``'yaml'``, ``'txt'``, ``'image'``).

    Returns:
        ``bytes`` for image output, ``str`` for everything else.

    Raises:
        ValueError: If *fmt* is ``'image'`` but the response contains no
            image data.
    """
    if fmt == 'image':
        if response.image:
            return response.image.data
        raise ValueError('Expected image output but model returned no image data.')

    # Text-based formats
    if response.text:
        return serialize_result(response.text, fmt)
    if response.tool_calls:
        return serialize_result(
            [tc.model_dump() for tc in response.tool_calls],
            fmt,
        )
    return ''


def aggregate_batch_results(
    results: list[BatchResult],
    fmt: str,
    file_name_tpl: str,
) -> dict[str, Any]:
    """Aggregate batch results into a single dict keyed by file stem.

    Each entry is keyed as ``{file_name_tpl}_{result.index}``.  Error results
    are stored as ``{'error': ...}``, text results are optionally parsed from
    JSON for structured formats, tool-call results are stored as dicts, and
    empty responses become ``None``.

    Args:
        results: List of :class:`BatchResult` objects.
        fmt: The output format (``'json'``, ``'yaml'``, ``'txt'``).
        file_name_tpl: File-name template used as the key prefix.

    Returns:
        A dict mapping stem names to their (possibly parsed) content.
    """
    aggregated: dict[str, Any] = {}
    for result in results:
        stem = f'{file_name_tpl}_{result.index}'
        if result.status == 'error':
            aggregated[stem] = {'error': result.error}
        elif result.response and result.response.text:
            # Try to parse JSON text for structured formats
            try:
                aggregated[stem] = json.loads(result.response.text)
            except (json.JSONDecodeError, TypeError):
                aggregated[stem] = result.response.text
        elif result.response and result.response.tool_calls:
            aggregated[stem] = [tc.model_dump() for tc in result.response.tool_calls]
        else:
            aggregated[stem] = None
    return aggregated
