"""Execution runner for dotpromptz prompts.

Provides the run() function that handles the render-then-generate pipeline,
including batch execution with concurrency control and retry logic.
"""

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

import jsonschema
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,
    DataPart,
    MediaPart,
    PromptOutputConfig,
    RetryConfig,
    RuntimeConfig,
    TextPart,
    ToolRequestPart,
    ToolResponsePart,
)

logger = structlog.get_logger(__name__)

_STRUCTURED_OUTPUT_ERROR_PREVIEW_CHARS = 4000


@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 _parse_structured_json_text(text: str) -> Any:
    """Parse structured model output text as JSON.

    Both ``json`` and ``yaml`` output modes require the model to return JSON
    when schema validation is enabled. YAML files are produced by converting
    validated JSON into YAML in ``serialize_result``.

    Args:
        text: Raw model text response.

    Returns:
        Parsed JSON value.

    Raises:
        ValueError: If the text is not valid JSON.
    """
    try:
        return json.loads(text)
    except json.JSONDecodeError as exc:
        raw_preview = _preview_value(text)
        raise ValueError(
            'Structured output requires valid JSON text from the model. '
            f'JSON parse error at line {exc.lineno}, column {exc.colno}: {exc.msg}. '
            f'Raw model text: {raw_preview}'
        ) from exc
    except TypeError as exc:
        raw_preview = _preview_value(text)
        raise ValueError(
            'Structured output requires valid JSON text from the model. '
            f'Received non-text response. Raw model text: {raw_preview}'
        ) from exc


def _preview_value(value: Any, *, max_chars: int = _STRUCTURED_OUTPUT_ERROR_PREVIEW_CHARS) -> str:
    """Return a safe, length-limited string preview for debugging output.

    Args:
        value: Any Python value to render into a preview string.
        max_chars: Maximum preview length in characters.

    Returns:
        A preview string suitable for exception messages.
    """
    try:
        rendered = json.dumps(value, ensure_ascii=False)
    except TypeError:
        rendered = repr(value)

    if len(rendered) <= max_chars:
        return rendered

    omitted = len(rendered) - max_chars
    return f'{rendered[:max_chars]}... [truncated {omitted} chars]'


def _validate_response_schema(response: GenerateResponse, output_config: PromptOutputConfig | None) -> None:
    """Validate model response text against ``output.schema`` when configured.

    Args:
        response: Adapter response from the model.
        output_config: Prompt output configuration.

    Raises:
        ValueError: If response text is missing/invalid or fails schema validation.
    """
    if output_config is None or output_config.schema_ is None:
        return
    if output_config.format not in ('json', 'yaml'):
        return
    if not response.text:
        raise ValueError('Schema validation failed: model returned no text content.')

    instance = _parse_structured_json_text(response.text)
    try:
        jsonschema.validate(instance=instance, schema=output_config.schema_)
    except jsonschema.exceptions.SchemaError as exc:
        raise ValueError(f'Invalid output.schema: {exc.message}') from exc
    except jsonschema.exceptions.ValidationError as exc:
        path = '.'.join(str(part) for part in exc.path) or '<root>'
        raw_preview = _preview_value(response.text)
        invalid_value_preview = _preview_value(exc.instance)
        raise ValueError(
            f'Schema validation failed at {path}: {exc.message}. '
            f'Invalid value: {invalid_value_preview}. '
            f'Raw model text: {raw_preview}'
        ) from exc


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,
    abort_event: asyncio.Event | None = 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).
        abort_event: Optional abort signal. When set, no further attempts
            are started and a ``RuntimeError`` is raised.

    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.
    """
    if retry_cfg is not None:
        if retry_cfg.max_retries < 0:
            raise ValueError('retry.max_retries must be >= 0')
        if retry_cfg.retry_delay < 0:
            raise ValueError('retry.retry_delay must be >= 0')

    max_attempts = 1 + (retry_cfg.max_retries if retry_cfg else 0)
    last_exc: Exception | None = None

    for attempt in range(max_attempts):
        if abort_event is not None and abort_event.is_set():
            raise RuntimeError('Aborted before execution.')
        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),
            )
            if abort_event is None:
                await asyncio.sleep(delay)
            else:
                try:
                    await asyncio.wait_for(abort_event.wait(), timeout=delay)
                except TimeoutError:
                    pass
                else:
                    raise RuntimeError('Aborted during retry wait.') from exc

    # 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,
    abort_event: asyncio.Event | 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.
        abort_event: Optional abort signal. Items not yet started are marked
            as errors when the event is set.

    Returns:
        A list of ``BatchResult`` objects sorted by ``index``.
    """
    if max_workers < 1:
        raise ValueError('max_workers must be >= 1')

    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:
        if abort_event is not None and abort_event.is_set():
            return BatchResult(
                index=index,
                input=variables,
                status='error',
                response=None,
                error='Aborted before execution.',
                attempts=0,
                elapsed=0.0,
            )

        async with sem:
            t0 = time.monotonic()
            attempts = 1
            try:
                if abort_event is not None and abort_event.is_set():
                    raise RuntimeError('Aborted before execution.')
                rendered = dp.render(
                    source,
                    input=variables,
                )

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

                    async def _generate_and_validate() -> GenerateResponse:
                        response = await current_adapter.generate(rendered)
                        _validate_response_schema(response, rendered.output)
                        return response

                    return _generate_and_validate()

                response, attempts = await _execute_with_retry(
                    _make_coro,
                    retry_cfg,
                    timeout,
                    abort_event=abort_event,
                )
                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,
                )
            return result

    async def _safe_process_one(index: int, variables: dict[str, Any]) -> BatchResult:
        try:
            return await _process_one(index, variables)
        except Exception as exc:
            return BatchResult(
                index=index,
                input=variables,
                status='error',
                response=None,
                error=str(exc),
            )

    final: list[BatchResult] = []
    tasks = [asyncio.create_task(_safe_process_one(i, v)) for i, v in enumerate(input_list)]
    for task in asyncio.as_completed(tasks):
        result = await task
        if on_item_complete is not None:
            try:
                on_item_complete(result)
            except Exception as exc:
                final.append(
                    BatchResult(
                        index=result.index,
                        input=result.input,
                        status='error',
                        response=None,
                        error=str(exc),
                        attempts=result.attempts,
                        elapsed=result.elapsed,
                    )
                )
                continue
        final.append(result)

    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,
    *,
    force: bool = False,
) -> 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).
        force: Whether to overwrite an existing file.

    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 filepath.exists() and not force:
        raise FileExistsError(f'Output file already exists: {filepath}. Use --force to overwrite.')
    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 ''
