"""
cortex/hippo/tools/expand.py — hippo_expand tool implementation.

Full parity with tools/lcm-expand-tool.ts (448 lines).
Traverses the summary DAG to retrieve children and source messages.

Issue: #807, #1118
"""

from __future__ import annotations

from typing import Any, Optional

from cortex.hippo.tools.common import json_result
from cortex.hippo.tools.conversation_scope import resolve_conversation_scope


# ═══════════════════════════════════════════════════════════════════════════
# Schema
# ═══════════════════════════════════════════════════════════════════════════

HIPPO_EXPAND_TOOL = {
    "name": "hippo_expand",
    "description": (
        "Expand compacted conversation summaries from Hippo LCM. "
        "Traverses the summary DAG to retrieve children and source messages. "
        "Use this to drill into previously-compacted context when you need detail "
        "that was summarised away. "
        "Provide either summaryIds (direct expansion) or query (grep-first, then expand top matches). "
        "Returns a compact text payload with cited IDs for follow-up."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "summaryIds": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Summary IDs to expand (sum_xxx format). Required if query not provided.",
            },
            "query": {
                "type": "string",
                "description": (
                    "Text query to grep for matching summaries before expanding. "
                    "If provided, summaryIds is ignored."
                ),
            },
            "maxDepth": {
                "type": "integer",
                "default": 3,
                "minimum": 1,
                "maximum": 10,
                "description": "Max traversal depth per summary.",
            },
            "tokenCap": {
                "type": "integer",
                "minimum": 1,
                "description": "Max tokens across the entire expansion result.",
            },
            "includeMessages": {
                "type": "boolean",
                "default": False,
                "description": "Include raw source messages at leaf level.",
            },
            "conversationId": {
                "type": "integer",
                "description": "Conversation ID to scope expansion to.",
            },
            "allConversations": {
                "type": "boolean",
                "default": False,
                "description": "Allow cross-conversation expansion.",
            },
        },
    },
}


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


def _normalize_summary_ids(ids: list[str]) -> list[str]:
    """Deduplicate and filter valid summary IDs."""
    seen: set[str] = set()
    result: list[str] = []
    for raw in ids:
        sid = raw.strip()
        if sid and sid.startswith("sum_") and sid not in seen:
            seen.add(sid)
            result.append(sid)
    return result


def _distill_for_output(result: dict[str, Any]) -> str:
    """Render expansion result as compact markdown text with cited IDs."""
    lines: list[str] = []
    cited: list[str] = []

    for expansion in result.get("expansions", []):
        sid = expansion.get("summary_id", "")
        cited.append(sid)
        lines.append(f"## {sid}")

        content = expansion.get("content", "")
        if content:
            lines.append(content)

        children = expansion.get("children", [])
        for child in children:
            child_id = child.get("summary_id", "")
            cited.append(child_id)
            depth_marker = "  " * child.get("depth_from_root", 1)
            lines.append(f"{depth_marker}### {child_id} (k={child.get('kind', '')})")
            child_content = child.get("content", "")
            if child_content:
                lines.append(f"{depth_marker}{child_content}")

        messages = expansion.get("messages", [])
        for msg in messages:
            lines.append(f"  [{msg.get('role', '?')}] {msg.get('content', '')}")

    if result.get("truncated"):
        lines.append("\n---\n*[truncated: token budget reached]*")

    if cited:
        lines.append(f"\ncited: {' '.join(cited)}")

    return "\n".join(lines) if lines else "(no expansions)"


# ═══════════════════════════════════════════════════════════════════════════
# DAG Expansion
# ═══════════════════════════════════════════════════════════════════════════


async def _expand_summary_dag(
    retrieval,
    summary_ids: list[str],
    *,
    max_depth: int = 3,
    token_cap: Optional[int] = None,
    include_messages: bool = False,
    owner_id: str = "default",
) -> dict[str, Any]:
    """Expand multiple summaries through the DAG.

    Returns a dict matching the ExpansionResult shape:
    {expansions, citedIds, totalTokens, truncated}
    """
    from cortex.hippo.retrieval import ExpandInput

    budget = token_cap if token_cap is not None else float("inf")
    total_tokens = 0
    truncated = False
    expansions: list[dict[str, Any]] = []
    cited_ids: list[str] = []

    for sid in summary_ids:
        if truncated:
            break

        remaining = int(budget - total_tokens) if budget != float("inf") else None
        if remaining is not None and remaining <= 0:
            truncated = True
            break

        expand_input = ExpandInput(
            summary_id=sid,
            depth=max_depth,
            include_messages=include_messages,
            token_cap=remaining,
        )
        result = await retrieval.expand(expand_input)

        expansion_entry: dict[str, Any] = {
            "summary_id": sid,
            "children": [
                {
                    "summary_id": c.summary_id,
                    "kind": c.kind,
                    "content": c.content,
                    "token_count": c.token_count,
                }
                for c in result.children
            ],
            "messages": [
                {
                    "message_id": m.message_id,
                    "role": m.role,
                    "content": m.content,
                    "token_count": m.token_count,
                }
                for m in result.messages
            ],
        }

        # Add the root summary content
        desc = await retrieval.describe(id=sid, owner_id=owner_id)
        if desc and desc.summary:
            expansion_entry["content"] = desc.summary.get("content", "")
            cited_ids.append(sid)

        cited_ids.extend(c.summary_id for c in result.children)
        total_tokens += result.estimated_tokens
        truncated = truncated or result.truncated
        expansions.append(expansion_entry)

    return {
        "expansions": expansions,
        "citedIds": cited_ids,
        "totalTokens": total_tokens,
        "truncated": truncated,
    }


# ═══════════════════════════════════════════════════════════════════════════
# Execute
# ═══════════════════════════════════════════════════════════════════════════


async def execute_expand(
    params: dict[str, Any],
    retrieval,
    *,
    session_id: Optional[str] = None,
    owner_id: str = "default",
    conversation_store=None,
) -> dict[str, Any]:
    """Execute the hippo_expand tool.

    Args:
        params: Tool parameters.
        retrieval: RetrievalEngine instance.
        session_id: Current session ID for conversation scope resolution.
        owner_id: Owner ID for multi-tenant isolation.
        conversation_store: Optional conversation store for scope resolution.

    Returns:
        Tool result dict with content and details.
    """
    summary_ids = params.get("summaryIds")
    query = params.get("query")
    if isinstance(query, str):
        query = query.strip() or None
    max_depth = params.get("maxDepth", 3)
    if isinstance(max_depth, (int, float)):
        max_depth = max(1, int(max_depth))
    else:
        max_depth = 3
    token_cap = params.get("tokenCap")
    if isinstance(token_cap, (int, float)):
        token_cap = max(1, int(token_cap))
    else:
        token_cap = None
    include_messages = bool(params.get("includeMessages", False))

    # Resolve conversation scope
    scope = await resolve_conversation_scope(
        params=params,
        session_id=session_id,
        owner_id=owner_id,
        conversation_store=conversation_store,
    )

    # ── Query path: grep first, then expand top matches ────────────
    if query:
        try:
            from cortex.hippo.retrieval import GrepInput

            grep_input = GrepInput(
                pattern=query,
                mode="full_text",
                scope="summaries",
                conversation_id=scope.conversation_id,
                all_conversations=scope.all_conversations,
                owner_id=owner_id,
            )
            grep_result = await retrieval.grep(grep_input)
            matched_ids = [s.summary_id for s in grep_result.summaries]

            if not matched_ids:
                result = {
                    "expansions": [],
                    "citedIds": [],
                    "totalTokens": 0,
                    "truncated": False,
                }
            else:
                result = await _expand_summary_dag(
                    retrieval,
                    matched_ids,
                    max_depth=max_depth,
                    token_cap=token_cap,
                    include_messages=include_messages,
                    owner_id=owner_id,
                )

            text = _distill_for_output(result)
            return {
                "content": [{"type": "text", "text": text}],
                "details": {
                    "expansionCount": len(result["expansions"]),
                    "citedIds": result["citedIds"],
                    "totalTokens": result["totalTokens"],
                    "truncated": result["truncated"],
                    "executionPath": "direct",
                    "queryMatchCount": len(matched_ids),
                },
            }
        except Exception as err:
            return json_result({"error": str(err)})

    # ── SummaryIds path: direct expansion ──────────────────────────
    if summary_ids and len(summary_ids) > 0:
        try:
            normalized = _normalize_summary_ids(summary_ids)
            if not normalized:
                return json_result({"error": "No valid summary IDs provided."})

            # Conversation scope check
            if scope.conversation_id is not None:
                out_of_scope: list[str] = []
                for sid in normalized:
                    described = await retrieval.describe(id=sid, owner_id=owner_id)
                    if (
                        described
                        and described.type == "summary"
                        and described.summary
                        and described.summary.get("conversation_id") != scope.conversation_id
                    ):
                        out_of_scope.append(sid)
                if out_of_scope:
                    return json_result({
                        "error": (
                            f"Some summaryIds are outside conversation "
                            f"{scope.conversation_id}: {', '.join(out_of_scope)}"
                        ),
                        "hint": "Use allConversations=true for cross-conversation expansion.",
                    })

            result = await _expand_summary_dag(
                retrieval,
                normalized,
                max_depth=max_depth,
                token_cap=token_cap,
                include_messages=include_messages,
                owner_id=owner_id,
            )

            text = _distill_for_output(result)
            return {
                "content": [{"type": "text", "text": text}],
                "details": {
                    "expansionCount": len(result["expansions"]),
                    "citedIds": result["citedIds"],
                    "totalTokens": result["totalTokens"],
                    "truncated": result["truncated"],
                    "executionPath": "direct",
                },
            }
        except Exception as err:
            return json_result({"error": str(err)})

    return json_result({"error": "Either summaryIds or query must be provided."})
