"""Context builder for nathan-core.

Assembles: role prompt + ALL memories + task → one complete LLM-ready prompt.
Memories are curated and small — attach ALL, no search.

Architecture ref: ARCHITECTURE.md Section 6 (Complete Flow), CLAUDE.md (Key Decisions)

Memory ordering: global and namespace memories are interleaved (zip_longest) before
trimming. This distributes budget pressure proportionally — under tight token limits,
both types survive rather than globals being dropped first.

NOTE: Interleaving works best when global:namespace memory ratio is < 1:20.
Under extreme ratios (e.g. 2 global, 100 namespace), namespace memories may still
dominate the budget. This is a safety valve — curate memories to stay small.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from itertools import zip_longest

from netanel_core.config import NathanConfig
from netanel_core.memory_store import Memory, MemoryStore

logger = logging.getLogger(__name__)

# Approximate tokens per character (conservative estimate for English text).
# Used for the safety valve — not for billing.
_CHARS_PER_TOKEN = 4


@dataclass
class Context:
    """Assembled context ready for LLM invocation.

    Contains all components that form the complete prompt.
    """

    system_prompt: str
    memories_text: str
    task: str
    memories_count: int = 0
    estimated_tokens: int = 0
    trimmed: bool = False

    @property
    def full_system_prompt(self) -> str:
        """System prompt with memories injected."""
        if not self.memories_text:
            return self.system_prompt

        label = "memory" if self.memories_count == 1 else "memories"
        return (
            f"{self.system_prompt}\n\n"
            f"## Past Learnings ({self.memories_count} {label})\n\n"
            f"{self.memories_text}"
        )


def _estimate_tokens(text: str) -> int:
    """Rough token estimate from character count.

    Conservative (overestimates slightly) — used for safety valve only.
    For exact counts, use tiktoken. But we don't need exact for a safety check.
    """
    return len(text) // _CHARS_PER_TOKEN


def _format_memories(memories: list[Memory]) -> str:
    """Format memories into a text block for injection into the prompt."""
    if not memories:
        return ""

    sections = []
    for memory in memories:
        sections.append(f"### {memory.title}\n{memory.content}")

    return "\n\n".join(sections)


def _trim_memories(
    memories: list[Memory],
    max_tokens: int,
    system_prompt: str,
    task: str,
) -> list[Memory]:
    """Drop oldest memories first if total context exceeds max_tokens.

    This is a SAFETY VALVE, not normal operation. If we're hitting this,
    it means memories need manual curation.

    Strategy: drop from the front (oldest first, since memories are sorted
    by filename which is chronological).
    """
    # Reserve space for system prompt + task + formatting overhead
    reserved = _estimate_tokens(system_prompt) + _estimate_tokens(task) + 500
    available = max_tokens - reserved

    if available <= 0:
        # System prompt + task alone exceed limit — return empty
        logger.warning(
            "_trim_memories: reserved tokens (%d = prompt + task + 500 overhead) exceeds "
            "max_tokens=%d — all memories dropped. Curate memories or increase "
            "max_tokens_per_task.",
            reserved,
            max_tokens,
        )
        return []

    # Build up from most recent (reverse), stop when full
    kept: list[Memory] = []
    running_tokens = 0

    for memory in reversed(memories):
        mem_tokens = _estimate_tokens(f"### {memory.title}\n{memory.content}")
        if running_tokens + mem_tokens > available:
            break
        kept.append(memory)
        running_tokens += mem_tokens

    # Restore chronological order
    kept.reverse()
    return kept


def _interleave(globals_: list[Memory], namespace: list[Memory]) -> list[Memory]:
    """Interleave global and namespace memories in round-robin order.

    Uses zip_longest so both lists are fully consumed regardless of length.
    This ensures globals survive proportionally under budget pressure — not all
    dropped at once when namespace memories outnumber globals (BUG-039 fix).

    Example: [G1, G2], [N1, N2, N3] → [G1, N1, G2, N2, N3]

    Assumption: works best when global:namespace ratio is < 1:20.
    Under extreme ratios, namespace memories may still dominate tight budgets.
    """
    result: list[Memory] = []
    for pair in zip_longest(globals_, namespace):
        result.extend(m for m in pair if m is not None)
    return result


class ContextBuilder:
    """Assembles complete LLM context from files.

    Reads role prompt + ALL memories + task → produces Context object
    ready for LLM invocation.

    Usage:
        builder = ContextBuilder(config, store)
        context = builder.build("Write a sorting function")
    """

    def __init__(self, config: NathanConfig, store: MemoryStore) -> None:
        self._config = config
        self._store = store

    def build(self, task: str) -> Context:
        """Build complete context for an LLM call.

        Args:
            task: The user's task/instruction for this call.

        Returns:
            Context object with assembled prompt, memories, and metadata.
        """
        # 3a. Load role prompt
        system_prompt = self._load_prompt()

        # 3b. Load ALL memories (global + namespace) — interleaved for proportional
        # trimming under budget pressure (BUG-039 fix).
        global_memories = self._store.read_global_memories()
        namespace_memories = self._store.read_all_memories()
        memories = _interleave(global_memories, namespace_memories)

        # 3d. Token check — safety valve
        max_tokens = self._config.safety.max_tokens_per_task
        total_estimate = self._estimate_total(system_prompt, memories, task)
        trimmed = False

        if total_estimate > max_tokens:
            memories = _trim_memories(memories, max_tokens, system_prompt, task)
            trimmed = True

        # 3c. Assemble
        memories_text = _format_memories(memories)
        estimated_tokens = _estimate_tokens(system_prompt + memories_text + task)

        return Context(
            system_prompt=system_prompt,
            memories_text=memories_text,
            task=task,
            memories_count=len(memories),
            estimated_tokens=estimated_tokens,
            trimmed=trimmed,
        )

    def _load_prompt(self) -> str:
        """Load the current role prompt, or return a sensible default."""
        prompt = self._store.read_prompt()
        if prompt is not None:
            return prompt

        # Default prompt for cold start
        return (
            f"You are an AI assistant operating in the '{self._config.namespace}' role.\n"
            f"Complete the given task to the best of your ability.\n"
            f"Be precise, thorough, and explain your reasoning."
        )

    def _estimate_total(
        self,
        system_prompt: str,
        memories: list[Memory],
        task: str,
    ) -> int:
        """Estimate total token count for the assembled context."""
        memories_text = _format_memories(memories)
        total_text = system_prompt + memories_text + task
        return _estimate_tokens(total_text)
