#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Linear Canal — Semantic Context Injector (UserPromptSubmit)

Reads the user's query, finds the most relevant ctx files + LanceDB chunks,
and outputs them as plain text so the PowerShell wrapper can inject them as
`additionalContext` BEFORE Claude processes the query.

From Claude's perspective: the context just "appears" — zero tool calls.
"""
from __future__ import annotations

import io
import json
import os
import re
import sys
from pathlib import Path
from typing import Any, Optional

# Force UTF-8 stdout so ctx file Unicode (arrows, em-dashes etc.) passes through
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
else:
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")

# ── Config ────────────────────────────────────────────────────────────────────

MAX_CTX_FILES   = 3        # max ctx files to inject per query
MAX_CTX_CHARS   = 2_500    # chars per ctx file (caps token budget)
MAX_CHUNKS      = 5        # max LanceDB chunks
MAX_CHUNK_CHARS = 500      # chars per chunk snippet
MAX_TOTAL_CHARS = 10_000   # hard cap on total output

STOP_WORDS = {
    "a","an","and","are","as","at","be","been","by","can","could","do","does",
    "for","from","has","have","how","i","in","is","it","its","of","on","or",
    "that","the","this","to","was","what","when","where","which","who","will",
    "with","work","works","would","you","your",
}

# These high-level ctx files are useful only when nothing more specific matches
_FALLBACK_CTX = ("ARCHITECTURE", "INTERFACES", "PATTERNS")

# ── Keyword helpers ───────────────────────────────────────────────────────────


def extract_keywords(text: str) -> list[str]:
    words = re.sub(r"[^a-z0-9_]", " ", text.lower()).split()
    return [w for w in words if w not in STOP_WORDS and len(w) > 2]


# ── Project discovery ─────────────────────────────────────────────────────────


def find_project(cwd: Path) -> tuple[Optional[str], Optional[Path]]:
    data_dir = cwd / ".linear_canal_data"
    if not data_dir.exists():
        return None, None
    for d in sorted(data_dir.iterdir()):
        if d.is_dir() and (d / "ctx" / "ARCHITECTURE.ctx").exists():
            return d.name, d / "ctx"
    return None, None


# ── Ctx file keyword matching ─────────────────────────────────────────────────


def match_ctx_files(keywords: list[str], ctx_dir: Path) -> list[Path]:
    """Score ctx files by keyword overlap, skip generic high-level files."""
    if not keywords:
        return []

    scored: list[tuple[float, Path]] = []
    for f in ctx_dir.glob("*.ctx"):
        stem = f.stem.upper()
        if stem in _FALLBACK_CTX:
            continue  # handled separately as fallback
        stem_l = f.stem.lower()
        score = 0.0
        for kw in keywords:
            kw_l = kw.lower()
            if kw_l == stem_l:
                score += 2.0          # exact match
            elif kw_l in stem_l or stem_l in kw_l:
                score += 1.0          # partial match
        if score > 0:
            scored.append((score, f))

    scored.sort(key=lambda x: x[0], reverse=True)
    matched = [f for _, f in scored[:MAX_CTX_FILES]]

    # If fewer than 2 specific matches, prepend ARCHITECTURE as context anchor
    if len(matched) < 2:
        arch = ctx_dir / "ARCHITECTURE.ctx"
        if arch.exists():
            matched.insert(0, arch)

    return matched[:MAX_CTX_FILES]


# ── LanceDB search (vector-first, FTS fallback) ───────────────────────────────

_DIM_TO_MODEL = {
    384: ("BAAI/bge-small-en-v1.5", "minimal"),
    768: ("nomic-ai/nomic-embed-text-v1.5", "standard"),
}


def _open_lancedb() -> tuple[Any, str, int] | tuple[None, None, None]:
    """Open the global LanceDB store. Returns (table, model_name, dim) or (None,None,None)."""
    try:
        import warnings  # noqa: PLC0415
        import lancedb  # noqa: PLC0415
        db_path = Path.home() / ".linear-canal" / "data" / "lancedb_global_store"
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            db = lancedb.connect(str(db_path))
            tables = db.table_names()
        for dim, (model_name, _) in _DIM_TO_MODEL.items():
            table_name = f"lc_chunks_{dim}"
            if table_name in tables:
                return db.open_table(table_name), model_name, dim
    except Exception:
        pass
    return None, None, None


def lancedb_search(query: str, project_id: str) -> list[dict]:
    """Vector search (semantic) with FTS fallback. Returns [] on any failure."""
    tbl, model_name, dim = _open_lancedb()
    if tbl is None:
        return []

    where = f"project_id = '{project_id}'"

    # ── Primary: vector (embedding) search ───────────────────────────────────
    try:
        from fastembed import TextEmbedding  # noqa: PLC0415
        model = TextEmbedding(model_name)
        vec = list(model.embed([query]))[0].tolist()
        rows = (
            tbl.search(vec)
               .where(where, prefilter=True)
               .limit(MAX_CHUNKS)
               .to_list()
        )
        if rows:
            # Attach similarity score for formatting
            for r in rows:
                dist = float(r.get("_distance", 1.0))
                r["_sim"] = max(0.0, 1.0 - dist / 2.0)
            return rows
    except Exception:
        pass

    # ── Fallback: FTS (BM25) ─────────────────────────────────────────────────
    try:
        rows = (
            tbl.search(query, query_type="fts")
               .where(where, prefilter=True)
               .limit(MAX_CHUNKS)
               .to_list()
        )
        if rows:
            return rows
    except Exception:
        pass

    return []


# ── Output formatting ─────────────────────────────────────────────────────────


def format_ctx_block(query: str, ctx_files: list[Path], chunks: list[dict],
                     project_id: str = "") -> str:
    parts: list[str] = []

    # Always prepend MCP tool instructions so the agent knows to use semantic search
    pid = project_id or "this project"
    parts.append(
        f"[Linear Canal — {pid}] Semantic context pre-loaded for your query.\n"
        f"MCP TOOLS AVAILABLE — call these instead of Glob/Grep/Read on source files:\n"
        f"  lc_context(task_description='<full task description>')  ← CALL THIS FIRST\n"
        f"  lc_search(query='<concept, symbol, or pattern>')        ← for targeted search\n"
        f"  lc_impact(file_path='<path>')                          ← before editing a file\n"
        f"The following context was retrieved via vector search (LanceDB) for your query:"
    )
    parts.append("")

    if ctx_files:
        parts.append(f'=== Pre-indexed context for: "{query[:80]}" ===')
        parts.append("")
        for f in ctx_files:
            try:
                content = f.read_text(encoding="utf-8", errors="ignore").strip()
                parts.append(f"=== {f.stem}.ctx ===")
                if len(content) > MAX_CTX_CHARS:
                    parts.append(content[:MAX_CTX_CHARS])
                    parts.append(f"(... full file: .linear_canal_data/*/ctx/{f.name})")
                else:
                    parts.append(content)
                parts.append("")
            except OSError:
                pass

    if chunks:
        parts.append("[Linear Canal] Semantic code search results:")
        parts.append("")
        for c in chunks:
            fp    = c.get("file_path", "?")
            sim   = c.get("_sim", c.get("_relevance_score", 0.0))
            text  = (c.get("content", "") or "").strip()[:MAX_CHUNK_CHARS]
            sym   = c.get("symbol_name", "")
            start = c.get("start_line", "")
            label = f"{fp}:{start}  [{sym}]" if sym else f"{fp}:{start}"
            sim_s = f"  (similarity: {sim:.2f})" if sim else ""
            parts.append(f"--- {label}{sim_s}")
            parts.append(text)
            parts.append("")

    output = "\n".join(parts).strip()
    return output[:MAX_TOTAL_CHARS]


# ── Main ──────────────────────────────────────────────────────────────────────


def basic_reminder(project_id: str) -> str:
    return (
        f"[Linear Canal active — project: {project_id}]\n"
        f"MANDATORY — use Linear Canal tools before ANY file exploration:\n"
        f"  MCP tool (preferred): lc_context(task_description='<your task here>')\n"
        f"  MCP tool (search):    lc_search(query='<concept or symbol>')\n"
        f"  Fallback (no MCP):    Read .linear_canal_data/{project_id}/ctx/ARCHITECTURE.ctx\n"
        f"Do NOT use Glob, Grep, or raw Read on source files before calling lc_context."
    )


def main() -> None:
    # ── Read stdin payload ────────────────────────────────────────────────────
    query = ""
    try:
        raw = sys.stdin.read().strip()
        if raw:
            payload = json.loads(raw)
            query = str(payload.get("prompt", "")).strip()
    except Exception:
        pass

    cwd = Path(os.getcwd())
    project_id, ctx_dir = find_project(cwd)

    # No project → silent exit
    if not project_id or not ctx_dir:
        sys.exit(0)

    # Short / slash-command queries → basic reminder only
    if not query or len(query) < 8 or query.lstrip().startswith("/"):
        print(basic_reminder(project_id))
        sys.exit(0)

    keywords = extract_keywords(query)

    # ── Ctx file keyword match ────────────────────────────────────────────────
    ctx_files = match_ctx_files(keywords, ctx_dir)

    # ── LanceDB semantic search (vector → FTS fallback) ──────────────────────
    chunks = lancedb_search(query, project_id)

    # ── Format and emit ───────────────────────────────────────────────────────
    output = format_ctx_block(query, ctx_files, chunks, project_id=project_id)

    if output:
        print(output)
    else:
        print(basic_reminder(project_id))

    sys.exit(0)


if __name__ == "__main__":
    main()
