"""Execution runner for dotpromptz prompts.

Provides a reusable ``run`` function that handles the render → generate
pipeline without any CLI or I/O concerns.

This is the primary entry point 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
    from dotpromptz.adapters import get_adapter


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

        # Single execution (batch of one)
        results = await run(source, [{'topic': 'AI'}], adapter)
        print(results[0].response.text)

        # Batch execution
        inputs = [{'topic': 'AI'}, {'topic': 'ML'}, {'topic': 'DL'}]
        results = await run(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 time
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime
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, render_template_str
from dotpromptz.typing import (
    AdapterConfig,
    DataArgument,
    DataPart,
    MediaPart,
    PromptOutputConfig,
    RetryConfig,
    RuntimeConfig,
    TextPart,
    ToolRequestPart,
    ToolResponsePart,
)

logger = structlog.get_logger(__name__)



@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) -> float:
    """Compute the delay before the next retry attempt.

    Uses a simple linear strategy: ``retry_delay * (attempt + 1)``.

    Args:
        attempt: Zero-based retry attempt number (0 = first retry).
        retry_cfg: The retry configuration.

    Returns:
        The delay in seconds.
    """
    return retry_cfg.retry_delay * (attempt + 1)

def _make_adapter_for_attempt(
    attempt: int,
    adapter: Adapter,
    pool: CredentialPool | None,
) -> Adapter:
    """Return the adapter to use for a given retry attempt.

    On the first attempt (``attempt == 0``) or when no credential pool is
    available, the original *adapter* is returned unchanged.  On subsequent
    attempts the pool is queried for the next credential matching the same
    adapter type and model, and a fresh adapter is constructed.

    Args:
        attempt: Zero-based attempt number.
        adapter: The original adapter instance.
        pool: Optional credential pool for rotation.

    Returns:
        An ``Adapter`` instance to use for this attempt.
    """
    if attempt == 0 or pool is None:
        return adapter
    adapter_name = adapter._credential.adapter
    model_name: str | None = getattr(adapter, '_default_model', None)
    try:
        credential = pool.next(adapter_name, model=model_name)
        return get_adapter(adapter_name, credential=credential)
    except (ValueError, KeyError):
        # No alternative credential available; fall back to original adapter.
        return adapter


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

    Any exception triggers a retry (no error classification).  The
    *coro_factory* receives the zero-based attempt index so callers can
    implement credential rotation or other per-attempt logic.

    Args:
        coro_factory: A callable accepting the attempt index and returning
            an awaitable (called fresh on each attempt).
        retry_cfg: Retry configuration.  ``None`` disables retrying.
        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(attempt)
            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

            if attempt >= retry_cfg.max_retries:
                raise

            delay = _compute_delay(attempt, retry_cfg)
            logger.warning(
                'Retrying after error.',
                attempt=attempt + 1,
                max_retries=retry_cfg.max_retries,
                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(
    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,
    pool: CredentialPool | 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`` and ``timeout`` fields are used
    to control per-item retry behaviour.  On retry, if a *pool* is supplied,
    credentials are rotated via ``_make_adapter_for_attempt``.

    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 and timeout settings.
        pool: Optional ``CredentialPool`` for credential rotation on retry.

    Returns:
        A list of ``BatchResult`` objects sorted by ``index``.
    """
    dp = dp or Dotprompt()
    retry_cfg = runtime.retry if runtime else None
    timeout = runtime.timeout if runtime else None
    sem = asyncio.Semaphore(max_workers)

    async def _process_one(index: int, variables: dict[str, Any]) -> BatchResult:
        async with sem:
            t0 = time.monotonic()
            attempts = 1
            try:
                rendered = dp.render(
                    source,
                    data=DataArgument(input=variables),
                    variables=variables,
                )

                def _make_coro(attempt: int) -> Any:
                    current_adapter = _make_adapter_for_attempt(attempt, adapter, pool)
                    return current_adapter.generate(rendered)

                response, attempts = await _execute_with_retry(
                    _make_coro,
                    retry_cfg,
                    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
                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)]
    raw_results = await asyncio.gather(*tasks, return_exceptions=True)

    final: list[BatchResult] = []
    for i, r in enumerate(raw_results):
        if 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)

    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,
    *,
    prompt_path: Path | None = None,
) -> dict[str, Any]:
    """Aggregate batch results into a single dict keyed by file stem.

    When *prompt_path* is provided, each key is resolved as a Handlebars
    template with ``NAME``, ``TIME``, ``INDEX``, and the item's input
    variables.  Otherwise falls back to ``{file_name_tpl}_{result.index}``.

    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.
        prompt_path: Optional path to the ``.prompt`` file for template resolution.

    Returns:
        A dict mapping stem names to their (possibly parsed) content.
    """
    aggregated: dict[str, Any] = {}
    for result in results:
        if prompt_path is not None:
            stem = _resolve_file_name_for_aggregation(
                file_name_tpl, prompt_path, item_vars=result.input, index=result.index,
            )
        else:
            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


def _resolve_file_name_for_aggregation(
    file_name_tpl: str,
    prompt_path: Path,
    *,
    item_vars: dict[str, Any] | None = None,
    index: int | None = None,
) -> str:
    """Resolve a file_name template for aggregation dict keys.

    UPPERCASE auto-variables (``NAME``, ``TIME``, ``INDEX``) are merged with
    input variables into a single context dict.

    Args:
        file_name_tpl: Raw file_name template string.
        prompt_path: Path to the ``.prompt`` file.
        item_vars: Input variables dict for the current batch item.
        index: Zero-based batch item index.

    Returns:
        Resolved file name string.
    """
    ctx: dict[str, Any] = {
        'NAME': prompt_path.stem,
        'TIME': datetime.now().strftime('%Y%m%d_%H%M%S'),
        **(item_vars or {}),
    }
    if index is not None:
        ctx['INDEX'] = str(index)
    return render_template_str(file_name_tpl, ctx)
