"""Session search tool for recalling past conversations.

Searches JSONL session history, ranks matching sessions, then returns concise
session summaries. If an agent runtime/provider is available, summaries are
generated by the model; otherwise deterministic fallback summaries are used.
"""

from __future__ import annotations

import asyncio
import json
import re
from pathlib import Path
from typing import Any

from kyber.agent.tools.base import Tool
from kyber.agent.tools.registry import registry


SESSIONS_DIR = Path.home() / ".kyber" / "sessions"
MAX_RESULTS = 5
MAX_CONTEXT_MESSAGES = 14
MAX_CONTEXT_CHARS = 18000
MAX_SUMMARY_CHARS = 2200


def _read_session(path: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    """Read one JSONL session file, returning metadata and message rows."""
    meta: dict[str, Any] = {}
    messages: list[dict[str, Any]] = []

    try:
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                raw = line.strip()
                if not raw:
                    continue
                try:
                    obj = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                if not isinstance(obj, dict):
                    continue
                if obj.get("_type") == "metadata":
                    meta = obj
                    continue
                if "content" not in obj:
                    continue
                messages.append(obj)
    except Exception:
        return {}, []

    return meta, messages


def _query_terms(query: str) -> list[str]:
    """Normalize query into lowercase terms."""
    words = re.findall(r"[a-z0-9][a-z0-9_-]{1,}", query.lower())
    return words[:12]


def _score_message(content: str, query_lower: str, terms: list[str]) -> tuple[int, int]:
    """Return (score, term_hits) for one message."""
    text = (content or "").lower()
    if not text:
        return 0, 0

    term_hits = 0
    score = 0

    if query_lower and query_lower in text:
        score += 8

    for term in terms:
        if term in text:
            hits = text.count(term)
            term_hits += hits
            score += min(4, hits + 1)

    return score, term_hits


def _shorten(text: str, limit: int = 360) -> str:
    """Shorten long message text while preserving useful tail context."""
    clean = " ".join((text or "").split())
    if len(clean) <= limit:
        return clean
    head = clean[: max(0, limit - 90)]
    tail = clean[-70:]
    return f"{head} ... {tail}"


def _format_message(msg: dict[str, Any]) -> str:
    role = str(msg.get("role", "unknown")).upper()
    content = _shorten(str(msg.get("content", "")))
    tool_name = msg.get("tool_name")
    if role == "TOOL" and tool_name:
        return f"[TOOL:{tool_name}] {content}"
    return f"[{role}] {content}"


def _build_context(messages: list[dict[str, Any]], match_indexes: list[int]) -> str:
    """Build a compact transcript centered around matched messages."""
    if not messages:
        return ""

    selected: list[int] = []
    seen: set[int] = set()
    radius = 2

    for idx in match_indexes[:5]:
        for pos in range(max(0, idx - radius), min(len(messages), idx + radius + 1)):
            if pos in seen:
                continue
            seen.add(pos)
            selected.append(pos)

    if not selected:
        selected = list(range(max(0, len(messages) - MAX_CONTEXT_MESSAGES), len(messages)))

    selected = sorted(selected)[:MAX_CONTEXT_MESSAGES]
    lines = [_format_message(messages[idx]) for idx in selected]
    text = "\n\n".join(lines)

    if len(text) > MAX_CONTEXT_CHARS:
        keep = MAX_CONTEXT_CHARS // 2
        text = text[:keep] + "\n\n...[truncated]...\n\n" + text[-keep:]

    return text


def _build_fallback_summary(context: str, match_count: int) -> str:
    """Fallback summary when LLM summarization is unavailable."""
    preview_lines = [line for line in context.splitlines() if line.strip()][:6]
    preview = "\n".join(preview_lines)
    summary = (
        f"- Matched {match_count} relevant message(s) in this session.\n"
        f"- Context preview:\n{preview}"
    ).strip()
    if len(summary) > MAX_SUMMARY_CHARS:
        summary = summary[: MAX_SUMMARY_CHARS - 3] + "..."
    return summary


async def _summarize_with_provider(
    *,
    agent_core: Any,
    query: str,
    session_id: str,
    updated_at: str,
    context: str,
    match_count: int,
) -> str | None:
    """Summarize one session using the active agent provider."""
    provider = getattr(agent_core, "provider", None)
    if provider is None:
        return None

    model = getattr(agent_core, "model", None)
    system_prompt = (
        "You summarize prior chat sessions for memory recall. Produce a compact, factual recap "
        "focused on what helps continue the current conversation: user goal, actions taken, "
        "key outputs, commands/files, and final outcome."
    )
    user_prompt = (
        f"Query: {query}\n"
        f"Session ID: {session_id}\n"
        f"Last Updated: {updated_at}\n"
        f"Matched messages: {match_count}\n\n"
        "Summarize this session context in 4-8 bullet points.\n\n"
        f"SESSION CONTEXT:\n{context}"
    )

    try:
        response = await asyncio.wait_for(
            provider.chat(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                tools=None,
                model=model,
                max_tokens=600,
                temperature=0.2,
            ),
            timeout=45.0,
        )
    except Exception:
        return None

    content = (response.content or "").strip()
    if not content:
        return None
    if len(content) > MAX_SUMMARY_CHARS:
        content = content[: MAX_SUMMARY_CHARS - 3] + "..."
    return content


def _search_sessions(query: str, limit: int) -> list[dict[str, Any]]:
    """Rank session files by relevance to query."""
    if not SESSIONS_DIR.exists():
        return []

    query_norm = " ".join(query.lower().split())
    terms = _query_terms(query)
    if not terms and not query_norm:
        return []

    hits: list[dict[str, Any]] = []

    for session_file in SESSIONS_DIR.glob("*.jsonl"):
        meta, messages = _read_session(session_file)
        if not messages:
            continue

        total_score = 0
        total_term_hits = 0
        match_indexes: list[int] = []

        for idx, msg in enumerate(messages):
            score, term_hits = _score_message(str(msg.get("content", "")), query_norm, terms)
            if score <= 0:
                continue
            total_score += score
            total_term_hits += term_hits
            match_indexes.append(idx)

        if total_score <= 0:
            continue

        context = _build_context(messages, match_indexes)
        hits.append(
            {
                "session_id": session_file.stem,
                "updated_at": str(meta.get("updated_at") or meta.get("created_at") or ""),
                "score": total_score,
                "match_count": len(match_indexes),
                "term_hits": total_term_hits,
                "context": context,
            }
        )

    hits.sort(
        key=lambda x: (x["score"], x["match_count"], x["term_hits"], x["updated_at"]),
        reverse=True,
    )
    return hits[:limit]


class SessionSearchTool(Tool):
    """Search past conversations and summarize relevant sessions."""

    @property
    def name(self) -> str:
        return "session_search"

    @property
    def description(self) -> str:
        return (
            "Search long-term conversation history and return concise session summaries.\n\n"
            "Use when the user references prior work: 'we did this before', 'last time', "
            "'remember when', or asks what was previously decided."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Keywords or phrase to find in prior sessions.",
                },
                "limit": {
                    "type": "integer",
                    "description": "Max sessions to return (default 3, max 5).",
                    "default": 3,
                    "minimum": 1,
                    "maximum": 5,
                },
            },
            "required": ["query"],
        }

    @property
    def toolset(self) -> str:
        return "memory"

    async def execute(self, query: str, limit: int = 3, **kwargs: Any) -> str:
        query_clean = (query or "").strip()
        if not query_clean:
            return json.dumps({"success": False, "error": "Query cannot be empty."}, ensure_ascii=False)

        clamped_limit = max(1, min(limit, MAX_RESULTS))
        ranked = _search_sessions(query_clean, clamped_limit)
        if not ranked:
            return json.dumps(
                {
                    "success": True,
                    "query": query_clean,
                    "count": 0,
                    "results": [],
                    "message": "No matching sessions found.",
                },
                ensure_ascii=False,
            )

        agent_core = kwargs.get("agent_core")
        summaries: list[str | None] = [None] * len(ranked)

        if agent_core is not None:
            jobs = [
                _summarize_with_provider(
                    agent_core=agent_core,
                    query=query_clean,
                    session_id=entry["session_id"],
                    updated_at=entry["updated_at"],
                    context=entry["context"],
                    match_count=entry["match_count"],
                )
                for entry in ranked
            ]
            out = await asyncio.gather(*jobs, return_exceptions=True)
            for i, item in enumerate(out):
                if isinstance(item, Exception):
                    summaries[i] = None
                else:
                    summaries[i] = item

        formatted = []
        for idx, entry in enumerate(ranked):
            summary = summaries[idx] or _build_fallback_summary(
                entry["context"],
                entry["match_count"],
            )
            formatted.append(
                {
                    "session_id": entry["session_id"],
                    "updated_at": entry["updated_at"] or None,
                    "score": entry["score"],
                    "match_count": entry["match_count"],
                    "summary": summary,
                }
            )

        return json.dumps(
            {
                "success": True,
                "query": query_clean,
                "count": len(formatted),
                "results": formatted,
            },
            ensure_ascii=False,
        )


registry.register(SessionSearchTool())
