"""Execution planning for compiled prompts."""

from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

import jinja2

from dotpromptz.compiler import CompiledPrompt, render_compiled_prompt
from dotpromptz.errors import ExecutionConfigError, InputVariableNamingError
from dotpromptz.typing import PromptOutputConfig, RenderedPrompt, RuntimeConfig

_AUTO_VARIABLE_PATTERN = re.compile(r'^[A-Z][A-Z0-9_]*$')
_NUMBER_VARIABLE_PATTERN = re.compile(r'(?<![A-Za-z0-9_])NUMBER(?![A-Za-z0-9_])')


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


@dataclass(frozen=True)
class ExecutionPlanItem:
    """One planned execution unit for batch processing."""

    index: int
    variables: dict[str, Any]
    rendered: RenderedPrompt
    output_stem: str
    output_path: Path


@dataclass(frozen=True)
class ExecutionPlan:
    """Concrete execution plan derived from a compiled prompt."""

    compiled: CompiledPrompt
    output_config: PromptOutputConfig
    runtime: RuntimeConfig | None
    output_dir: Path
    fmt: str
    batch_time: str
    force: bool
    items: list[ExecutionPlanItem]


def _validate_input_variable_keys(keys: set[str]) -> None:
    """Validate output-template input vars against ALL_CAPS reservation."""
    invalid = {name for name in keys if _AUTO_VARIABLE_PATTERN.fullmatch(name) is not None}
    if invalid:
        raise InputVariableNamingError(invalid)


def _render_template_str(
    template: str,
    context: dict[str, Any],
    *,
    globals_: dict[str, Any] | None = None,
) -> str:
    """Render a simple Jinja2 template string for output path/name templates."""
    if '{{' not in template and '{%' not in template:
        return template
    env = jinja2.Environment(autoescape=False, undefined=jinja2.Undefined)
    if globals_:
        env.globals.update(dict(globals_))
    return env.from_string(template).render(context)


def _uses_mumber(template: str | None) -> bool:
    """Return whether a template references the ``NUMBER`` auto variable."""
    if not template:
        return False
    return _NUMBER_VARIABLE_PATTERN.search(template) is not None


def _template_globals(
    prompt_path: Path,
    *,
    time_value: str | None = None,
    index: int | None = None,
    mumber: int | None = None,
) -> dict[str, str]:
    """Build output-template globals for output path/name rendering."""
    globals_: dict[str, str] = {
        'NAME': prompt_path.stem,
        'NUMBER': str(mumber or 1),
        'TIME': time_value or datetime.now().strftime('%Y%m%d_%H%M%S'),
    }
    if index is not None:
        globals_['INDEX'] = str(index)
    return globals_


def _render_output_template(
    template: str,
    prompt_path: Path,
    *,
    item_vars: dict[str, Any] | None = None,
    index: int | None = None,
    time_value: str | None = None,
    mumber: int | None = None,
) -> str:
    """Render an output template with auto globals and item variables."""
    input_vars = item_vars or {}
    _validate_input_variable_keys(set(input_vars.keys()))
    return _render_template_str(
        template,
        input_vars,
        globals_=_template_globals(prompt_path, time_value=time_value, index=index, mumber=mumber),
    )


def _resolve_output_dir(
    output_dir_tpl: str | None,
    prompt_path: Path,
    *,
    time_value: str | None = None,
    mumber: int | None = None,
) -> Path:
    """Resolve output directory from template or default value."""
    if output_dir_tpl is None:
        output_dir_tpl = 'output/{{NAME}}_{{TIME}}'

    resolved_dir = Path(
        _render_output_template(
            output_dir_tpl,
            prompt_path,
            time_value=time_value,
            mumber=mumber,
        )
    )
    if resolved_dir.is_absolute():
        return resolved_dir
    return prompt_path.parent / resolved_dir


def _resolve_file_name(
    file_name_tpl: str,
    prompt_path: Path,
    *,
    item_vars: dict[str, Any] | None = None,
    index: int | None = None,
    time_value: str | None = None,
    mumber: int | None = None,
) -> str:
    """Resolve output file name template using input vars and auto globals."""
    return _render_output_template(
        file_name_tpl,
        prompt_path,
        item_vars=item_vars,
        index=index,
        time_value=time_value,
        mumber=mumber,
    )


def _resolve_unique_output_target(
    file_name_tpl: str,
    prompt_path: Path,
    *,
    output_dir: Path,
    ext: str,
    item_vars: dict[str, Any] | None = None,
    index: int | None = None,
    time_value: str | None = None,
    seen_paths: dict[Path, int],
    force: bool,
    start_mumber: int = 1,
) -> tuple[str, Path, int | None]:
    """Resolve one output file target, incrementing ``NUMBER`` when needed."""
    if not _uses_mumber(file_name_tpl):
        stem = _resolve_file_name(
            file_name_tpl,
            prompt_path,
            item_vars=item_vars,
            index=index,
            time_value=time_value,
        )
        return stem, output_dir / f'{stem}.{ext}', None

    mumber = max(1, start_mumber)
    while True:
        stem = _resolve_file_name(
            file_name_tpl,
            prompt_path,
            item_vars=item_vars,
            index=index,
            time_value=time_value,
            mumber=mumber,
        )
        target = output_dir / f'{stem}.{ext}'
        if target in seen_paths or (target.exists() and not force):
            mumber += 1
            continue
        return stem, target, mumber


def _plan_output_items(
    compiled: CompiledPrompt,
    *,
    file_name_tpl: str,
    output_dir: Path,
    ext: str,
    batch_time: str,
    force: bool,
    dir_mumber: int | None,
) -> list[ExecutionPlanItem] | None:
    """Plan output items for one resolved output directory.

    Returns ``None`` when an existing file collision should retry with the next
    directory ``NUMBER`` candidate.
    """
    is_single = len(compiled.input_list) == 1
    seen_paths: dict[Path, int] = {}
    items: list[ExecutionPlanItem] = []
    retry_on_existing_file = dir_mumber is not None and not _uses_mumber(file_name_tpl)

    for idx, variables in enumerate(compiled.input_list):
        stem, target, file_mumber = _resolve_unique_output_target(
            file_name_tpl,
            compiled.prompt_path,
            output_dir=output_dir,
            ext=ext,
            item_vars=variables,
            index=None if is_single else idx,
            time_value=batch_time,
            seen_paths=seen_paths,
            force=force,
            start_mumber=dir_mumber or 1,
        )
        previous_idx = seen_paths.get(target)
        if previous_idx is not None:
            raise ExecutionConfigError(
                (
                    'Duplicate output target detected in batch. '
                    f'index={idx}, previous_index={previous_idx}, path={target}'
                ),
                code='output_target_duplicate',
            )
        seen_paths[target] = idx
        if target.exists() and not force:
            if retry_on_existing_file:
                return None
            raise ExecutionConfigError(
                f'Output file already exists: {target}. Use --force to overwrite.',
                code='output_exists',
            )
        render_mumber = file_mumber or dir_mumber
        rendered = render_compiled_prompt(
            compiled,
            variables,
            auto_vars_override={'NUMBER': str(render_mumber)} if render_mumber is not None else None,
        )
        items.append(
            ExecutionPlanItem(
                index=idx,
                variables=variables,
                rendered=rendered,
                output_stem=stem,
                output_path=target,
            )
        )

    return items


def plan_prompt(
    compiled: CompiledPrompt,
    *,
    force: bool = False,
    batch_time: str | None = None,
) -> ExecutionPlan:
    """Build an execution plan from a compiled prompt.

    Args:
        compiled: Compiled prompt object.
        force: Whether overwrite checks should be skipped.
        batch_time: Optional fixed timestamp used for deterministic planning.

    Returns:
        Execution plan with rendered items and output targets.

    Raises:
        ExecutionConfigError: If output configuration is invalid.
    """
    parsed_prompt = compiled.parsed_prompt
    runtime_config = parsed_prompt.runtime
    output_config = parsed_prompt.output

    if output_config is None:
        raise ExecutionConfigError(
            'No output configuration found in .prompt frontmatter. output.format is required.',
            code='output_config_missing',
        )
    if not output_config.file_name:
        raise ExecutionConfigError(
            'output.file_name is required in .prompt frontmatter.',
            code='output_filename_missing',
        )

    effective_batch_time = batch_time or datetime.now().strftime('%Y%m%d_%H%M%S')
    file_name_tpl = output_config.file_name
    fmt = output_config.format
    ext = _resolve_output_extension(fmt)
    if _uses_mumber(output_config.output_dir):
        dir_mumber = 1
        while True:
            output_dir = _resolve_output_dir(
                output_config.output_dir,
                compiled.prompt_path,
                time_value=effective_batch_time,
                mumber=dir_mumber,
            )
            items = _plan_output_items(
                compiled,
                file_name_tpl=file_name_tpl,
                output_dir=output_dir,
                ext=ext,
                batch_time=effective_batch_time,
                force=force,
                dir_mumber=dir_mumber,
            )
            if items is not None:
                break
            dir_mumber += 1
    else:
        dir_mumber = None
        output_dir = _resolve_output_dir(
            output_config.output_dir,
            compiled.prompt_path,
            time_value=effective_batch_time,
        )
        items = _plan_output_items(
            compiled,
            file_name_tpl=file_name_tpl,
            output_dir=output_dir,
            ext=ext,
            batch_time=effective_batch_time,
            force=force,
            dir_mumber=dir_mumber,
        )
        assert items is not None

    return ExecutionPlan(
        compiled=compiled,
        output_config=output_config,
        runtime=runtime_config,
        output_dir=output_dir,
        fmt=fmt,
        batch_time=effective_batch_time,
        force=force,
        items=items,
    )
