"""
Budget Allocator — Token budget distribution across memory layers.

PRD §8 — Token Budget Allocation:
    Total: max_tokens = 2000 (configurable)

    Default allocation:
        experience:  800  (40%)
        context:     400  (20%)
        user_memory: 300  (15%)
        task_memory: 300  (15%)
        team_memory: 200  (10%)

    Dynamic adjustment by task type:
        deployment:  experience ↑ (+200), context ↓ (-200)
        debug:       context ↑ (+200), experience ↓ (-200)
        exploration: context ↑ (+150), experience ↓ (-150)
        execution:   experience ↑ (+150), context ↓ (-150)
"""

import logging
from typing import Dict

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

# Default token allocation (PRD §8)
_DEFAULT_ALLOC: Dict[str, int] = {
    "experience":  800,
    "context":     400,
    "user_memory": 300,
    "task_memory": 300,
    "team_memory": 200,
}

# Total = 2000

# Delta adjustments per task type (PRD §8 dynamic adjustment)
_TASK_ADJUSTMENTS: Dict[str, Dict[str, int]] = {
    "deployment": {
        "experience":  +200,
        "context":     -200,
    },
    "debugging": {
        "context":     +200,
        "experience":  -200,
    },
    "research": {
        "context":     +150,
        "experience":  -150,
    },
    "planning": {
        "context":     +100,
        "experience":  -100,
    },
    "coding": {
        "experience":  +150,
        "context":     -100,
        "task_memory": -50,
    },
}


class BudgetAllocator:
    """
    Allocates token budget across memory layers.

    Usage:
        allocator = BudgetAllocator(max_tokens=2000)
        budget = allocator.allocate("deployment")
        # → {"experience": 1000, "context": 200, "user_memory": 300, ...}
    """

    def __init__(self, max_tokens: int = 2000):
        self.max_tokens = max_tokens
        self._scale = max_tokens / 2000.0  # scale from default 2000

    def allocate(self, task_type: str = "coding") -> Dict[str, int]:
        """
        Return token budget per memory layer for the given task type.

        Budget values are scaled to max_tokens and never go below 50.
        """
        base = {k: int(v * self._scale) for k, v in _DEFAULT_ALLOC.items()}
        deltas = _TASK_ADJUSTMENTS.get(task_type, {})
        scaled_deltas = {k: int(v * self._scale) for k, v in deltas.items()}

        result = {}
        for layer in base:
            alloc = base[layer] + scaled_deltas.get(layer, 0)
            result[layer] = max(50, alloc)  # never below 50 tokens

        # Normalise to max_tokens
        total = sum(result.values())
        if total != self.max_tokens:
            # Adjust the largest bucket to compensate
            largest = max(result, key=result.get)
            result[largest] += self.max_tokens - total

        logger.debug("allocate task=%s budget=%s total=%d",
                     task_type, result, sum(result.values()))
        return result

    def remaining(self, budget: Dict[str, int], used: Dict[str, int]) -> Dict[str, int]:
        """Compute remaining budget per layer after actual usage."""
        return {layer: max(0, budget.get(layer, 0) - used.get(layer, 0))
                for layer in budget}

    def utilization(self, budget: Dict[str, int],
                    used: Dict[str, int]) -> float:
        """Compute token utilization rate (0.0 – 1.0). Target: > 85%."""
        total_budget = sum(budget.values())
        total_used = sum(used.values())
        if total_budget == 0:
            return 0.0
        return min(1.0, total_used / total_budget)
