"""Planner configuration for {{ agent_name }}."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import asyncio

from penguiflow.catalog import build_catalog
from penguiflow.artifacts import ArtifactRetentionConfig, InMemoryArtifactStore
from penguiflow.planner import ReactPlanner{% if with_background_tasks %}, BackgroundTasksConfig{% endif %}{% if reflection_enabled %}, ReflectionConfig{% if reflection_criteria %}, ReflectionCriteria{% endif %}{% endif %}

from penguiflow.planner.memory import MemoryBudget, MemoryIsolation, ShortTermMemoryConfig
from penguiflow.rich_output import DEFAULT_ALLOWLIST, RichOutputConfig, attach_rich_output_nodes, get_runtime
{% if with_background_tasks %}
from penguiflow.sessions.task_tools import build_task_tool_specs
{% endif %}

from .config import Config
from .tools import build_catalog_bundle
{% if has_external_tools %}
from .external_tools import get_external_tool_specs, setup_external_tools
{% endif %}


SYSTEM_PROMPT_EXTRA = """{{ system_prompt_extra }}"""
{% if include_memory_prompt %}
MEMORY_PROMPT = """{{ memory_prompt }}"""
{% endif %}


@dataclass
class PlannerBundle:
    """Container for planner and registry."""

    planner: ReactPlanner
    registry: Any


def _build_system_prompt(config: Config, rich_output_prompt: str = "") -> str:
    """Build the system prompt, optionally including memory and rich output context."""
    prompt = SYSTEM_PROMPT_EXTRA
{% if include_memory_prompt %}
    if config.memory_enabled:
        prompt = prompt + "\n\n" + MEMORY_PROMPT
{% endif %}
    if rich_output_prompt:
        prompt = prompt + "\n\n" + rich_output_prompt
    return prompt


def _build_rich_output_config(config: Config) -> RichOutputConfig:
    """Build rich output config for component rendering."""
    allowlist = config.rich_output_allowlist or list(DEFAULT_ALLOWLIST)
    return RichOutputConfig(
        enabled=config.rich_output_enabled,
        allowlist=allowlist,
        include_prompt_catalog=config.rich_output_include_prompt_catalog,
        include_prompt_examples=config.rich_output_include_prompt_examples,
        max_payload_bytes=config.rich_output_max_payload_bytes,
        max_total_bytes=config.rich_output_max_total_bytes,
    )


def _build_short_term_memory(config: Config) -> ShortTermMemoryConfig | None:
    """Build built-in short-term memory configuration."""
    if not config.short_term_memory_enabled:
        return None

    budget = MemoryBudget(
        full_zone_turns=config.short_term_memory_full_zone_turns,
        summary_max_tokens=config.short_term_memory_summary_max_tokens,
        total_max_tokens=config.short_term_memory_total_max_tokens,
        overflow_policy=config.short_term_memory_overflow_policy,
    )
    isolation = MemoryIsolation(
        tenant_key=config.short_term_memory_tenant_key,
        user_key=config.short_term_memory_user_key,
        session_key=config.short_term_memory_session_key,
        require_explicit_key=config.short_term_memory_require_explicit_key,
    )

    return ShortTermMemoryConfig(
        strategy=config.short_term_memory_strategy,
        budget=budget,
        isolation=isolation,
        summarizer_model=config.short_term_memory_summarizer_model,
        include_trajectory_digest=config.short_term_memory_include_trajectory_digest,
        recovery_backlog_limit=config.short_term_memory_recovery_backlog_limit,
        retry_attempts=config.short_term_memory_retry_attempts,
        retry_backoff_base_s=config.short_term_memory_retry_backoff_base_s,
        degraded_retry_interval_s=config.short_term_memory_degraded_retry_interval_s,
    )

{% if with_background_tasks %}

def _build_background_tasks_config(config: Config) -> BackgroundTasksConfig | None:
    if not config.background_tasks_enabled:
        return None
    return BackgroundTasksConfig(
        enabled=True,
        include_prompt_guidance=config.background_tasks_include_prompt_guidance,
        allow_tool_background=config.background_tasks_allow_tool_background,
        default_mode=config.background_tasks_default_mode,
        default_merge_strategy=config.background_tasks_default_merge_strategy,
        context_depth=config.background_tasks_context_depth,
        propagate_on_cancel=config.background_tasks_propagate_on_cancel,
        spawn_requires_confirmation=config.background_tasks_spawn_requires_confirmation,
        max_concurrent_tasks=config.background_tasks_max_concurrent_tasks,
        max_tasks_per_session=config.background_tasks_max_tasks_per_session,
        task_timeout_s=config.background_tasks_task_timeout_s,
        max_pending_steering=config.background_tasks_max_pending_steering,
    )

{% endif %}


def _build_artifact_store(config: Config) -> InMemoryArtifactStore | None:
    """Build artifact store from config.

    When disabled, ReactPlanner falls back to NoOpArtifactStore.
    """
    if not config.artifact_store_enabled:
        return None

    return InMemoryArtifactStore(
        retention=ArtifactRetentionConfig(
            ttl_seconds=config.artifact_store_ttl_seconds,
            max_artifact_bytes=config.artifact_store_max_artifact_bytes,
            max_session_bytes=config.artifact_store_max_session_bytes,
            max_trace_bytes=config.artifact_store_max_trace_bytes,
            max_artifacts_per_trace=config.artifact_store_max_artifacts_per_trace,
            max_artifacts_per_session=config.artifact_store_max_artifacts_per_session,
            cleanup_strategy=config.artifact_store_cleanup_strategy,
        ),
    )


def build_planner(config: Config, *, event_callback=None) -> PlannerBundle:
    """Build ReactPlanner with generated tool catalog."""

    nodes, registry = build_catalog_bundle()
    rich_output_config = _build_rich_output_config(config)
    nodes.extend(attach_rich_output_nodes(registry, config=rich_output_config))
    catalog = build_catalog(nodes, registry)
{% if has_external_tools %}
    # Connect external tool sources and merge their specs using the current event loop
    loop = asyncio.get_event_loop()
    if loop.is_running():
        raise RuntimeError("build_planner must be called before the event loop starts when external tools are configured")
    else:
        external_nodes = loop.run_until_complete(setup_external_tools(registry))
    catalog.extend(get_external_tool_specs(external_nodes))
{% endif %}

{% if with_background_tasks %}
    background_tasks_config = _build_background_tasks_config(config)
    if background_tasks_config is not None:
        task_tool_specs = build_task_tool_specs()
        catalog.extend(task_tool_specs)
{% endif %}

{% if reflection_enabled %}
    reflection_cfg = ReflectionConfig(
        enabled=True,
        quality_threshold=config.reflection_quality_threshold,
        max_revisions=config.reflection_max_revisions,
{% if reflection_criteria %}
        criteria=ReflectionCriteria(
            completeness="{{ reflection_criteria.completeness }}",
            accuracy="{{ reflection_criteria.accuracy }}",
            clarity="{{ reflection_criteria.clarity }}",
        ),
{% endif %}
    )
{% else %}
    reflection_cfg = None
{% endif %}

{% if summarizer_enabled %}
    summarizer_llm = config.summarizer_model or config.llm_model
{% else %}
    summarizer_llm = None
{% endif %}

{% if reflection_enabled %}
    reflection_llm = config.reflection_model or config.llm_model
{% else %}
    reflection_llm = None
{% endif %}

    artifact_store = _build_artifact_store(config)

    rich_output_prompt = get_runtime().prompt_section()

    planner = ReactPlanner(
        llm=config.llm_model,
        catalog=catalog,
        system_prompt_extra=_build_system_prompt(config, rich_output_prompt),
        max_iters=config.planner_max_iters,
        hop_budget=config.planner_hop_budget,
        absolute_max_parallel=config.planner_absolute_max_parallel,
        multi_action_sequential=config.planner_multi_action_sequential,
        multi_action_read_only_only=config.planner_multi_action_read_only_only,
        multi_action_max_tools=config.planner_multi_action_max_tools,
        summarizer_llm=summarizer_llm,
        reflection_config=reflection_cfg,
        reflection_llm=reflection_llm,
        planning_hints={{ planning_hints_literal }},
        artifact_store=artifact_store,
        event_callback=event_callback,
        stream_final_response=config.planner_stream_final_response,
        short_term_memory=_build_short_term_memory(config),
{% if with_background_tasks %}
        background_tasks=background_tasks_config,
{% endif %}
    )

    return PlannerBundle(planner=planner, registry=registry)
