"""Command-line interface for dotpromptz.

Usage::

    runprompt my_prompt.prompt
    runprompt batch_prompts.prompt
"""

import argparse
import asyncio
import json
import sys
from pathlib import Path
from typing import Any

import structlog
from dotenv import load_dotenv

try:
    import yaml
except ImportError:
    yaml = None  # type: ignore

from dotpromptz import Dotprompt
from dotpromptz.adapters import get_adapter, list_adapters
from dotpromptz.input_resolver import resolve_input
from dotpromptz.typing import AdapterConfig, DataArgument, PromptOutputConfig

logger = structlog.get_logger(__name__)

def _load_env_file(env_file: str | None) -> None:
    """Load a ``.env`` file if it exists.

    Args:
        env_file: Path to the env file. ``None`` means try ``.env`` in cwd.
    """
    path = Path(env_file) if env_file else Path('.env')
    if path.is_file():
        load_dotenv(path, override=False)


# Model-name prefixes used to auto-infer a suitable adapter.
_MODEL_ADAPTER_PREFIXES: list[tuple[str, str]] = [
    ('gpt-', 'openai'),
    ('o1', 'openai'),
    ('o3', 'openai'),
    ('o4', 'openai'),
    ('chatgpt', 'openai'),
    ('claude', 'anthropic'),
    ('gemini', 'google'),
]


def _infer_adapter_from_model(model: str | None) -> str | None:
    """Guess the adapter name from the model identifier.

    Args:
        model: The model string from the prompt frontmatter.

    Returns:
        An adapter name, or ``None`` if it cannot be inferred.
    """
    if not model:
        return None
    lower = model.lower()
    for prefix, adapter_name in _MODEL_ADAPTER_PREFIXES:
        if lower.startswith(prefix):
            return adapter_name
    return None


def _resolve_adapter(rendered: Any) -> Any:
    """Resolve adapter from RenderedPrompt's adapter field.

    Handles three scenarios:
    1. AdapterConfig object: use name and base_url/api_key
    2. String adapter name: use directly
    3. None: auto-infer from config.model

    Args:
        rendered: The RenderedPrompt with adapter configuration.

    Returns:
        An Adapter instance ready to generate responses.

    Raises:
        SystemExit: If no adapter can be determined.
    """
    adapter_cfg = rendered.adapter  # str | AdapterConfig | None
    adapter_name: str | None = None
    adapter_kwargs: dict[str, Any] = {}

    if isinstance(adapter_cfg, AdapterConfig):
        adapter_name = adapter_cfg.name
        if adapter_cfg.base_url:
            adapter_kwargs['base_url'] = adapter_cfg.base_url
        if adapter_cfg.api_key:
            adapter_kwargs['api_key'] = adapter_cfg.api_key
    elif isinstance(adapter_cfg, str):
        adapter_name = adapter_cfg
    else:
        # Auto-infer from config.model.
        config = rendered.config if isinstance(rendered.config, dict) else {}
        model_name = config.get('model')
        adapter_name = _infer_adapter_from_model(model_name)

    if not adapter_name:
        available = ', '.join(list_adapters())
        logger.error(
            'Cannot determine which adapter to use. '
            'Add "adapter: <name>" to the .prompt frontmatter, '
            'or use a model name that can be auto-detected.',
            available_adapters=available,
        )
        sys.exit(2)

    return get_adapter(adapter_name, **adapter_kwargs)


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.
    """
    from dotpromptz.typing import (
        DataPart,
        MediaPart,
        TextPart,
        ToolRequestPart,
        ToolResponsePart,
    )

    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.

    Args:
        content: Text or binary data to write.
        output_config: The output configuration.
        output_dir: Resolved output directory.
        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 _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.

    Raises:
        RuntimeError: If ``fmt`` is ``'yaml'`` but PyYAML is not installed.
    """
    if fmt == 'json':
        if isinstance(data, str):
            # Try to parse then re-serialize for pretty output
            try:
                data = json.loads(data)
            except (json.JSONDecodeError, TypeError):
                pass
        return json.dumps(data, indent=2, ensure_ascii=False)
    if fmt == 'yaml':
        if yaml is None:
            raise RuntimeError("'yaml' output format requires PyYAML: pip install pyyaml")
        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 cli(args: list[str] | None = None) -> None:
    """Run a .prompt file with input configured in frontmatter.

    Auto-detects single vs batch mode based on input in frontmatter:
    - No input or dict input -> single execution
    - List input -> batch execution with concurrency control

    Batch execution settings (max_workers, output_dir) are read from
    the .prompt file's 'runtime:' section.

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

    Examples::

      runprompt explain.prompt
      runprompt batch_prompts.prompt
    """
    structlog.configure(
        processors=[
            structlog.contextvars.merge_contextvars,
            structlog.processors.add_log_level,
            structlog.processors.TimeStamper(fmt='%Y-%m-%d %H:%M:%S'),
            structlog.dev.ConsoleRenderer(
                colors=False,
                pad_level=False,
            ),
        ],
        wrapper_class=structlog.make_filtering_bound_logger(20),
        logger_factory=lambda *_: structlog.PrintLogger(file=sys.stderr),
        cache_logger_on_first_use=False,
    )


    parser = argparse.ArgumentParser(
        prog='runprompt',
        description='Run a .prompt file with input configured in frontmatter.',
    )
    parser.add_argument(
        'prompt_file',
        metavar='PROMPT_FILE',
        help='Path to a .prompt file.',
    )
    parsed = parser.parse_args(args)

    prompt_file: str = parsed.prompt_file
    if not Path(prompt_file).is_file():
        logger.error('Path does not exist.', path=prompt_file)
        sys.exit(2)

    asyncio.run(_run_cli_async(prompt_file=prompt_file))


async def _run_cli_async(*, prompt_file: str) -> None:
    """Async implementation of the unified CLI command."""
    from dotpromptz.parse import parse_document
    from dotpromptz.version import PromptVersionError, adapt_prompt

    # 1. Load .env from cwd if present.
    _load_env_file(None)

    # 2. Read the prompt source and parse.
    prompt_path = Path(prompt_file)
    source = prompt_path.read_text(encoding='utf-8')
    parsed = parse_document(source)

    # 2b. 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, is_batch = 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. Auto-detect mode and execute.
    if not records and defaults:
        # defaults-only (no data records) → single execution with defaults
        await _run_single(source, defaults)
    elif not records:
        # No input at all → run with empty input (template variable
        # validation in RenderFunc.__call__ will catch missing vars).
        await _run_single(source, {})
    elif is_batch:
        # Batch execution mode
        await _run_batch(source, records)
    else:
        # Single execution mode
        await _run_single(source, records[0])


async def _run_single(
    source: str,
    input_data: dict,
) -> None:
    """Execute single prompt rendering and generation.

    All results are written to files according to the ``output`` config.
    Stdout is silent.

    Args:
        source: The .prompt file content.
        input_data: Dict of input variables.
    """
    # 1. Create Dotprompt instance and render.
    dp = Dotprompt()
    rendered = dp.render(source, data=DataArgument(input=input_data), variables=input_data)

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

    # 2. Resolve adapter.
    adapter = _resolve_adapter(rendered)

    # 3. Generate.
    try:
        response = await adapter.generate(rendered)
    except ImportError as exc:
        logger.error('Error.', error=str(exc))
        sys.exit(1)

    # 4. Determine output location.
    output_dir = Path(output_config.output_dir) if output_config.output_dir else Path('.')
    file_name = output_config.file_name or 'output'
    fmt = output_config.format

    # 5. Write output.
    if fmt == 'image':
        if response.image:
            path = _write_output(response.image.data, output_config, output_dir, file_name)
            logger.info('Image saved.', path=str(path))
        else:
            logger.error('Expected image output but model returned no image data.')
            sys.exit(1)
    else:
        # Text-based formats
        if response.text:
            content = _serialize_result(response.text, fmt)
        elif response.tool_calls:
            content = _serialize_result(
                [tc.model_dump() for tc in response.tool_calls], fmt,
            )
        else:
            content = ''
        path = _write_output(content, output_config, output_dir, file_name)
        logger.info('Output saved.', path=str(path))

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

    All results are written to files. When ``aggregate`` is enabled, text-based
    results are merged into a single file after all items complete.
    Stdout is silent.

    Args:
        source: The .prompt file content.
        input_list: List of input variable dicts.
    """
    # 1. Create Dotprompt instance.
    dp = Dotprompt()

    # 2. Parse the prompt once to get runtime and output config.
    from dotpromptz.parse import parse_document

    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)

    # 3. Get batch settings.
    max_workers = runtime_config.max_workers if runtime_config else 5
    output_dir = Path(output_config.output_dir) if output_config.output_dir else Path('.')
    file_name_tpl = output_config.file_name or 'output'
    fmt = output_config.format
    aggregate = output_config.aggregate

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

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

    # 6. Create semaphore for concurrency control.
    sem = asyncio.Semaphore(max_workers)

    # 7. Process one input item.
    async def _process_one(index: int, variables: dict[str, Any]) -> dict[str, Any]:
        async with sem:
            try:
                rendered = dp.render(source, data=DataArgument(input=variables), variables=variables)
                response = await adapter.generate(rendered)
                return {
                    'index': index,
                    'input': variables,
                    'status': 'success',
                    'text': response.text,
                    'image_data': response.image.data if response.image else None,
                    'tool_calls': [tc.model_dump() for tc in response.tool_calls] if response.tool_calls else None,
                    'error': None,
                }
            except Exception as exc:
                return {
                    'index': index,
                    'input': variables,
                    'status': 'error',
                    'text': None,
                    'image_data': None,
                    'tool_calls': None,
                    'error': str(exc),
                }

    # 8. Run all tasks concurrently.
    tasks = [_process_one(i, v) for i, v in enumerate(input_list)]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    # 9. Normalise results.
    final_results: list[dict[str, Any]] = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            final_results.append({
                'index': i,
                'input': input_list[i],
                'status': 'error',
                'text': None,
                'image_data': None,
                'tool_calls': None,
                'error': str(result),
            })
        else:
            final_results.append(result)
    final_results.sort(key=lambda r: r['index'])

    # 10. Write output files.
    if aggregate and fmt != 'image':
        # Aggregate: collect all text results into a single dict keyed by file stem.
        aggregated: dict[str, Any] = {}
        for result in final_results:
            stem = f'{file_name_tpl}_{result["index"]}'
            if result['status'] == 'error':
                aggregated[stem] = {'error': result['error']}
            elif result['text']:
                # Try to parse JSON text for structured formats
                if fmt in ('json', 'yaml'):
                    try:
                        aggregated[stem] = json.loads(result['text'])
                    except (json.JSONDecodeError, TypeError):
                        aggregated[stem] = result['text']
                else:
                    aggregated[stem] = result['text']
            elif result['tool_calls']:
                aggregated[stem] = result['tool_calls']
            else:
                aggregated[stem] = None
        content = _serialize_result(aggregated, fmt)
        path = _write_output(content, output_config, output_dir, file_name_tpl)
        logger.info('Aggregated output saved.', path=str(path))
    else:
        # Individual files per result.
        for result in final_results:
            idx = result['index']
            stem = f'{file_name_tpl}_{idx}'
            if fmt == 'image':
                if result['image_data']:
                    path = _write_output(result['image_data'], output_config, output_dir, stem)
                    logger.info('Image saved.', path=str(path))
                elif result['status'] == 'error':
                    logger.warning('Item failed.', index=idx, error=result['error'])
                else:
                    logger.warning('Expected image but got no image data.', index=idx)
            else:
                if result['status'] == 'error':
                    data: Any = {'error': result['error']}
                elif result['text']:
                    data = result['text']
                elif result['tool_calls']:
                    data = result['tool_calls']
                else:
                    data = ''
                content = _serialize_result(data, fmt)
                path = _write_output(content, output_config, output_dir, stem)
                logger.info('Output saved.', path=str(path))

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