"""
prompt_pack.py — Epigenetic optimization: optimize prompts BEFORE touching weights.

A PromptPack is the compiled output of optimization — ready to deploy:
  - Optimized system instructions
  - Selected skills (highest fitness)
  - Few-shot examples (from best traces)
  - Tool policies
  - Output schema hints
  - Token budget compliance guaranteed
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
from purpose_agent.skills.schema import SkillCard
from purpose_agent.memory_homeostasis import MemoryBudget


@dataclass
class PromptPack:
    """Compiled prompt optimization output — deployable artifact."""
    name: str = "default"
    version: int = 1
    system_instructions: list[str] = field(default_factory=list)
    skills: list[dict[str, Any]] = field(default_factory=list)
    examples: list[dict[str, str]] = field(default_factory=list)  # {input, output}
    tool_policies: list[str] = field(default_factory=list)
    output_hints: list[str] = field(default_factory=list)
    token_estimate: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_system_prompt(self) -> str:
        """Compile into a single system prompt string."""
        parts = []
        if self.system_instructions:
            parts.append("## Instructions\n" + "\n".join(f"- {i}" for i in self.system_instructions))
        if self.skills:
            parts.append("## Skills\n" + "\n".join(
                f"- {s.get('trigger','')}: {s.get('procedure','')}" for s in self.skills[:5]))
        if self.examples:
            parts.append("## Examples")
            for i, ex in enumerate(self.examples[:3], 1):
                parts.append(f"### Example {i}\nInput: {ex.get('input','')}\nOutput: {ex.get('output','')}")
        if self.tool_policies:
            parts.append("## Tool Policies\n" + "\n".join(f"- {p}" for p in self.tool_policies))
        return "\n\n".join(parts)

    @property
    def total_chars(self) -> int:
        return len(self.to_system_prompt())


class PromptPackBuilder:
    """
    Builds optimized PromptPacks from skills, traces, and memory.
    
    Usage:
        builder = PromptPackBuilder(budget=MemoryBudget(max_injected_tokens=500))
        pack = builder.build(
            skills=active_skills,
            instructions=["Always validate input"],
            examples=[{"input": "fib(5)", "output": "5"}],
        )
        system_prompt = pack.to_system_prompt()
    """

    def __init__(self, budget: MemoryBudget | None = None):
        self.budget = budget or MemoryBudget()

    def build(
        self,
        skills: list[SkillCard] | None = None,
        instructions: list[str] | None = None,
        examples: list[dict[str, str]] | None = None,
        tool_policies: list[str] | None = None,
    ) -> PromptPack:
        """Build a token-budget-compliant PromptPack."""
        pack = PromptPack(
            system_instructions=instructions or [],
            tool_policies=tool_policies or [],
        )

        # Add skills sorted by fitness, under budget
        token_used = self.budget.estimate_tokens(pack.to_system_prompt())
        max_tokens = self.budget.max_injected_tokens

        if skills:
            sorted_skills = sorted(skills, key=lambda s: -s.fitness_score)
            for skill in sorted_skills:
                skill_text = f"{skill.trigger}: {' → '.join(skill.procedure[:3])}"
                skill_tokens = self.budget.estimate_tokens(skill_text)
                if token_used + skill_tokens > max_tokens:
                    break
                pack.skills.append({"trigger": skill.trigger, "procedure": skill_text,
                                    "fitness": skill.fitness_score})
                token_used += skill_tokens

        # Add examples under remaining budget
        if examples:
            for ex in examples[:5]:
                ex_text = f"{ex.get('input','')} {ex.get('output','')}"
                ex_tokens = self.budget.estimate_tokens(ex_text)
                if token_used + ex_tokens > max_tokens:
                    break
                pack.examples.append(ex)
                token_used += ex_tokens

        pack.token_estimate = token_used
        return pack