"""Full-text search over agent sessions using tantivy."""

from __future__ import annotations

import json
import logging
import shutil
from dataclasses import dataclass
from pathlib import Path

import tantivy
from platformdirs import user_cache_path
from tantivy import Query, Occur

from .models import SessionMetadata

logger = logging.getLogger(__name__)

SCHEMA_VERSION = 1

CACHE_DIR = user_cache_path("agent-devlog") / "search_index"
MANIFEST_FILE = CACHE_DIR / "manifest.json"

# Field display names for the UI
FIELD_LABELS = {
    "project_name": "Project",
    "summary": "Summary",
    "first_message": "First message",
    "user_messages": "User message",
    "assistant_messages": "Assistant response",
    "git_branch": "Branch",
}

# Boost weights: higher = more relevant
FIELD_BOOSTS = {
    "project_name": 5.0,
    "summary": 3.0,
    "first_message": 2.5,
    "user_messages": 2.0,
    "assistant_messages": 1.0,
    "git_branch": 1.5,
}

SEARCH_FIELDS = list(FIELD_BOOSTS.keys())


@dataclass
class SearchHighlight:
    """A highlighted snippet from a matching field."""

    field_name: str
    field_label: str
    snippet_html: str


@dataclass
class SearchResult:
    """A search result with session reference and highlighted snippets."""

    session_id: str
    project_dir: str
    score: float
    highlights: list[SearchHighlight]


def _build_schema() -> tantivy.Schema:
    builder = tantivy.SchemaBuilder()
    builder.add_text_field("session_id", stored=True, tokenizer_name="raw")
    builder.add_text_field("project_dir", stored=True, tokenizer_name="raw")
    builder.add_text_field("project_name", stored=True, tokenizer_name="en_stem")
    builder.add_text_field("summary", stored=True, tokenizer_name="en_stem")
    builder.add_text_field("first_message", stored=True, tokenizer_name="en_stem")
    builder.add_text_field("user_messages", stored=True, tokenizer_name="en_stem")
    builder.add_text_field("assistant_messages", stored=True, tokenizer_name="en_stem")
    builder.add_text_field("git_branch", stored=True, tokenizer_name="default")
    return builder.build()


def _extract_session_text(
    file_path: Path,
) -> tuple[str, list[str], list[str]]:
    """Extract text content from a JSONL session file without full Pydantic parsing.

    Returns:
        (first_user_message, user_texts, assistant_texts)
    """
    first_user_message = ""
    user_texts: list[str] = []
    assistant_texts: list[str] = []

    try:
        with open(file_path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    raw = json.loads(line)
                except json.JSONDecodeError:
                    continue

                rtype = raw.get("type", "")
                if rtype not in ("user", "assistant"):
                    continue

                content = raw.get("message", {}).get("content", "")
                if isinstance(content, str):
                    text = content
                elif isinstance(content, list):
                    text = "\n".join(
                        item.get("text", "")
                        for item in content
                        if isinstance(item, dict) and item.get("type") == "text"
                    )
                else:
                    continue

                if not text.strip():
                    continue

                if rtype == "user":
                    if not first_user_message:
                        first_user_message = text
                    user_texts.append(text)
                elif rtype == "assistant":
                    assistant_texts.append(text)
    except (OSError, UnicodeDecodeError):
        logger.warning(f"Failed to read {file_path}", exc_info=True)

    return first_user_message, user_texts, assistant_texts


def _build_manifest(sessions: list[SessionMetadata]) -> dict:
    """Build a manifest dict mapping file paths to their mtime_ns."""
    files = {}
    for meta in sessions:
        try:
            p = Path(meta.file_path)
            files[meta.file_path] = p.stat().st_mtime_ns
        except OSError:
            pass
    return {"schema_version": SCHEMA_VERSION, "files": files}


def _load_manifest() -> dict | None:
    """Load the stored manifest from disk, or None if missing/corrupt."""
    try:
        with open(MANIFEST_FILE) as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return None


def _save_manifest(manifest: dict) -> None:
    """Save manifest to disk."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    with open(MANIFEST_FILE, "w") as f:
        json.dump(manifest, f)


class SessionSearchIndex:
    """Tantivy-backed full-text search index for sessions."""

    def __init__(self) -> None:
        self._schema = _build_schema()
        self._index: tantivy.Index | None = None
        self._ready = False

    @property
    def ready(self) -> bool:
        return self._ready

    def load_or_build(self, sessions: list[SessionMetadata]) -> None:
        """Load existing index from disk if fresh, otherwise rebuild."""
        current_manifest = _build_manifest(sessions)
        stored_manifest = _load_manifest()

        if stored_manifest == current_manifest and (CACHE_DIR / "meta.json").exists():
            try:
                self._index = tantivy.Index(self._schema, path=str(CACHE_DIR))
                self._index.reload()
                self._ready = True
                logger.info("Search index loaded from disk")
                return
            except Exception:
                logger.warning("Failed to load cached index, rebuilding", exc_info=True)

        self._rebuild(sessions, current_manifest)

    def _rebuild(
        self, sessions: list[SessionMetadata], manifest: dict
    ) -> None:
        """Build a fresh index from all sessions."""
        # Clear any existing index directory
        if CACHE_DIR.exists():
            shutil.rmtree(CACHE_DIR)
        CACHE_DIR.mkdir(parents=True, exist_ok=True)

        self._index = tantivy.Index(self._schema, path=str(CACHE_DIR))
        writer = self._index.writer()

        for meta in sessions:
            try:
                first_msg, user_texts, assistant_texts = _extract_session_text(
                    Path(meta.file_path)
                )
                project_name = " ".join(
                    filter(None, [meta.project_short_name, meta.project_path])
                )
                writer.add_document(
                    tantivy.Document(
                        session_id=[meta.session_id],
                        project_dir=[meta.project_dir],
                        project_name=[project_name],
                        summary=[meta.summary or ""],
                        first_message=[first_msg or meta.first_message_preview or ""],
                        user_messages=["\n\n".join(user_texts)],
                        assistant_messages=["\n\n".join(assistant_texts)],
                        git_branch=[meta.git_branch or ""],
                    )
                )
            except Exception:
                logger.warning(
                    f"Failed to index session {meta.session_id}", exc_info=True
                )

        writer.commit()
        writer.wait_merging_threads()
        self._index.reload()

        _save_manifest(manifest)

        self._ready = True
        logger.info(f"Search index built: {len(sessions)} sessions indexed")

    def search(self, query_text: str, limit: int = 50) -> list[SearchResult]:
        """Search for sessions matching the query text."""
        if not self._ready or not self._index or not query_text.strip():
            return []

        searcher = self._index.searcher()
        query = self._build_boosted_query(query_text)
        if query is None:
            return []

        search_result = searcher.search(query, limit)

        results: list[SearchResult] = []
        for score, doc_address in search_result.hits:
            doc = searcher.doc(doc_address)
            session_id = doc["session_id"][0]
            project_dir = doc["project_dir"][0]
            highlights = self._generate_highlights(searcher, query, doc)

            results.append(
                SearchResult(
                    session_id=session_id,
                    project_dir=project_dir,
                    score=score,
                    highlights=highlights,
                )
            )

        return results

    def _build_boosted_query(self, query_text: str) -> tantivy.Query | None:
        """Build a disjunction_max query with per-field boosting."""
        assert self._index is not None
        sub_queries = []
        for field_name, boost in FIELD_BOOSTS.items():
            try:
                field_query = self._index.parse_query(
                    query_text,
                    default_field_names=[field_name],
                )
                sub_queries.append(Query.boost_query(field_query, boost))
            except ValueError:
                continue

        if not sub_queries:
            return None

        return Query.boolean_query(
            [
                (
                    Occur.Should,
                    Query.disjunction_max_query(sub_queries, tie_breaker=0.3),
                )
            ]
        )

    def _generate_highlights(
        self,
        searcher: tantivy.Searcher,
        query: tantivy.Query,
        doc: tantivy.Document,
        max_highlights: int = 3,
    ) -> list[SearchHighlight]:
        """Generate highlighted snippets for the best matching fields."""
        highlights: list[SearchHighlight] = []

        for field_name in SEARCH_FIELDS:
            try:
                snippet_gen = tantivy.SnippetGenerator.create(
                    searcher, query, self._schema, field_name
                )
                snippet = snippet_gen.snippet_from_doc(doc)
                html = snippet.to_html()

                if "<b>" in html and html.strip():
                    highlights.append(
                        SearchHighlight(
                            field_name=field_name,
                            field_label=FIELD_LABELS[field_name],
                            snippet_html=html,
                        )
                    )
                    if len(highlights) >= max_highlights:
                        break
            except Exception:
                continue

        return highlights
