"""
Prompt Builder — Assembles final prompt from ranked, budgeted memory layers.

PRD §10 — Final Prompt Structure:
    SYSTEM (soul.md)
    USER PROFILE         ← user_memory
    CONTEXT SUMMARY      ← context_memory
    TASK MEMORY          ← task_memory
    EXPERIENCE           ← experience
    USER INPUT           ← original prompt

PRD §4 — Output:
    {
      "context": "...",
      "user_memory": "...",
      "task_memory": "...",
      "experience": "...",
      "token_usage": 1200
    }
"""

import logging
from typing import Dict, Optional
from openmemo.context_memory.context_budget import estimate_tokens

logger = logging.getLogger("openmemo.orchestrator")


class PromptBuilder:
    """
    Assembles the final injected prompt from memory layer blocks.

    Usage:
        builder = PromptBuilder()
        result = builder.build(
            user_input="Deploy my backend",
            user_memory_block="USER PROFILE\\n- prefers FastAPI",
            context_block="CONTEXT SUMMARY\\n- Goal: deploy",
            task_memory_block="TASK MEMORY\\n- service ready",
            experience_block="EXPERIENCE\\n- deploy|fastapi|docker",
            budget={"user_memory": 300, "context": 400, "experience": 800},
        )
        print(result["context"])   # the full injected prompt
        print(result["token_usage"])
    """

    def build(
        self,
        user_input: str,
        user_memory_block: str = "",
        context_block: str = "",
        task_memory_block: str = "",
        team_memory_block: str = "",
        experience_block: str = "",
        budget: Optional[Dict[str, int]] = None,
        system_block: str = "",
    ) -> dict:
        """
        Build the final prompt following PRD §10 order:
            SYSTEM → USER PROFILE → CONTEXT SUMMARY → TASK MEMORY
            → TEAM MEMORY → EXPERIENCE → USER INPUT

        Returns a structured dict (PRD §4 output format).
        """
        budget = budget or {}

        # Apply per-layer budget truncation
        user_memory_block = self._fit(user_memory_block,
                                       budget.get("user_memory", 300))
        context_block     = self._fit(context_block,
                                       budget.get("context", 400))
        task_memory_block = self._fit(task_memory_block,
                                       budget.get("task_memory", 300))
        team_memory_block = self._fit(team_memory_block,
                                       budget.get("team_memory", 200))
        experience_block  = self._fit(experience_block,
                                       budget.get("experience", 800))

        # Build ordered sections (PRD §10)
        sections = []
        if system_block:
            sections.append(system_block)
        if user_memory_block:
            sections.append(user_memory_block)
        if context_block:
            sections.append(context_block)
        if task_memory_block:
            sections.append(task_memory_block)
        if team_memory_block:
            sections.append(team_memory_block)
        if experience_block:
            sections.append(experience_block)
        # User input always last
        sections.append(user_input)

        full_prompt = "\n\n".join(sections)

        # Calculate token usage
        token_usage = estimate_tokens(full_prompt)
        layer_tokens = {
            "user_memory": estimate_tokens(user_memory_block),
            "context":     estimate_tokens(context_block),
            "task_memory": estimate_tokens(task_memory_block),
            "team_memory": estimate_tokens(team_memory_block),
            "experience":  estimate_tokens(experience_block),
        }

        return {
            "context":      full_prompt,     # complete injected prompt (PRD §4)
            "user_memory":  user_memory_block,
            "context_block": context_block,
            "task_memory":  task_memory_block,
            "team_memory":  team_memory_block,
            "experience":   experience_block,
            "user_input":   user_input,
            "token_usage":  token_usage,
            "layer_tokens": layer_tokens,
        }

    def _fit(self, text: str, max_tokens: int) -> str:
        """Truncate text to max_tokens (word-based approximation)."""
        if not text:
            return ""
        current = estimate_tokens(text)
        if current <= max_tokens:
            return text
        # _WORDS_PER_TOKEN = 0.75
        max_words = int(max_tokens * 0.75)
        words = text.split()
        if len(words) <= max_words:
            return text
        return " ".join(words[:max_words]) + " ..."
