"""
cortex/hippo/expansion.py — DAG expansion orchestrator with global token budget.

expand(): Walk summaryIds through RetrievalEngine with global token cap.
describe_and_expand(): Grep first, then expand top matches.
distill_for_subagent(): Format ExpansionResult as compact text payload.

Issue: #807
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Callable, Awaitable, Literal, Optional

from cortex.hippo.config import HippoConfig
from cortex.hippo.retrieval import (
    ChildItem,
    ExpandInput,
    GrepInput,
    MessageItem,
    RetrievalEngine,
)

logger = logging.getLogger(__name__)

# Type alias for the LLM complete function used by expand_query
LLMCompleteFn = Callable[..., Awaitable[dict]]


# ═══════════════════════════════════════════════════════════════════════════
# Data types
# ═══════════════════════════════════════════════════════════════════════════


@dataclass
class ExpansionRequest:
    summary_ids: list[str]
    conversation_id: int = 0
    max_depth: int = 3
    token_cap: Optional[int] = None
    include_messages: bool = False
    owner_id: str = "default"


@dataclass
class ExpansionEntry:
    summary_id: str
    children: list[ChildItem] = field(default_factory=list)
    messages: list[MessageItem] = field(default_factory=list)


@dataclass
class ExpansionResult:
    expansions: list[ExpansionEntry] = field(default_factory=list)
    cited_ids: list[str] = field(default_factory=list)
    total_tokens: int = 0
    truncated: bool = False


@dataclass
class ExpandQueryResult:
    answer: str
    cited_ids: list[str] = field(default_factory=list)
    total_tokens: int = 0
    truncated: bool = False


# ═══════════════════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════════════════


def resolve_expansion_token_cap(
    requested_token_cap: Optional[int],
    max_expand_tokens: int,
) -> int:
    """Clamp requested cap to [1, max_expand_tokens]."""
    if requested_token_cap is None:
        return max(1, max_expand_tokens)
    return min(max(1, requested_token_cap), max(1, max_expand_tokens))


# ═══════════════════════════════════════════════════════════════════════════
# ExpansionOrchestrator
# ═══════════════════════════════════════════════════════════════════════════


class ExpansionOrchestrator:
    """Orchestrates DAG expansion with global token budget.

    expand(): Walk summaryIds through RetrievalEngine with global token cap.
    describe_and_expand(): Grep first, then expand top matches.
    distill_for_subagent(): Format ExpansionResult as compact text payload.
    """

    def __init__(
        self,
        retrieval: RetrievalEngine,
        config: HippoConfig,
    ) -> None:
        self.retrieval = retrieval
        self.config = config

    async def expand(self, request: ExpansionRequest) -> ExpansionResult:
        """Expand each summary ID through the DAG, enforcing global token cap."""
        token_cap = resolve_expansion_token_cap(
            request.token_cap, self.config.max_expand_tokens
        )

        result = ExpansionResult()
        cited_set: set[str] = set()

        for summary_id in request.summary_ids:
            if result.truncated:
                break

            remaining = token_cap - result.total_tokens
            if remaining <= 0:
                result.truncated = True
                break

            raw = await self.retrieval.expand(
                ExpandInput(
                    summary_id=summary_id,
                    depth=request.max_depth,
                    include_messages=request.include_messages,
                    token_cap=remaining,
                )
            )

            entry = ExpansionEntry(
                summary_id=summary_id,
                children=raw.children,
                messages=raw.messages,
            )
            result.expansions.append(entry)
            result.total_tokens += raw.estimated_tokens

            # Collect cited IDs
            cited_set.add(summary_id)
            for child in raw.children:
                cited_set.add(child.summary_id)

            if raw.truncated:
                result.truncated = True

        result.cited_ids = sorted(cited_set)
        return result

    async def describe_and_expand(
        self,
        query: str,
        owner_id: str = "default",
        mode: Literal["regex", "full_text"] = "full_text",
        conversation_id: Optional[int] = None,
        max_depth: int = 3,
        token_cap: Optional[int] = None,
        include_messages: bool = False,
    ) -> ExpansionResult:
        """Grep for summaries matching query, then expand top results."""
        grep_result = await self.retrieval.grep(
            GrepInput(
                pattern=query,
                mode=mode,
                scope="summaries",
                conversation_id=conversation_id,
                owner_id=owner_id,
                limit=20,  # Grab top 20 matches for expansion
            )
        )

        # Sort by created_at desc (recency-first)
        sorted_summaries = sorted(
            grep_result.summaries,
            key=lambda s: s.created_at or "",
            reverse=True,
        )

        summary_ids = [s.summary_id for s in sorted_summaries]
        if not summary_ids:
            return ExpansionResult()

        return await self.expand(
            ExpansionRequest(
                summary_ids=summary_ids,
                conversation_id=conversation_id or 0,
                max_depth=max_depth,
                token_cap=token_cap,
                include_messages=include_messages,
                owner_id=owner_id,
            )
        )

    @staticmethod
    def distill_for_subagent(result: ExpansionResult) -> str:
        """Format ExpansionResult as compact text payload for tool response."""
        lines = [
            f"## Expansion Results ({len(result.expansions)} summaries, "
            f"{result.total_tokens} total tokens)",
            "",
        ]

        for entry in result.expansions:
            kind = "condensed" if entry.children else "leaf"
            token_sum = (
                sum(c.token_count for c in entry.children)
                + sum(m.token_count for m in entry.messages)
            )
            lines.append(f"### {entry.summary_id} ({kind}, {token_sum} tokens)")

            if entry.children:
                lines.append(
                    f"Children: {', '.join(c.summary_id for c in entry.children)}"
                )
            if entry.messages:
                msg_parts = [
                    f"msg#{m.message_id} ({m.role}, {m.token_count} tokens)"
                    for m in entry.messages
                ]
                lines.append(f"Messages: {', '.join(msg_parts)}")

            # Show one snippet
            for child in entry.children:
                if child.content:
                    lines.append(f"[Snippet: {child.content[:200]}]")
                    break

            lines.append("")

        if result.cited_ids:
            lines.append(f"Cited IDs for follow-up: {', '.join(result.cited_ids)}")
        lines.append(f"[Truncated: {'yes' if result.truncated else 'no'}]")

        return "\n".join(lines)


# ═══════════════════════════════════════════════════════════════════════════
# hippo_expand_query execution
# ═══════════════════════════════════════════════════════════════════════════


async def execute_expand_query(
    prompt: str,
    orchestrator: ExpansionOrchestrator,
    complete_fn: LLMCompleteFn,
    config: HippoConfig,
    query: Optional[str] = None,
    summary_ids: Optional[list[str]] = None,
    conversation_id: Optional[int] = None,
    owner_id: str = "default",
    token_cap: Optional[int] = None,
    max_tokens: int = 2000,
) -> ExpandQueryResult:
    """Execute a hippo_expand_query: expand context, then answer with LLM.

    Parameters
    ----------
    prompt : str
        Question to answer using expanded context.
    orchestrator : ExpansionOrchestrator
        For DAG traversal.
    complete_fn : callable
        LLM completion function.
    config : HippoConfig
        Configuration.
    query : str, optional
        Search query for grep-first path. Defaults to prompt.
    summary_ids : list[str], optional
        Direct summary IDs to expand.
    conversation_id : int, optional
        Conversation scope.
    owner_id : str
        Tenant filter.
    token_cap : int, optional
        Token budget for expansion.
    max_tokens : int
        Max tokens for LLM answer.
    """
    effective_cap = resolve_expansion_token_cap(token_cap, config.max_expand_tokens)
    effective_max = min(max_tokens, 4000)

    # Step 1: Expand
    if summary_ids:
        expansion = await orchestrator.expand(
            ExpansionRequest(
                summary_ids=summary_ids,
                conversation_id=conversation_id or 0,
                token_cap=effective_cap,
                owner_id=owner_id,
            )
        )
    else:
        search_query = query or prompt
        expansion = await orchestrator.describe_and_expand(
            query=search_query,
            owner_id=owner_id,
            mode="full_text",
            conversation_id=conversation_id,
            token_cap=effective_cap,
        )

    if not expansion.expansions:
        return ExpandQueryResult(
            answer=f"No summaries found matching: {query or summary_ids or prompt}",
            cited_ids=[],
            total_tokens=0,
            truncated=False,
        )

    # Step 2: Build context for LLM answer
    distilled = ExpansionOrchestrator.distill_for_subagent(expansion)
    answer_prompt = (
        f"Based on the following compacted conversation history, answer this question:\n\n"
        f"Question: {prompt}\n\n"
        f"Context from conversation history:\n{distilled}\n\n"
        f"Provide a focused answer citing specific summary IDs where relevant."
    )

    # Step 3: Call LLM for focused answer
    try:
        result = await complete_fn(
            messages=[{"role": "user", "content": answer_prompt}],
            max_tokens=effective_max,
            temperature=0.1,
        )
        answer_text = ""
        content = result.get("content", [])
        if isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    answer_text += block.get("text", "")
        elif isinstance(content, str):
            answer_text = content
    except Exception as e:
        logger.warning("LLM call failed in expand_query: %s", e)
        answer_text = f"Expansion found {len(expansion.expansions)} summaries but LLM answer generation failed."

    return ExpandQueryResult(
        answer=answer_text,
        cited_ids=expansion.cited_ids,
        total_tokens=expansion.total_tokens,
        truncated=expansion.truncated,
    )
