"""
cortex/hippo/tools/expand_query.py — hippo_expand_query tool implementation.

Full parity with tools/lcm-expand-query-tool.ts (594 lines).
Grep → expand → synthesize pipeline with gateway delegation support.

Issues: #807, #1119, #1146
"""

from __future__ import annotations

import json
import logging
import re
import uuid
from typing import Any, Optional, Protocol, runtime_checkable

from cortex.hippo.tools.common import json_result
from cortex.hippo.tools.conversation_scope import resolve_conversation_scope
from cortex.hippo.tools.expand import _expand_summary_dag, _normalize_summary_ids
from cortex.hippo.tools.recursion_guard import (
    clear_delegated_expansion_context,
    evaluate_expansion_recursion_guard,
    record_expansion_delegation_telemetry,
    resolve_expansion_request_id,
    resolve_next_expansion_depth,
    stamp_delegated_expansion_context,
)

logger = logging.getLogger(__name__)

DEFAULT_MAX_ANSWER_TOKENS = 2_000
DEFAULT_EXPANSION_TOKEN_CAP = 50_000

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

HIPPO_EXPAND_QUERY_TOOL = {
    "name": "hippo_expand_query",
    "description": (
        "Answer a focused question using delegated Hippo expansion. "
        "Find candidate summaries (by IDs or query), expand them, and return a "
        "compact prompt-focused answer with cited summary IDs. "
        "Higher-level than hippo_expand — use this when you have a question to answer, "
        "not just IDs to traverse."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "prompt": {
                "type": "string",
                "description": "Question to answer using expanded context.",
            },
            "query": {
                "type": "string",
                "description": (
                    "Text query used to find summaries via grep before expansion. "
                    "Required when summaryIds is not provided."
                ),
            },
            "summaryIds": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Summary IDs to expand (sum_xxx). Required when query is not provided.",
            },
            "maxTokens": {
                "type": "integer",
                "default": DEFAULT_MAX_ANSWER_TOKENS,
                "minimum": 1,
                "description": "Maximum answer tokens to target.",
            },
            "tokenCap": {
                "type": "integer",
                "minimum": 1,
                "description": "Expansion retrieval token budget.",
            },
            "conversationId": {
                "type": "integer",
                "description": "Conversation ID to scope expansion to.",
            },
            "allConversations": {
                "type": "boolean",
                "default": False,
                "description": "Allow cross-conversation lookup.",
            },
        },
        "required": ["prompt"],
    },
}


# ═══════════════════════════════════════════════════════════════════════════
# Protocols / callback types
# ═══════════════════════════════════════════════════════════════════════════


@runtime_checkable
class SynthesizeFn(Protocol):
    """Callback that synthesizes an answer from expanded context.

    Matches the summarize_fn pattern used in compaction —
    can be a direct LLM call or any async callable.
    """

    async def __call__(self, context: str, prompt: str, max_tokens: int) -> str: ...


@runtime_checkable
class GatewayDelegator(Protocol):
    """Protocol for gateway-based delegation.

    When available, expand_query uses this to spawn a sub-agent session
    for autonomous retrieval navigation instead of local synthesis.
    """

    async def delegate(
        self,
        system_prompt: str,
        user_message: str,
        timeout: float,
    ) -> str: ...


# ═══════════════════════════════════════════════════════════════════════════
# Expand-query reply parsing
# ═══════════════════════════════════════════════════════════════════════════


def _parse_delegated_reply(
    raw_reply: str | None,
    fallback_expanded_count: int,
) -> dict[str, Any]:
    """Parse delegated sub-agent reply.

    Accepts plain JSON or fenced JSON blocks (```json ... ```).
    Falls back to treating the raw text as the answer.
    """
    fallback = {
        "answer": (raw_reply or "").strip(),
        "citedIds": [],
        "expandedSummaryCount": fallback_expanded_count,
        "totalSourceTokens": 0,
        "truncated": False,
    }

    reply = (raw_reply or "").strip()
    if not reply:
        return fallback

    candidates: list[str] = [reply]
    fenced = re.search(r"```(?:json)?\s*([\s\S]*?)```", reply, re.IGNORECASE)
    if fenced and fenced.group(1):
        candidates.insert(0, fenced.group(1).strip())

    for candidate in candidates:
        try:
            parsed = json.loads(candidate)
            answer = parsed.get("answer", "")
            if isinstance(answer, str):
                answer = answer.strip()
            else:
                answer = ""

            cited = parsed.get("citedIds", [])
            if isinstance(cited, list):
                cited = _normalize_summary_ids([c for c in cited if isinstance(c, str)])
            else:
                cited = []

            expanded_count = parsed.get("expandedSummaryCount", fallback_expanded_count)
            if not (isinstance(expanded_count, (int, float)) and expanded_count == expanded_count):
                expanded_count = fallback_expanded_count
            expanded_count = max(0, int(expanded_count))

            total_tokens = parsed.get("totalSourceTokens", 0)
            if not (isinstance(total_tokens, (int, float)) and total_tokens == total_tokens):
                total_tokens = 0
            total_tokens = max(0, int(total_tokens))

            truncated = parsed.get("truncated") is True

            return {
                "answer": answer or fallback["answer"],
                "citedIds": cited,
                "expandedSummaryCount": expanded_count,
                "totalSourceTokens": total_tokens,
                "truncated": truncated,
            }
        except (json.JSONDecodeError, AttributeError, TypeError):
            continue

    return fallback


# ═══════════════════════════════════════════════════════════════════════════
# Delegation task builder (matches TS buildDelegatedExpandQueryTask)
# ═══════════════════════════════════════════════════════════════════════════


def _build_delegation_task(
    *,
    summary_ids: list[str],
    conversation_id: int,
    query: str | None,
    prompt: str,
    max_tokens: int,
    token_cap: int,
    request_id: str,
    expansion_depth: int,
    origin_session_key: str,
) -> str:
    """Build the system+user message for a delegated expansion sub-agent."""
    seed_ids = ", ".join(summary_ids) if summary_ids else "(none)"
    lines = [
        "You are an autonomous LCM retrieval navigator. Plan and execute retrieval before answering.",
        "",
        "Available tools: lcm_describe, lcm_expand, lcm_grep",
        f"Conversation scope: {conversation_id}",
        f"Expansion token budget (total across this run): {token_cap}",
        f"Seed summary IDs: {seed_ids}",
    ]
    if query:
        lines.append(f"Routing query: {query}")
    lines.extend([
        "",
        "Strategy:",
        "1. Start with `lcm_describe` on seed summaries to inspect subtree manifests and branch costs.",
        "2. If additional candidates are needed, use `lcm_grep` scoped to summaries.",
        "3. Select branches that fit remaining budget; prefer high-signal paths first.",
        "4. Call `lcm_expand` selectively (do not expand everything blindly).",
        "5. Keep includeMessages=false by default; use includeMessages=true only for specific leaf evidence.",
        f"6. Stay within {token_cap} total expansion tokens across all lcm_expand calls.",
        "",
        "User prompt to answer:",
        prompt,
        "",
        "Delegated expansion metadata (for tracing):",
        f"- requestId: {request_id}",
        f"- expansionDepth: {expansion_depth}",
        f"- originSessionKey: {origin_session_key}",
        "",
        "Return ONLY JSON with this shape:",
        "{",
        '  "answer": "string",',
        '  "citedIds": ["sum_xxx"],',
        '  "expandedSummaryCount": 0,',
        '  "totalSourceTokens": 0,',
        '  "truncated": false',
        "}",
        "",
        "Rules:",
        "- In delegated context, call `lcm_expand` directly for source retrieval.",
        "- DO NOT call `lcm_expand_query` from this delegated session.",
        "- Synthesize the final answer from retrieved evidence, not assumptions.",
        f"- Keep answer concise and focused (target <= {max_tokens} tokens).",
        "- citedIds must be unique summary IDs.",
        "- expandedSummaryCount should reflect how many summaries were expanded/used.",
        "- totalSourceTokens should estimate total tokens consumed from expansion calls.",
        "- truncated should indicate whether source expansion appears truncated.",
    ])
    return "\n".join(lines)


# ═══════════════════════════════════════════════════════════════════════════
# Summary candidate resolution
# ═══════════════════════════════════════════════════════════════════════════


async def _resolve_summary_candidates(
    retrieval,
    *,
    explicit_ids: list[str],
    query: str | None,
    conversation_id: int | None,
    owner_id: str = "default",
) -> list[dict[str, Any]]:
    """Resolve candidate summaries from explicit IDs and/or grep query.

    Returns list of {summaryId, conversationId} dicts.
    """
    candidates: dict[str, dict[str, Any]] = {}

    # Resolve explicit IDs via describe
    for sid in explicit_ids:
        desc = await retrieval.describe(id=sid, owner_id=owner_id)
        if not desc or desc.type != "summary" or not desc.summary:
            raise ValueError(f"Summary not found: {sid}")
        conv_id = desc.summary.get("conversation_id")
        candidates[sid] = {"summaryId": sid, "conversationId": conv_id}

    # Grep for additional candidates
    if query:
        from cortex.hippo.retrieval import GrepInput

        grep_input = GrepInput(
            pattern=query,
            mode="full_text",
            scope="summaries",
            conversation_id=conversation_id,
            owner_id=owner_id,
        )
        grep_result = await retrieval.grep(grep_input)
        for s in grep_result.summaries:
            if s.summary_id not in candidates:
                candidates[s.summary_id] = {
                    "summaryId": s.summary_id,
                    "conversationId": s.conversation_id,
                }

    return list(candidates.values())


def _resolve_source_conversation_id(
    *,
    scoped_conversation_id: int | None,
    all_conversations: bool,
    candidates: list[dict[str, Any]],
) -> int:
    """Resolve a single source conversation ID for delegated expansion.

    Raises ValueError when unable to resolve unambiguously.
    """
    if scoped_conversation_id is not None:
        mismatched = [
            c["summaryId"]
            for c in candidates
            if c["conversationId"] != scoped_conversation_id
        ]
        if mismatched:
            raise ValueError(
                f"Some summaryIds are outside conversation "
                f"{scoped_conversation_id}: {', '.join(mismatched)}"
            )
        return scoped_conversation_id

    conv_ids = list({c["conversationId"] for c in candidates})
    if len(conv_ids) == 1 and isinstance(conv_ids[0], int):
        return conv_ids[0]

    if all_conversations and len(conv_ids) > 1:
        raise ValueError(
            "Query matched summaries from multiple conversations. "
            "Provide conversationId or narrow the query."
        )

    raise ValueError(
        "Unable to resolve a single conversation scope. "
        "Provide conversationId or set a narrower summary scope."
    )


# ═══════════════════════════════════════════════════════════════════════════
# Local synthesis (default when no gateway delegator)
# ═══════════════════════════════════════════════════════════════════════════


async def _local_synthesize(
    *,
    retrieval,
    summary_ids: list[str],
    prompt: str,
    max_tokens: int,
    token_cap: int | None,
    owner_id: str,
    synthesize_fn: SynthesizeFn | None,
) -> dict[str, Any]:
    """Expand summaries locally and synthesize an answer.

    This is the fallback path when no gateway delegator is available.
    Uses the expand tool's DAG traversal, then calls synthesize_fn.
    """
    from cortex.hippo.tools.expand import _distill_for_output

    result = await _expand_summary_dag(
        retrieval,
        summary_ids,
        max_depth=3,
        token_cap=token_cap,
        include_messages=False,
        owner_id=owner_id,
    )

    context_text = _distill_for_output(result)

    if synthesize_fn is not None:
        answer = await synthesize_fn(context_text, prompt, max_tokens)
    else:
        # No synthesize_fn — return raw context as answer
        answer = context_text

    return {
        "answer": answer,
        "citedIds": result.get("citedIds", []),
        "expandedSummaryCount": len(result.get("expansions", [])),
        "totalSourceTokens": result.get("totalTokens", 0),
        "truncated": result.get("truncated", False),
    }


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


async def execute_expand_query(
    params: dict[str, Any],
    retrieval,
    *,
    session_id: str | None = None,
    session_key: str | None = None,
    owner_id: str = "default",
    conversation_store=None,
    synthesize_fn: SynthesizeFn | None = None,
    gateway_delegator: GatewayDelegator | None = None,
    max_expand_tokens: int = DEFAULT_EXPANSION_TOKEN_CAP,
) -> dict[str, Any]:
    """Execute the hippo_expand_query tool.

    Core flow:
    1. Validate params (prompt required, query or summaryIds required)
    2. Check recursion guard
    3. Resolve conversation scope + summary candidates
    4. If gateway_delegator available: delegate to sub-agent for autonomous retrieval
    5. Otherwise: local expand → synthesize via synthesize_fn

    Args:
        params: Tool parameters.
        retrieval: RetrievalEngine instance.
        session_id: Current session ID for conversation scope resolution.
        session_key: Session key for recursion guard scoping.
        owner_id: Owner ID for multi-tenant isolation.
        conversation_store: Optional conversation store for scope resolution.
        synthesize_fn: Callback for local synthesis (context, prompt, max_tokens) -> answer.
        gateway_delegator: Optional gateway delegator for sub-agent delegation.
        max_expand_tokens: Default expansion token budget from config.

    Returns:
        Tool result dict with answer, citedIds, and expansion stats.
    """
    # ── Parse params ───────────────────────────────────────────────
    prompt = params.get("prompt", "")
    if isinstance(prompt, str):
        prompt = prompt.strip()
    else:
        prompt = ""

    if not prompt:
        return json_result({"error": "prompt is required."})

    raw_summary_ids = params.get("summaryIds")
    explicit_ids = _normalize_summary_ids(raw_summary_ids) if isinstance(raw_summary_ids, list) else []

    query = params.get("query", "")
    if isinstance(query, str):
        query = query.strip()
    else:
        query = ""

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

    raw_max_tokens = params.get("maxTokens")
    if isinstance(raw_max_tokens, (int, float)) and raw_max_tokens == raw_max_tokens:
        max_tokens = max(1, int(raw_max_tokens))
    else:
        max_tokens = DEFAULT_MAX_ANSWER_TOKENS

    raw_token_cap = params.get("tokenCap")
    if isinstance(raw_token_cap, (int, float)) and raw_token_cap == raw_token_cap:
        expansion_token_cap = max(1, int(raw_token_cap))
    else:
        expansion_token_cap = max(1, int(max_expand_tokens))

    # ── Recursion guard ────────────────────────────────────────────
    caller_session_key = (session_key or session_id or "").strip()
    request_id = resolve_expansion_request_id(caller_session_key or None)

    recursion_check = evaluate_expansion_recursion_guard(
        session_key=caller_session_key or None,
        request_id=request_id,
    )

    record_expansion_delegation_telemetry(
        component="lcm_expand_query",
        event="start",
        request_id=request_id,
        session_key=caller_session_key or None,
        expansion_depth=recursion_check.expansion_depth,
        origin_session_key=recursion_check.origin_session_key,
    )

    if recursion_check.blocked:
        record_expansion_delegation_telemetry(
            component="lcm_expand_query",
            event="block",
            request_id=request_id,
            session_key=caller_session_key or None,
            expansion_depth=recursion_check.expansion_depth,
            origin_session_key=recursion_check.origin_session_key,
            reason=recursion_check.reason,
        )
        return json_result({
            "errorCode": recursion_check.code,
            "error": recursion_check.message,
            "requestId": recursion_check.request_id,
            "expansionDepth": recursion_check.expansion_depth,
            "originSessionKey": recursion_check.origin_session_key,
            "reason": recursion_check.reason,
        })

    # ── Conversation scope ─────────────────────────────────────────
    scope = await resolve_conversation_scope(
        params=params,
        session_id=session_id,
        owner_id=owner_id,
        conversation_store=conversation_store,
    )

    scoped_conv_id = scope.conversation_id

    if not scope.all_conversations and scoped_conv_id is None:
        return json_result({
            "error": (
                "No LCM conversation found for this session. "
                "Provide conversationId or set allConversations=true."
            ),
        })

    # ── Resolve candidates ─────────────────────────────────────────
    child_session_key = ""

    try:
        candidates = await _resolve_summary_candidates(
            retrieval,
            explicit_ids=explicit_ids,
            query=query or None,
            conversation_id=scoped_conv_id,
            owner_id=owner_id,
        )

        if not candidates:
            if scoped_conv_id is None:
                return json_result({"error": "No matching summaries found."})
            return json_result({
                "answer": "No matching summaries found for this scope.",
                "citedIds": [],
                "sourceConversationId": scoped_conv_id,
                "expandedSummaryCount": 0,
                "totalSourceTokens": 0,
                "truncated": False,
            })

        source_conv_id = _resolve_source_conversation_id(
            scoped_conversation_id=scoped_conv_id,
            all_conversations=scope.all_conversations,
            candidates=candidates,
        )

        summary_ids = _normalize_summary_ids([
            c["summaryId"]
            for c in candidates
            if c["conversationId"] == source_conv_id
        ])

        if not summary_ids:
            return json_result({
                "error": "No summaryIds available after applying conversation scope.",
            })

        # ── Delegation vs local synthesis ──────────────────────────
        if gateway_delegator is not None:
            # Gateway delegation path
            child_session_key = f"agent:hippo:subagent:{uuid.uuid4()}"
            child_depth = resolve_next_expansion_depth(caller_session_key or None)
            origin_key = recursion_check.origin_session_key or caller_session_key or "main"

            stamp_delegated_expansion_context(
                session_key=child_session_key,
                request_id=request_id,
                expansion_depth=child_depth,
                origin_session_key=origin_key,
                stamped_by="lcm_expand_query",
            )

            task = _build_delegation_task(
                summary_ids=summary_ids,
                conversation_id=source_conv_id,
                query=query or None,
                prompt=prompt,
                max_tokens=max_tokens,
                token_cap=expansion_token_cap,
                request_id=request_id,
                expansion_depth=child_depth,
                origin_session_key=origin_key,
            )

            try:
                raw_reply = await gateway_delegator.delegate(
                    system_prompt=task,
                    user_message=prompt,
                    timeout=120.0,
                )
                parsed = _parse_delegated_reply(raw_reply, len(summary_ids))

                record_expansion_delegation_telemetry(
                    component="lcm_expand_query",
                    event="success",
                    request_id=request_id,
                    session_key=caller_session_key or None,
                    expansion_depth=child_depth,
                    origin_session_key=origin_key,
                )

                return json_result({
                    "answer": parsed["answer"],
                    "citedIds": parsed["citedIds"],
                    "sourceConversationId": source_conv_id,
                    "expandedSummaryCount": parsed["expandedSummaryCount"],
                    "totalSourceTokens": parsed["totalSourceTokens"],
                    "truncated": parsed["truncated"],
                })
            except Exception as err:
                logger.warning(
                    "Gateway delegation failed, falling back to local synthesis: %s", err,
                )
                # Fall through to local synthesis

        # Local synthesis path
        synth_result = await _local_synthesize(
            retrieval=retrieval,
            summary_ids=summary_ids,
            prompt=prompt,
            max_tokens=max_tokens,
            token_cap=expansion_token_cap,
            owner_id=owner_id,
            synthesize_fn=synthesize_fn,
        )

        record_expansion_delegation_telemetry(
            component="lcm_expand_query",
            event="success",
            request_id=request_id,
            session_key=caller_session_key or None,
            expansion_depth=recursion_check.expansion_depth,
            origin_session_key=recursion_check.origin_session_key,
        )

        return json_result({
            "answer": synth_result["answer"],
            "citedIds": synth_result["citedIds"],
            "sourceConversationId": source_conv_id,
            "expandedSummaryCount": synth_result["expandedSummaryCount"],
            "totalSourceTokens": synth_result["totalSourceTokens"],
            "truncated": synth_result["truncated"],
        })

    except Exception as err:
        return json_result({"error": str(err)})

    finally:
        if child_session_key:
            clear_delegated_expansion_context(child_session_key)
