"""Shared multi-word search logic used by both claude-history and codex-history."""

import json as _rg_json
import re
import shutil
import subprocess
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Callable

try:
    import orjson as _json

    def _loads(line: str) -> Any:
        return _json.loads(line)

except Exception:
    import json as _json

    def _loads(line: str) -> Any:
        return _json.loads(line)


@dataclass
class SearchMatch:
    """Search match within a conversation."""

    session_id: str
    title: str
    timestamp: datetime | None
    snippet: str


# Callback that extracts (role, text, timestamp) from a parsed JSONL entry,
# or returns None if the entry should be skipped.
ExtractFn = Callable[[dict[str, Any]], tuple[str, str, datetime | None] | None]

# Callback that returns (session_id, title) for a given file path.
TitleFn = Callable[[Path], tuple[str, str]]


def _tokenize_query(query: str) -> list[str]:
    """Split a query into individual words for multi-word matching."""
    words = re.split(r"\s+", query.strip())
    return [w for w in words if w]


def _build_rg_pattern(words: list[str]) -> tuple[str, bool]:
    """Build a ripgrep pattern from query words.

    Returns (pattern, is_regex). Single word uses fixed-string mode,
    multiple words build a regex alternation matching any word.
    """
    if len(words) <= 1:
        return words[0] if words else "", False
    escaped = [re.escape(w) for w in words]
    return "|".join(escaped), True


def _score_match(text_lower: str, words_lower: list[str]) -> int:
    """Count how many query words appear in the text."""
    return sum(1 for w in words_lower if w in text_lower)


def _best_snippet(
    text: str, text_lower: str, words_lower: list[str], query: str
) -> str:
    """Find the best snippet — prefer the region with the most word hits."""
    # Try exact full-query match first.
    idx = text_lower.find(query.lower())
    if idx >= 0:
        start = max(0, idx - 50)
        end = min(len(text), idx + len(query) + 50)
        snippet = text[start:end]
        if start > 0:
            snippet = "..." + snippet
        if end < len(text):
            snippet = snippet + "..."
        return snippet

    # Otherwise find the first word match.
    best_idx = len(text)
    for w in words_lower:
        idx = text_lower.find(w)
        if 0 <= idx < best_idx:
            best_idx = idx

    if best_idx >= len(text):
        best_idx = 0

    start = max(0, best_idx - 50)
    end = min(len(text), best_idx + 100)
    snippet = text[start:end]
    if start > 0:
        snippet = "..." + snippet
    if end < len(text):
        snippet = snippet + "..."
    return snippet


def rg_available() -> bool:
    return shutil.which("rg") is not None


def search_with_ripgrep(
    paths: list[Path],
    query: str,
    limit: int,
    extract_fn: ExtractFn,
    title_fn: TitleFn,
    *,
    glob_pattern: str | None = None,
) -> list[SearchMatch]:
    """Search using ripgrep with multi-word support.

    Args:
        paths: Files or directories to search.
        query: User query (may be multiple words).
        limit: Max results to return.
        extract_fn: Callback to extract (role, text, timestamp) from a parsed entry.
        title_fn: Callback to get (session_id, title) for a file path.
        glob_pattern: Optional ripgrep glob filter (e.g. "*.jsonl").
    """
    if not paths or limit <= 0:
        return []

    words = _tokenize_query(query)
    if not words:
        return []

    pattern, is_regex = _build_rg_pattern(words)
    words_lower = [w.lower() for w in words]

    cmd = ["rg", "--json", "-i"]
    if not is_regex:
        cmd += ["-F"]
    if glob_pattern:
        cmd += ["-g", glob_pattern]
    cmd += [pattern, *[str(p) for p in paths]]

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        encoding="utf-8",
        errors="replace",
    )

    collect_limit = limit * 10
    candidates: list[tuple[int, SearchMatch]] = []
    title_cache: dict[str, tuple[str, str]] = {}
    seen: set[tuple[str, str]] = set()

    try:
        assert proc.stdout is not None
        for raw in proc.stdout:
            raw = raw.strip()
            if not raw:
                continue

            try:
                event = _rg_json.loads(raw)
            except ValueError:
                continue

            if event.get("type") != "match":
                continue

            data = event.get("data", {})
            path_text = data.get("path", {}).get("text")
            line_text = data.get("lines", {}).get("text")
            if not path_text or not line_text:
                continue

            try:
                entry = _loads(line_text.strip())
            except ValueError:
                continue

            if not isinstance(entry, dict):
                continue

            extracted = extract_fn(entry)
            if extracted is None:
                continue

            role, text, timestamp = extracted
            if not text:
                continue

            text_lower = text.lower()
            score = _score_match(text_lower, words_lower)
            if score == 0:
                continue

            snippet = _best_snippet(text, text_lower, words_lower, query)

            cached = title_cache.get(path_text)
            if cached is None:
                cached = title_fn(Path(path_text))
                title_cache[path_text] = cached
            session_id, title = cached

            key = (session_id, snippet)
            if key in seen:
                continue
            seen.add(key)

            candidates.append(
                (
                    score,
                    SearchMatch(
                        session_id=session_id,
                        title=title,
                        timestamp=timestamp,
                        snippet=snippet,
                    ),
                )
            )

            if len(candidates) >= collect_limit:
                proc.terminate()
                break
    finally:
        if proc.stdout:
            proc.stdout.close()
        if proc.stderr:
            proc.stderr.close()
        try:
            proc.wait(timeout=2)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait(timeout=2)

    candidates.sort(key=lambda x: x[0], reverse=True)
    return [match for _, match in candidates[:limit]]


def search_python_fallback(
    files: list[Path],
    query: str,
    limit: int,
    extract_fn: ExtractFn,
    title_fn: TitleFn,
) -> list[SearchMatch]:
    """Pure-Python multi-word search fallback when ripgrep is unavailable."""
    if not files or limit <= 0:
        return []

    words = _tokenize_query(query)
    if not words:
        return []

    words_lower = [w.lower() for w in words]
    collect_limit = limit * 10
    candidates: list[tuple[int, SearchMatch]] = []

    for file_path in files:
        cached_title: tuple[str, str] | None = None

        try:
            with open(file_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    # Quick pre-check: does the raw line contain any query word?
                    line_lower = line.lower()
                    if not any(w in line_lower for w in words_lower):
                        continue

                    try:
                        entry = _loads(line)
                    except ValueError:
                        continue

                    if not isinstance(entry, dict):
                        continue

                    extracted = extract_fn(entry)
                    if extracted is None:
                        continue

                    role, text, timestamp = extracted
                    if not text:
                        continue

                    text_lower = text.lower()
                    score = _score_match(text_lower, words_lower)
                    if score == 0:
                        continue

                    snippet = _best_snippet(text, text_lower, words_lower, query)

                    if cached_title is None:
                        cached_title = title_fn(file_path)
                    session_id, title = cached_title

                    candidates.append(
                        (
                            score,
                            SearchMatch(
                                session_id=session_id,
                                title=title,
                                timestamp=timestamp,
                                snippet=snippet,
                            ),
                        )
                    )

                    if len(candidates) >= collect_limit:
                        break
        except Exception:
            continue

        if len(candidates) >= collect_limit:
            break

    candidates.sort(key=lambda x: x[0], reverse=True)
    return [match for _, match in candidates[:limit]]


def search_files(
    files: list[Path],
    query: str,
    limit: int,
    extract_fn: ExtractFn,
    title_fn: TitleFn,
    *,
    glob_pattern: str | None = None,
) -> list[SearchMatch]:
    """Search conversation files, using ripgrep if available."""
    if not files or limit <= 0:
        return []

    if rg_available():
        try:
            return search_with_ripgrep(
                files, query, limit, extract_fn, title_fn, glob_pattern=glob_pattern
            )
        except Exception:
            pass

    return search_python_fallback(files, query, limit, extract_fn, title_fn)


def search_dirs(
    dirs: list[Path],
    query: str,
    limit: int,
    extract_fn: ExtractFn,
    title_fn: TitleFn,
    *,
    glob_pattern: str | None = None,
) -> list[SearchMatch]:
    """Search directories, using ripgrep if available."""
    if not dirs or limit <= 0:
        return []

    if rg_available():
        try:
            return search_with_ripgrep(
                dirs, query, limit, extract_fn, title_fn, glob_pattern=glob_pattern
            )
        except Exception:
            pass

    # Fallback: collect files from dirs and use Python search
    files: list[Path] = []
    for d in dirs:
        if glob_pattern:
            files.extend(
                sorted(
                    d.rglob(glob_pattern), key=lambda p: p.stat().st_mtime, reverse=True
                )
            )
        else:
            files.extend(
                sorted(d.rglob("*"), key=lambda p: p.stat().st_mtime, reverse=True)
            )
    return search_python_fallback(files, query, limit, extract_fn, title_fn)
