"""Execution runtime and plan orchestration for prompt batches."""

from __future__ import annotations

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

import jsonschema

from dotpromptz.adapters import get_adapter, list_adapters
from dotpromptz.adapters._base import Adapter, GenerateResponse
from dotpromptz.credentials import CredentialPool, load_credentials
from dotpromptz.errors import ExecutionConfigError
from dotpromptz.logging import get_logger
from dotpromptz.planner import ExecutionPlan
from dotpromptz.typing import (
    AdapterConfig,
    DataPart,
    MediaPart,
    PromptOutputConfig,
    RenderedPrompt,
    RetryConfig,
    RuntimeConfig,
    TextPart,
)
from dotpromptz.yamlio import dump_yaml

logger = get_logger(__name__)

_STRUCTURED_OUTPUT_ERROR_PREVIEW_CHARS = 4000


def _resolve_output_extension(fmt: str) -> str:
    """Resolve output file extension from ``output.format``."""
    if fmt == 'image':
        return 'png'
    return fmt


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

    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."""
    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."""
    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."""
    if output_config is None or output_config.schema_value 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_value)
    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 delay before next retry attempt."""
    return retry_cfg.retry_delay * (attempt + 1)


def _make_adapter_for_attempt(
    attempt: int,
    adapter: Adapter,
    pool: CredentialPool | None,
) -> Adapter:
    """Return adapter for current retry attempt, rotating credentials when possible."""
    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):
        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."""
    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

    raise last_exc  # type: ignore[misc]


async def execute_batch(
    rendered_prompts: list[RenderedPrompt],
    input_list: list[dict[str, Any]],
    adapter: Adapter,
    *,
    max_workers: int = 5,
    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 a rendered prompt batch with concurrency control."""
    if max_workers < 1:
        raise ValueError('max_workers must be >= 1')
    if len(rendered_prompts) != len(input_list):
        raise ValueError('rendered_prompts and input_list must have the same length')
    if not rendered_prompts:
        return []

    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], rendered: RenderedPrompt) -> 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.')

                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
                return BatchResult(
                    index=index,
                    input=variables,
                    status='success',
                    response=response,
                    error=None,
                    attempts=attempts,
                    elapsed=elapsed,
                )
            except Exception as exc:
                elapsed = time.monotonic() - t0
                return BatchResult(
                    index=index,
                    input=variables,
                    status='error',
                    response=None,
                    error=str(exc),
                    attempts=attempts,
                    elapsed=elapsed,
                )

    async def _safe_process_one(index: int, variables: dict[str, Any], rendered: RenderedPrompt) -> BatchResult:
        try:
            return await _process_one(index, variables, rendered)
        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(index, variables, rendered_prompts[index]))
        for index, variables 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 item: item.index)
    return final


def resolve_batch_adapter(rendered: RenderedPrompt, *, pool: CredentialPool | None = None) -> Adapter:
    """Resolve an adapter from a ``RenderedPrompt`` adapter field."""
    adapter_cfg = rendered.adapter
    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())
        raise ValueError(
            f'No adapter specified. Add "adapter: <name>" to the .prompt frontmatter. Available adapters: {available}'
        )

    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_output(data: Any, fmt: str) -> str:
    """Serialize data into the target output format."""
    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 dump_yaml(data, allow_unicode=True, default_flow_style=False, sort_keys=False)
    return str(data)


def part_to_output_dict(part: Any) -> dict[str, Any]:
    """Convert a content part model to a simple serializable dict."""
    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}
    return {'type': type(part).__name__}


def write_output_file(
    content: str | bytes,
    output_config: PromptOutputConfig,
    output_dir: Path,
    file_name: str,
    *,
    force: bool = False,
    extension_override: str | None = None,
) -> Path:
    """Write one result to disk using ``output.format`` as extension."""
    output_dir.mkdir(parents=True, exist_ok=True)
    ext = extension_override or _resolve_output_extension(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_output_content(response: GenerateResponse, fmt: str) -> str | bytes:
    """Extract and serialize response payload for output."""
    if fmt == 'image':
        if not response.images:
            raise ValueError('Image output requires at least one generated image.')
        first = response.images[0]
        if first.mime_type.lower() != 'image/png':
            raise ValueError(f'Image output requires MIME image/png; got {first.mime_type!r}.')
        return first.data
    if response.text:
        return serialize_output(response.text, fmt)
    return ''


@dataclass(frozen=True)
class ExecutionReport:
    """Execution result for one planned run."""

    batch_results: list[BatchResult]
    written_files: list[Path]
    output_dir: Path

    @property
    def all_failed(self) -> bool:
        """Return whether every batch item failed."""
        return bool(self.batch_results) and all(item.status == 'error' for item in self.batch_results)


def _write_planned_output(
    content: str | bytes,
    output_config: PromptOutputConfig,
    output_dir: Path,
    stem: str,
    *,
    force: bool,
    extension_override: str | None = None,
) -> Path:
    """Write one planned output file."""
    return write_output_file(
        content,
        output_config,
        output_dir,
        stem,
        force=force,
        extension_override=extension_override,
    )


async def execute_plan(plan: ExecutionPlan) -> ExecutionReport:
    """Execute an ``ExecutionPlan`` and persist outputs.

    Args:
        plan: Planned prompt execution workload.

    Returns:
        Execution report with per-item statuses and written artifacts.

    Raises:
        ExecutionConfigError: If adapter resolution or output writing fails.
    """
    if not plan.items:
        return ExecutionReport(batch_results=[], written_files=[], output_dir=plan.output_dir)

    try:
        adapter = resolve_batch_adapter(plan.items[0].rendered)
    except ValueError as exc:
        raise ExecutionConfigError(str(exc), code='adapter_resolution_failed') from exc

    plan.output_dir.mkdir(parents=True, exist_ok=True)

    abort_event = asyncio.Event()
    loop = asyncio.get_running_loop()
    sigint_count = 0
    signal_handler_installed = False

    def _sigint_handler() -> None:
        nonlocal sigint_count
        sigint_count += 1
        if sigint_count == 1:
            logger.warning(
                'Interrupt received. Aborting after in-flight tasks finish... (press Ctrl+C again to force exit)',
            )
            abort_event.set()
            return
        raise KeyboardInterrupt

    try:
        loop.add_signal_handler(signal.SIGINT, _sigint_handler)
        signal_handler_installed = True
    except NotImplementedError:
        pass

    fatal_error: ExecutionConfigError | None = None
    written_files: list[Path] = []
    stems = {item.index: item.output_stem for item in plan.items}
    fmt = plan.fmt

    def _mark_fatal(message: str, *, code: str) -> None:
        nonlocal fatal_error
        if fatal_error is not None:
            return
        fatal_error = ExecutionConfigError(message, code=code)
        abort_event.set()

    def _write_or_mark_fatal(
        content: str | bytes,
        stem: str,
        *,
        extension_override: str | None = None,
    ) -> Path | None:
        try:
            path = _write_planned_output(
                content,
                plan.output_config,
                plan.output_dir,
                stem,
                force=plan.force,
                extension_override=extension_override,
            )
            written_files.append(path)
            return path
        except FileExistsError as exc:
            _mark_fatal(str(exc), code='output_exists')
            return None

    def _on_item_complete(result: BatchResult) -> None:
        if fatal_error is not None:
            return

        idx = result.index
        stem = stems[idx]
        if result.status == 'error':
            logger.error(
                'Batch item execution failed.',
                index=idx,
                attempts=result.attempts,
                elapsed=f'{result.elapsed:.3f}s',
                error=result.error,
            )
            return

        if result.response is None:
            _mark_fatal(f'Missing response for successful batch result at index={idx}.', code='missing_response')
            return

        try:
            item_content = extract_output_content(result.response, fmt)
        except (ValueError, AttributeError) as exc:
            if fmt == 'image':
                logger.warning(
                    'Failed to extract image response content.',
                    index=idx,
                    attempts=result.attempts,
                    elapsed=f'{result.elapsed:.3f}s',
                    error=str(exc),
                )
                return
            logger.warning('Failed to extract response content, writing empty output.', index=idx, error=str(exc))
            item_content = serialize_output('', fmt)
        _write_or_mark_fatal(item_content, stem)

    try:
        batch_results = await execute_batch(
            [item.rendered for item in plan.items],
            [item.variables for item in plan.items],
            adapter,
            max_workers=plan.runtime.max_workers if plan.runtime else 5,
            on_item_complete=_on_item_complete,
            runtime=plan.runtime,
            abort_event=abort_event,
        )
    finally:
        if signal_handler_installed:
            loop.remove_signal_handler(signal.SIGINT)

    if fatal_error is not None:
        raise fatal_error

    return ExecutionReport(
        batch_results=batch_results,
        written_files=written_files,
        output_dir=plan.output_dir,
    )
