"""ContextPacker: assembles curated context from memories, decisions, and sessions."""

from memotrail.core.searcher import Searcher, ChatSearchResult
from memotrail.storage import SQLiteStore, ChromaStore
from memotrail.core.embedder import Embedder
from memotrail.utils import config, get_logger
from memotrail.utils.tokens import count_tokens, truncate_to_tokens

logger = get_logger("memotrail.core.context_packer")


class ContextPacker:
    """Assembles curated context within a token budget.

    Packs decisions, memories, session summaries, and search results
    into a structured block that fits within a configurable token limit.
    """

    def __init__(
        self,
        sqlite_store: SQLiteStore | None = None,
        chroma_store: ChromaStore | None = None,
        searcher: Searcher | None = None,
        embedder: Embedder | None = None,
    ):
        self.sqlite = sqlite_store or SQLiteStore()
        self.chroma = chroma_store or ChromaStore()
        self.embedder = embedder or Embedder()
        self.searcher = searcher or Searcher(self.sqlite, self.chroma, self.embedder)

    def pack_project_context(
        self,
        project: str,
        budget: int | None = None,
    ) -> str:
        """Assemble curated context for a project.

        Includes decisions, memories, and recent session summaries
        proportionally allocated within the token budget.

        Args:
            project: Project name to gather context for
            budget: Total token budget (default from config)

        Returns:
            Formatted context string
        """
        budget = budget or config.context_budget_total
        sections: list[str] = []
        tokens_used = 0

        # 1. Decisions (up to context_budget_decisions tokens)
        decision_budget = min(config.context_budget_decisions, budget // 4)
        decisions_section = self._pack_decisions(project, decision_budget)
        if decisions_section:
            section_tokens = count_tokens(decisions_section)
            if tokens_used + section_tokens <= budget:
                sections.append(decisions_section)
                tokens_used += section_tokens

        # 2. Memories (up to 1/4 of budget)
        memory_budget = min(budget // 4, budget - tokens_used)
        memories_section = self._pack_memories(memory_budget)
        if memories_section:
            section_tokens = count_tokens(memories_section)
            if tokens_used + section_tokens <= budget:
                sections.append(memories_section)
                tokens_used += section_tokens

        # 3. Recent session summaries (remaining budget)
        session_budget = min(config.context_budget_recent, budget - tokens_used)
        sessions_section = self._pack_session_summaries(project, session_budget)
        if sessions_section:
            section_tokens = count_tokens(sessions_section)
            if tokens_used + section_tokens <= budget:
                sections.append(sessions_section)
                tokens_used += section_tokens

        if not sections:
            return f"No context available for project '{project}'."

        header = f"## Project Context: {project}\n\n"
        return header + "\n\n".join(sections)

    def pack_query_context(
        self,
        query: str,
        project: str | None = None,
        budget: int | None = None,
    ) -> str:
        """Retrieve semantically relevant past context for a query.

        Searches across conversations and memories, then assembles
        results within a token budget.

        Args:
            query: User prompt or question
            project: Optional project filter
            budget: Total token budget (default from config)

        Returns:
            Formatted context string
        """
        budget = budget or config.context_budget_total
        sections: list[str] = []
        tokens_used = 0

        # 1. Relevant decisions
        decision_budget = min(config.context_budget_decisions, budget // 5)
        if project:
            decisions_section = self._pack_decisions(project, decision_budget)
        else:
            decisions_section = self._pack_decisions(None, decision_budget)
        if decisions_section:
            section_tokens = count_tokens(decisions_section)
            if tokens_used + section_tokens <= budget:
                sections.append(decisions_section)
                tokens_used += section_tokens

        # 2. Search results from conversations
        retrieved_budget = min(config.context_budget_retrieved, budget - tokens_used)
        retrieved_section = self._pack_search_results(query, project, retrieved_budget)
        if retrieved_section:
            section_tokens = count_tokens(retrieved_section)
            if tokens_used + section_tokens <= budget:
                sections.append(retrieved_section)
                tokens_used += section_tokens

        # 3. Relevant memories
        memory_budget = min(budget // 5, budget - tokens_used)
        memories_section = self._pack_memory_search(query, memory_budget)
        if memories_section:
            section_tokens = count_tokens(memories_section)
            if tokens_used + section_tokens <= budget:
                sections.append(memories_section)
                tokens_used += section_tokens

        if not sections:
            return "No relevant past context found."

        header = "## Relevant Past Context\n\n"
        return header + "\n\n".join(sections)

    # ── Private helpers ─────────────────────────────────────

    def _pack_decisions(self, project: str | None, budget: int) -> str:
        """Pack decisions within budget."""
        decisions = self.sqlite.get_decisions(project=project, limit=20)
        if not decisions:
            return ""

        lines = ["### Decisions\n"]
        tokens_used = count_tokens(lines[0])

        for d in decisions:
            line = f"- [{d.category}] {d.decision_text}\n"
            line_tokens = count_tokens(line)
            if tokens_used + line_tokens > budget:
                break
            lines.append(line)
            tokens_used += line_tokens

        if len(lines) <= 1:
            return ""
        return "".join(lines)

    def _pack_memories(self, budget: int) -> str:
        """Pack all memories within budget."""
        # Get memories from SQLite
        rows = self.sqlite.conn.execute(
            "SELECT content, tags FROM memories ORDER BY updated_at DESC"
        ).fetchall()

        if not rows:
            return ""

        lines = ["### Memories\n"]
        tokens_used = count_tokens(lines[0])

        for row in rows:
            content = row["content"]
            line = f"- {content}\n"
            line_tokens = count_tokens(line)
            if tokens_used + line_tokens > budget:
                break
            lines.append(line)
            tokens_used += line_tokens

        if len(lines) <= 1:
            return ""
        return "".join(lines)

    def _pack_session_summaries(self, project: str | None, budget: int) -> str:
        """Pack recent session summaries within budget."""
        if project:
            rows = self.sqlite.conn.execute(
                "SELECT id, summary, started_at, message_count FROM sessions "
                "WHERE project = ? AND summary IS NOT NULL "
                "ORDER BY started_at DESC LIMIT 20",
                (project,),
            ).fetchall()
        else:
            rows = self.sqlite.conn.execute(
                "SELECT id, summary, started_at, message_count FROM sessions "
                "WHERE summary IS NOT NULL "
                "ORDER BY started_at DESC LIMIT 20",
            ).fetchall()

        if not rows:
            return ""

        lines = ["### Recent Sessions\n"]
        tokens_used = count_tokens(lines[0])

        for row in rows:
            line = f"- [{row['started_at'][:10]}] {row['summary']}\n"
            line_tokens = count_tokens(line)
            if tokens_used + line_tokens > budget:
                break
            lines.append(line)
            tokens_used += line_tokens

        if len(lines) <= 1:
            return ""
        return "".join(lines)

    def _pack_search_results(
        self, query: str, project: str | None, budget: int,
    ) -> str:
        """Pack semantic search results within budget."""
        results = self.searcher.search_chats(query, limit=10, project=project)
        if not results:
            return ""

        lines = ["### Relevant Conversations\n"]
        tokens_used = count_tokens(lines[0])

        for r in results:
            text = r.chunk_text
            # Truncate individual chunks if needed
            chunk_budget = min(budget // 3, budget - tokens_used)
            if count_tokens(text) > chunk_budget:
                text = truncate_to_tokens(text, chunk_budget)

            block = (
                f"**Session {r.session_id}** "
                f"(project: {r.project or 'unknown'}, score: {r.score:.2f})\n"
                f"{text}\n\n"
            )
            block_tokens = count_tokens(block)
            if tokens_used + block_tokens > budget:
                break
            lines.append(block)
            tokens_used += block_tokens

        if len(lines) <= 1:
            return ""
        return "".join(lines)

    def _pack_memory_search(self, query: str, budget: int) -> str:
        """Pack memory search results within budget."""
        results = self.searcher.search_memories(query, limit=10)
        if not results:
            return ""

        lines = ["### Relevant Memories\n"]
        tokens_used = count_tokens(lines[0])

        for r in results:
            line = f"- (score: {r.score:.2f}) {r.text}\n"
            line_tokens = count_tokens(line)
            if tokens_used + line_tokens > budget:
                break
            lines.append(line)
            tokens_used += line_tokens

        if len(lines) <= 1:
            return ""
        return "".join(lines)
