"""Command-line interface for dotpromptz.

Usage::

    prompt run my_prompt.prompt
    prompt build my_prompt.prompt
"""

import argparse
import asyncio
import signal
import subprocess
import sys
from datetime import datetime
from importlib.metadata import version
from pathlib import Path
from typing import Any

import structlog
import yaml

from dotpromptz import Dotprompt
from dotpromptz.dotprompt import render_template_str
from dotpromptz.inputs import resolve_input
from dotpromptz.logging import configure_logging, generate_build_output_filename
from dotpromptz.parse import parse_document
from dotpromptz.runner import (
    aggregate_batch_results as _runner_aggregate_batch_results,
    extract_response_content as _runner_extract_response_content,
    part_to_dict as _runner_part_to_dict,
    resolve_adapter as _runner_resolve_adapter,
    run,
    serialize_result as _runner_serialize_result,
    write_output as _runner_write_output,
)
from dotpromptz.typing import DataArgument, PromptOutputConfig
from dotpromptz.version import PromptVersionError, adapt_prompt, check_for_update

logger = structlog.get_logger(__name__)


class _ReadableDumper(yaml.Dumper):
    """Custom YAML dumper that uses literal block style for multiline strings."""


def _readable_str_representer(dumper: yaml.Dumper, data: str) -> yaml.ScalarNode:
    """Represent multiline strings with literal block style (|) so newlines render as-is."""
    if '\n' in data:
        return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
    return dumper.represent_scalar('tag:yaml.org,2002:str', data)


_ReadableDumper.add_representer(str, _readable_str_representer)



def _resolve_adapter(rendered: Any) -> Any:
    """CLI wrapper around :func:`runner.resolve_adapter`.

    Calls ``sys.exit(2)`` on failure instead of raising ``ValueError``.
    """
    try:
        return _runner_resolve_adapter(rendered)
    except ValueError as exc:
        logger.error(str(exc))
        sys.exit(2)


def _part_to_dict(part: Any) -> dict[str, Any]:
    """CLI wrapper — delegates to :func:`runner.part_to_dict`."""
    return _runner_part_to_dict(part)


def _write_output(
    content: str | bytes,
    output_config: PromptOutputConfig,
    output_dir: Path,
    file_name: str,
) -> Path:
    """CLI wrapper — delegates to :func:`runner.write_output`."""
    return _runner_write_output(content, output_config, output_dir, file_name)


def _serialize_result(data: Any, fmt: str) -> str:
    """CLI wrapper — delegates to :func:`runner.serialize_result`."""
    return _runner_serialize_result(data, fmt)


def _build_rendered_output(rendered: Any) -> dict[str, Any]:
    """Build a structured dict from a rendered prompt for inspection output.

    Args:
        rendered: The rendered prompt object.

    Returns:
        A dict containing config, messages, output, tools, and input schemas.
    """
    output: dict[str, Any] = {}
    if rendered.config:
        output['config'] = rendered.config
    output['messages'] = []
    for msg in rendered.messages:
        msg_dict: dict[str, Any] = {'role': msg.role}
        if msg.content:
            parts = []
            for part in msg.content:
                parts.append(_part_to_dict(part))
            msg_dict['content'] = parts
        output['messages'].append(msg_dict)
    if rendered.output:
        output['output'] = rendered.output.model_dump(exclude_none=True)
    if rendered.tools:
        output['tools'] = [t.model_dump(exclude_none=True) for t in rendered.tools]
    if rendered.input:
        output['input'] = rendered.input.model_dump(exclude_none=True)
    return output


def _save_build_output(rendered: Any, prompt_path: Path, log_dir: Path) -> Path:
    """Save rendered prompt output to a YAML file and log the path.

    Args:
        rendered: The rendered prompt object.
        prompt_path: Path to the .prompt file.
        log_dir: Directory where the build output YAML is written.

    Returns:
        Path to the written file.
    """
    output = _build_rendered_output(rendered)
    output_filename = generate_build_output_filename(prompt_path.stem)
    output_file = log_dir / output_filename
    output_file.write_text(
        yaml.dump(
            output, Dumper=_ReadableDumper, allow_unicode=True,
            default_flow_style=False, sort_keys=False, width=120,
        ),
        encoding='utf-8',
    )
    logger.info(f'Build output saved \u2192 {output_file}')
    return output_file


def _validate_prompt_file(prompt_file: str) -> Path:
    """Validate that the prompt file exists and return its Path.

    Args:
        prompt_file: Path string to the .prompt file.

    Returns:
        Resolved Path object.
    """
    path = Path(prompt_file)
    if not path.is_file():
        logger.error('Path does not exist.', path=prompt_file)
        sys.exit(2)
    return path


def _get_version() -> str:
    """Return the installed package version."""
    try:
        return version('dotpromptz-py')
    except Exception:
        return 'unknown'


def _run_upgrade() -> None:
    """Upgrade dotpromptz-py to the latest version via uv."""
    cmd = ['uv', 'tool', 'upgrade', 'dotpromptz-py']
    logger.info('Running upgrade command.', command=' '.join(cmd))
    try:
        result = subprocess.run(cmd, check=False)  # noqa: S603
        sys.exit(result.returncode)
    except FileNotFoundError:
        logger.error('uv not found. Please install uv first: https://docs.astral.sh/uv/')
        sys.exit(1)


def cli(args: list[str] | None = None) -> None:
    """Main CLI entry point with subcommands.

    Subcommands:
        run   — Render and execute a .prompt file (call LLM).
        build — Render a .prompt file without executing (no LLM call).

    Args:
        args: Command-line arguments (defaults to sys.argv[1:]).

    Examples::

      prompt run explain.prompt
      prompt build explain.prompt
    """
    parser = argparse.ArgumentParser(
        prog='prompt',
        description='Dotpromptz CLI — build and run .prompt files.',
    )
    parser.add_argument('--version', action='version', version=f'%(prog)s {_get_version()}')
    parser.add_argument(
        '--log-level',
        default='INFO',
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
        help='Set logging level (default: INFO).',
    )
    parser.add_argument(
        '--upgrade',
        '-U',
        action='store_true',
        default=False,
        help='Upgrade dotpromptz-py to the latest version via uv.',
    )
    subparsers = parser.add_subparsers(dest='command', help='Available commands.')

    # --- prompt run ---
    run_parser = subparsers.add_parser(
        'run',
        help='Render and execute a .prompt file (calls LLM).',
        description='Render a .prompt file with input from frontmatter, then call the configured LLM and write output.',
    )
    run_parser.add_argument(
        'prompt_file',
        metavar='PROMPT_FILE',
        help='Path to a .prompt file.',
    )

    # --- prompt build ---
    build_parser = subparsers.add_parser(
        'build',
        help='Render a .prompt file without executing (no LLM call).',
        description=(
            'Parse and render a .prompt file, resolve input and config, '
            'then print the rendered prompt to stdout. '
            'Useful for inspecting the final prompt before sending it to an LLM.'
        ),
    )
    build_parser.add_argument(
        'prompt_file',
        metavar='PROMPT_FILE',
        help='Path to a .prompt file.',
    )

    parsed = parser.parse_args(args)

    # Configure logging AFTER parsing args so --log-level is available.
    log_dir = configure_logging(level=parsed.log_level)

    # Best-effort update check (never blocks or fails the CLI).
    check_for_update()

    # Handle --upgrade / -U before subcommand dispatch.
    if parsed.upgrade:
        _run_upgrade()
        return

    if parsed.command is None:
        parser.print_help()
        sys.exit(2)

    if parsed.command == 'run':
        _validate_prompt_file(parsed.prompt_file)
        asyncio.run(_run_cli_async(prompt_file=parsed.prompt_file, log_dir=log_dir))
    elif parsed.command == 'build':
        _validate_prompt_file(parsed.prompt_file)
        _build_cli(prompt_file=parsed.prompt_file, log_dir=log_dir)


def _build_cli(*, prompt_file: str, log_dir: Path) -> None:
    """Parse, validate, and render a .prompt file without calling an LLM.

    Writes the rendered prompt (messages, config, metadata) to a YAML file
    under the log directory and logs the file path at INFO level.

    Args:
        prompt_file: Path to the .prompt file.
        log_dir: Directory where the build output YAML is written.
    """
    # 1. Read and parse.
    prompt_path = Path(prompt_file)
    source = prompt_path.read_text(encoding='utf-8')
    parsed = parse_document(source)

    # 2. Version compatibility check.
    try:
        adapt_prompt(parsed)
    except PromptVersionError as exc:
        logger.error(str(exc))
        sys.exit(2)

    # 3. Resolve input data from frontmatter.
    try:
        input_cfg = parsed.input
        defaults = input_cfg.defaults if input_cfg else None
        records = resolve_input(
            input_cfg.data if input_cfg else None,
            prompt_path,
            defaults=defaults,
        )
    except ValueError as exc:
        logger.error('Input resolution failed.', error=str(exc))
        sys.exit(1)
    # 4. Determine input data for rendering.
    if not records and defaults:
        input_data = defaults
    elif not records:
        input_data = {}
    elif len(records) > 1:
        # For batch, render the first item as a representative sample.
        input_data = records[0]
        logger.info('Batch mode detected. Rendering first item as sample.', total_items=len(records))
    else:
        input_data = records[0]
    # 5. Render (no LLM call).
    dp = Dotprompt()
    rendered = dp.render(source, data=DataArgument(input=input_data), variables=input_data)

    # 6. Save build output to YAML file.
    _save_build_output(rendered, prompt_path, log_dir)


async def _run_cli_async(*, prompt_file: str, log_dir: Path) -> None:
    """Async implementation of the run command.

    Always runs through ``_run_batch`` \u2014 single-item input is treated
    as a batch of one.
    """
    # 1. Read the prompt source and parse.
    prompt_path = Path(prompt_file)
    source = prompt_path.read_text(encoding='utf-8')
    parsed = parse_document(source)

    # 2. Version compatibility check (fail early before input resolution).
    try:
        adapt_prompt(parsed)
    except PromptVersionError as exc:
        logger.error(str(exc))
        sys.exit(2)

    # 3. Resolve input data from frontmatter.
    try:
        input_cfg = parsed.input
        defaults = input_cfg.defaults if input_cfg else None
        records = resolve_input(
            input_cfg.data if input_cfg else None,
            prompt_path,
            defaults=defaults,
        )
    except ValueError as exc:
        logger.error('Input resolution failed.', error=str(exc))
        sys.exit(1)

    # 4. Build input list (always a list).
    if not records and defaults:
        input_list = [defaults]
    elif not records:
        input_list = [{}]
    else:
        input_list = records

    # 5. Save build output (same as build command).
    dp = Dotprompt()
    first_input = input_list[0]
    rendered = dp.render(source, data=DataArgument(input=first_input), variables=first_input)
    _save_build_output(rendered, prompt_path, log_dir)

    # 6. Execute as batch (single item = batch of 1).
    await _run_batch(source, input_list, prompt_path)


async def _run_batch(
    source: str,
    input_list: list[dict],
    prompt_path: Path,
) -> None:
    """Execute batch prompt rendering and generation with concurrency control.

    Delegates core batch execution to ``runner.run`` and handles
    file I/O, aggregation, SIGINT handling, and exit codes.

    All results are written to files. When ``aggregate`` is enabled,
    results are merged into a single file after all items complete.
    For ``txt`` format aggregation, the output is serialized as JSON.
    Stdout is silent.

    Args:
        source: The .prompt file content.
        input_list: List of input variable dicts.
        prompt_path: Path to the .prompt file (for output dir template resolution).
    """
    # 1. Parse the prompt once to get runtime and output config.
    parsed_prompt = parse_document(source)
    runtime_config = parsed_prompt.runtime  # RuntimeConfig | None
    output_config: PromptOutputConfig | None = parsed_prompt.output

    if output_config is None:
        logger.error('No output configuration found in .prompt frontmatter. output.format is required.')
        sys.exit(2)

    if not output_config.file_name:
        logger.error('output.file_name is required in .prompt frontmatter.')
        sys.exit(2)

    # 2. Get batch settings.
    max_workers = runtime_config.max_workers if runtime_config else 5
    output_dir = _resolve_output_dir(output_config.output_dir, prompt_path)
    file_name_tpl = output_config.file_name
    fmt = output_config.format
    aggregate = output_config.aggregate

    # 3. Render first item to resolve adapter.
    dp = Dotprompt()
    first_rendered = dp.render(source, data=DataArgument(input=input_list[0]), variables=input_list[0])
    adapter = _resolve_adapter(first_rendered)

    # 4. Create output directory.
    output_dir.mkdir(parents=True, exist_ok=True)

    # 5. Set up SIGINT (Ctrl+C) handler for graceful abort.
    abort_event = asyncio.Event()
    loop = asyncio.get_running_loop()
    sigint_count = 0

    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()
        else:
            logger.warning('Force exit.')
            sys.exit(130)

    try:
        loop.add_signal_handler(signal.SIGINT, _sigint_handler)
    except NotImplementedError:
        # Windows does not support add_signal_handler; fall back to no-op.
        pass

    # 6. Run batch using runner.
    batch_results = await run(
        source,
        input_list,
        adapter,
        max_workers=max_workers,
        dp=dp,
        runtime=runtime_config,
    )

    # 7. Write output files.
    is_single = len(batch_results) == 1

    if aggregate and fmt != 'image':
        # Aggregate: collect all text results into a single dict keyed by resolved file stem.
        aggregated = _runner_aggregate_batch_results(
            batch_results,
            fmt,
            file_name_tpl,
            prompt_path=prompt_path,
        )
        # For txt aggregate, serialize as JSON.
        agg_fmt = 'json' if fmt == 'txt' else fmt
        content = _serialize_result(aggregated, agg_fmt)
        # Resolve the aggregate output file name (use first item's vars as context).
        agg_stem = _resolve_file_name(
            file_name_tpl,
            prompt_path,
            item_vars=batch_results[0].input if batch_results else None,
        )
        # Write with the appropriate extension.
        agg_output_config = (
            output_config
            if fmt != 'txt'
            else PromptOutputConfig(
                format='json',
                file_name=agg_stem,
                aggregate=True,
            )
        )
        path = _write_output(content, agg_output_config, output_dir, agg_stem)
        logger.info(f'Aggregated output saved \u2192 {path}')
    else:
        # Individual files per result.
        for result in batch_results:
            idx = result.index
            # Resolve file_name template with item variables and auto-vars.
            stem = _resolve_file_name(
                file_name_tpl,
                prompt_path,
                item_vars=result.input,
                index=None if is_single else idx,
            )
            if result.status == 'error':
                if fmt == 'image':
                    logger.warning('Item failed.', index=idx, error=result.error)
                else:
                    err_content = _serialize_result({'error': result.error}, fmt)
                    path = _write_output(err_content, output_config, output_dir, stem)
                    logger.info(f'Output saved \u2192 {path}')
                continue
            try:
                item_content = _runner_extract_response_content(result.response, fmt)
            except (ValueError, AttributeError):
                if fmt == 'image':
                    logger.warning('Expected image but got no image data.', index=idx)
                else:
                    item_content = _serialize_result('', fmt)
                    path = _write_output(item_content, output_config, output_dir, stem)
                    logger.info(f'Output saved \u2192 {path}')
                continue
            path = _write_output(item_content, output_config, output_dir, stem)
            if fmt == 'image':
                logger.info(f'Image saved \u2192 {path}')
            else:
                logger.info(f'Output saved \u2192 {path}')

    # 8. Exit with error code if ALL items failed.
    failed = sum(1 for r in batch_results if r.status == 'error')
    if failed == len(batch_results):
        sys.exit(1)


def _resolve_output_dir(output_dir_tpl: str | None, prompt_path: Path) -> Path:
    """Resolve output directory from template or default.

    Template variables:
        ``{{NAME}}`` — the ``.prompt`` file stem (without extension).
        ``{{TIME}}`` — current timestamp as ``YYYYMMDD_HHMMSS``.

    When *output_dir_tpl* is ``None``, defaults to
    ``output/{{NAME}}_{{TIME}}``.

    Args:
        output_dir_tpl: Output directory template string, or ``None``.
        prompt_path: Path to the .prompt file.

    Returns:
        Resolved ``Path`` object.
    """
    if output_dir_tpl is None:
        output_dir_tpl = 'output/{{NAME}}_{{TIME}}'

    ctx: dict[str, Any] = {
        'NAME': prompt_path.stem,
        'TIME': datetime.now().strftime('%Y%m%d_%H%M%S'),
    }
    return Path(render_template_str(output_dir_tpl, ctx))


def _resolve_file_name(
    file_name_tpl: str,
    prompt_path: Path,
    *,
    item_vars: dict[str, Any] | None = None,
    index: int | None = None,
) -> str:
    """Resolve a ``file_name`` template with auto-variables and input variables.

    UPPERCASE auto-variables (``NAME``, ``TIME``, ``INDEX``) are merged with
    input variables into a single context dict.  No name shadowing occurs
    because auto-vars are UPPERCASE and user input vars are lowercase.

    Args:
        file_name_tpl: The raw file_name string from frontmatter (may contain ``{{…}}``).
        prompt_path: Path to the ``.prompt`` file.
        item_vars: Input variables dict for the current batch item.
        index: Zero-based batch item index (omit for single-item runs).

    Returns:
        The 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)
