"""Busca semântica RAG do projeto.

Consulta `<schema>.code_chunks` no Postgres (via túnel SSH) combinando dois
sinais e devolvendo faixas de linha, para o agente ler só o trecho relevante
em vez do arquivo inteiro.

Uso:
    rag/search_tigre.sh "como funciona o export FCPXML"
    rag/search_tigre.sh "FCPXMLValidator" --snippet
    rag/search_tigre.sh "wizard" --module TigreAppUI --json
    rag/search_tigre.sh --map TigreAI          # inventário do módulo

Como módulo:
    from search import rag_search
    rag_search("consulta", top_k=5)

Por que busca híbrida
---------------------
A busca puramente densa erra nomes exatos: procurar `FCPXMLValidator` não
trazia `FCPXMLValidator.swift`, e `SilenceCutPipeline.swift` não aparecia em
nenhum top-5. Por isso rodamos duas listas em paralelo — densa (embedding) e
lexical (pg_trgm sobre os símbolos declarados) — e fundimos com RRF, que
soma 1/(k+posição) de cada lista e portanto não exige normalizar escalas
diferentes. Ver rag/bench.py para os números antes/depois.
"""

import argparse
import json
import os
import sys

import psycopg2
from dotenv import load_dotenv

RAG_DIR = os.path.dirname(os.path.abspath(__file__))
# Mesma mecânica do indexador: respeita RAG_ENV_FILE para permitir wrapper
# por sistema (ex: search_tigre.sh com RAG_ENV_FILE=tigre.env).
_load_path = os.environ.get("RAG_ENV_FILE", os.path.join(RAG_DIR, ".env"))
load_dotenv(_load_path)

if RAG_DIR not in sys.path:
    sys.path.insert(0, RAG_DIR)

from index_code import QUERY_PREFIX, _embed  # noqa: E402

# Constante do Reciprocal Rank Fusion e pesos por lista. Os valores saíram de
# uma varredura sobre o conjunto dourado (rag/bench.py), não da literatura: o
# k=60 clássico achata demais a curva para um corpus deste tamanho e deixava
# um 1º lugar isolado perder para um medíocre presente em duas listas.
# k=8 e peso 1.5 no resumo vêm de uma varredura (k × pesos) sobre as 18
# consultas de rag/bench.py. Não é o pico da grade — é o meio de um platô
# estável (k de 5 a 8 dá o mesmo recall, e em k=8 o resultado não muda com o
# peso do resumo entre 1.0 e 2.0), escolhido justamente para não sobreajustar
# um conjunto dourado pequeno. Ao mexer nesses números, rode rag/bench.sh.
RRF_K = 8
W_DENSE = 1.0
W_LEX = 1.0       # multiplicado pela similaridade bruta do casamento
W_SUMMARY = 1.5   # o resumo é curto e preciso: um acerto ali vale mais

# Piso do casamento lexical. Calibrado por medição: com os probes abaixo, um
# nome de tipo procurado diretamente pontua 1.00 e o segundo colocado no
# máximo 0.95, enquanto NENHUMA consulta conceitual em português passa de
# 0.35. Ou seja, 0.5 faz o lado lexical se calar quando não tem nada a dizer,
# em vez de afogar a lista densa com ruído.
LEX_MIN = 0.5

# Quantos chunks do mesmo arquivo podem ocupar o top-k. Sem esse limite um
# arquivo grande toma todos os lugares: numa medição o WizardModel.swift
# ficou com 3 dos 5 resultados, gastando token sem trazer nada novo.
MAX_PER_FILE = 2

# Quanto código o modo --snippet mostra por resultado.
SNIPPET_CHARS = 300

# Colunas extras introduzidas pela migration 002. O banco legado `doza` não
# as tem — este mesmo código atende os dois, então descobrimos quais existem
# e substituímos as ausentes por NULL, em vez de quebrar a busca do legado.
_EXTRA_COLS = ("start_line", "end_line", "symbols", "module", "kind")
_cols_cache = {}


def _available_cols(cur, schema):
    if schema not in _cols_cache:
        cur.execute(
            """SELECT column_name FROM information_schema.columns
               WHERE table_schema = %s AND table_name = 'code_chunks'""",
            (schema,),
        )
        _cols_cache[schema] = {r[0] for r in cur.fetchall()}
    return _cols_cache[schema]


def _select_cols(cur, schema):
    have = _available_cols(cur, schema)
    extra = ", ".join(f"c.{c}" if c in have else f"NULL AS {c}"
                      for c in _EXTRA_COLS)
    return f"c.id, c.file_path, c.content, {extra}"


def _db_connect():
    return psycopg2.connect(
        host=os.environ.get("RAG_DB_HOST", "127.0.0.1"),
        port=os.environ.get("RAG_DB_PORT", "5433"),
        dbname=os.environ.get("RAG_DB_NAME", "rag_doza"),
        user=os.environ.get("RAG_DB_USER", "rag_doza_indexer"),
        password=os.environ["RAG_DB_PASSWORD"],
        connect_timeout=5,
    )


def _schema():
    return os.environ.get("RAG_DB_SCHEMA", "doza")


def _qid(schema):
    """Schema como identificador SQL seguro (aspas duplas) — aceita guífen
    (ex: `jhonny-rag`) sem quebrar a sintaxe. Retroativo: produtos sem guífen
    voltam idênticos (`doza` -> `"doza"`)."""
    return '"' + schema.replace('"', '""') + '"'


def _filters(module, path, ext, have=()):
    """Cláusulas de escopo aplicadas antes do ranqueamento."""
    clauses, params = [], []
    if module and "module" in have:
        clauses.append("c.module = %s")
        params.append(module)
    if path:
        clauses.append("c.file_path ILIKE %s")
        params.append(f"%{path}%")
    if ext:
        clauses.append("c.file_path LIKE %s")
        params.append(f"%{ext}")
    return (" AND " + " AND ".join(clauses) if clauses else ""), params


def _probe_terms(query):
    """Termos que o lado lexical tenta casar contra os símbolos.

    Só dois: a consulta inteira e a versão sem espaços — é esta que faz
    "silence cut pipeline" casar 1.00 com `SilenceCutPipeline`.

    Palavra a palavra NÃO funciona, por mais tentador que pareça: termos
    portugueses comuns casam com qualquer lista de símbolos ("fora" pontuou
    1.00 contra `WizardStep`, "forma" 0.83 contra meia base) e o ruído
    expulsava do top-5 acertos legítimos do lado denso — o `FacePerceiver`
    era o 1º da lista densa para "detectar rostos" e sumia do resultado.
    """
    terms = [query, query.replace(" ", "")]
    return list(dict.fromkeys(t for t in terms if t))


def _row_to_dict(row):
    return {
        "id": row[0], "file_path": row[1], "content": row[2],
        "start_line": row[3], "end_line": row[4],
        "symbols": row[5], "module": row[6], "kind": row[7],
    }


def _dense(cur, schema, q_emb, limit, where, params):
    # O HNSW só explora `ef_search` candidatos; 64 dá folga sobre um top-k
    # pequeno sem custar latência perceptível nesta escala.
    cur.execute("SET LOCAL hnsw.ef_search = 64")
    cur.execute(
        f"""
        SELECT {_select_cols(cur, schema)}, 1 - (c.embedding <=> %s::vector) AS s
        FROM {_qid(schema)}.code_chunks c
        WHERE c.embedding IS NOT NULL {where}
        ORDER BY c.embedding <=> %s::vector
        LIMIT %s
        """,
        [q_emb] + params + [q_emb, limit],
    )
    return [(_row_to_dict(r), float(r[8])) for r in cur.fetchall()]


def _lexical(cur, schema, terms, limit, where, params):
    # Casa contra `symbols` (nome do arquivo + tipo + membros declarados),
    # nunca contra o caminho inteiro: nomes de diretório como `Pipeline/`
    # casariam com meia base e afogariam o sinal.
    have = _available_cols(cur, schema)
    stem = "regexp_replace(c.file_path, '^.*/|\\.[^.]*$', '', 'g')"
    target = f"coalesce(c.symbols, {stem})" if "symbols" in have else stem
    # Desempate estável: pela linha inicial onde ela existe, senão pelo id.
    tiebreak = "c.start_line" if "start_line" in have else "c.id"
    score = f"(SELECT max(word_similarity(t, {target})) FROM unnest(%s::text[]) t)"
    cur.execute(
        f"""
        SELECT {_select_cols(cur, schema)}, {score} AS s
        FROM {_qid(schema)}.code_chunks c
        WHERE {score} >= %s {where}
        ORDER BY s DESC, {tiebreak} ASC
        LIMIT %s
        """,
        [terms] + [terms, LEX_MIN] + params + [limit],
    )
    return [(_row_to_dict(r), float(r[8])) for r in cur.fetchall()]


def _summary_dense(cur, schema, q_emb, limit, where, params):
    """Terceira lista: busca sobre o resumo do arquivo, não sobre o código.

    Devolve o melhor trecho de cada arquivo cujo resumo casa com a consulta.
    Existe porque a lista por trecho é dominada por arquivos grandes — eles
    casam morno com tudo e empurram para baixo arquivos pequenos e exatos.
    """
    cur.execute("SELECT to_regclass(%s)", (f"{_qid(schema)}.file_index",))
    if cur.fetchone()[0] is None:
        return []
    cur.execute("SET LOCAL hnsw.ef_search = 64")
    cur.execute(
        f"""
        SELECT file_path, 1 - (summary_embedding <=> %s::vector) AS s
        FROM {_qid(schema)}.file_index
        WHERE summary_embedding IS NOT NULL
        ORDER BY summary_embedding <=> %s::vector
        LIMIT %s
        """,
        (q_emb, q_emb, limit),
    )
    hits = cur.fetchall()
    if not hits:
        return []
    order = {fp: i for i, (fp, _) in enumerate(hits)}
    raw = {fp: s for fp, s in hits}

    # Para cada arquivo achado pelo resumo, o trecho mais próximo da consulta.
    cur.execute(
        f"""
        SELECT DISTINCT ON (c.file_path) {_select_cols(cur, schema)}
        FROM {_qid(schema)}.code_chunks c
        WHERE c.file_path = ANY(%s) AND c.embedding IS NOT NULL {where}
        ORDER BY c.file_path, c.embedding <=> %s::vector
        """,
        [list(order)] + params + [q_emb],
    )
    rows = [_row_to_dict(r) for r in cur.fetchall()]
    rows.sort(key=lambda r: order[r["file_path"]])
    return [(r, raw[r["file_path"]]) for r in rows]


def _fuse(ranked_lists):
    """Reciprocal Rank Fusion ponderada sobre listas já ordenadas.

    RRF puro só olha a posição, e isso empatava o 1º lugar lexical com o 1º
    dense — deixando de fora casos em que só o lado lexical acertava
    (`AtomicFileIO` para "gravar arquivo de forma atomica"). Por isso cada
    lista traz uma função de peso: a lexical escala com a similaridade
    bruta, de modo que um casamento quase exato de nome vence o empate.
    """
    scores, best = {}, {}
    for results, weight_of in ranked_lists:
        for rank, (row, raw) in enumerate(results, 1):
            contrib = weight_of(raw) / (RRF_K + rank)
            scores[row["id"]] = scores.get(row["id"], 0.0) + contrib
            best.setdefault(row["id"], row)
    ordered = sorted(scores.items(), key=lambda kv: -kv[1])
    return [dict(best[i], score=s) for i, s in ordered]


def _dedupe(rows, top_k, max_per_file=MAX_PER_FILE):
    """Limita chunks por arquivo; o excedente vira uma nota de localização."""
    kept, counts, extras = [], {}, {}
    for row in rows:
        fp = row["file_path"]
        if counts.get(fp, 0) < max_per_file:
            counts[fp] = counts.get(fp, 0) + 1
            row["also_at"] = []
            kept.append(row)
        else:
            extras.setdefault(fp, []).append((row["start_line"], row["end_line"]))
    for row in kept:
        row["also_at"] = extras.get(row["file_path"], [])[:3]
    return kept[:top_k]


def _attach_map(cur, schema, rows):
    """Anexa tipo principal e resumo de `file_index`, quando existir."""
    cur.execute("SELECT to_regclass(%s)", (f"{_qid(schema)}.file_index",))
    if cur.fetchone()[0] is None or not rows:
        return rows
    paths = list({r["file_path"] for r in rows})
    cur.execute(
        f"SELECT file_path, main_type, summary FROM {_qid(schema)}.file_index "
        f"WHERE file_path = ANY(%s)",
        (paths,),
    )
    info = {p: (t, s) for p, t, s in cur.fetchall()}
    for row in rows:
        main_type, summary = info.get(row["file_path"], (None, None))
        row["main_type"] = main_type
        row["summary"] = summary
    return rows


def rag_search(query, top_k=5, module=None, path=None, ext=None, dense_only=False):
    """Busca híbrida. Devolve dicts com file_path, faixa de linhas e score."""
    schema = _schema()
    # Ampliamos o pool antes de fundir e deduplicar: o item certo pode estar
    # em 8º numa lista e em 1º na outra, e a fusão é que o traz para cima.
    pool = max(top_k * 4, 20)

    conn = _db_connect()
    try:
        with conn.cursor() as cur:
            where, params = _filters(module, path, ext,
                                     _available_cols(cur, schema))
            q_emb = _embed(query, prefix=QUERY_PREFIX)
            lists = [(_dense(cur, schema, q_emb, pool, where, params),
                      lambda raw: W_DENSE)]
            if not dense_only:
                lists.append((_lexical(cur, schema, _probe_terms(query),
                                       pool, where, params),
                              lambda raw: W_LEX * (0.5 + raw)))
                lists.append((_summary_dense(cur, schema, q_emb, pool,
                                             where, params),
                              lambda raw: W_SUMMARY))
            rows = _dedupe(_fuse(lists), top_k)
            rows = _attach_map(cur, schema, rows)
        conn.commit()
    finally:
        conn.close()
    return rows


def map_files(term=None, module=None, limit=60):
    """Inventário de arquivos — responde "onde fica X" sem corpo de código."""
    schema = _schema()
    conn = _db_connect()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT to_regclass(%s)", (f"{_qid(schema)}.file_index",))
            if cur.fetchone()[0] is None:
                return []
            clauses, params = [], []
            if module:
                clauses.append("module = %s")
                params.append(module)
            if term:
                clauses.append("(file_path ILIKE %s OR main_type ILIKE %s "
                               "OR summary ILIKE %s)")
                params += [f"%{term}%"] * 3
            where = "WHERE " + " AND ".join(clauses) if clauses else ""
            cur.execute(
                f"""SELECT file_path, module, main_type, summary, n_lines
                    FROM {_qid(schema)}.file_index {where}
                    ORDER BY module, file_path LIMIT %s""",
                params + [limit],
            )
            return [
                {"file_path": r[0], "module": r[1], "main_type": r[2],
                 "summary": r[3], "n_lines": r[4]}
                for r in cur.fetchall()
            ]
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Formatação
# ---------------------------------------------------------------------------

# Linhas que não descrevem nada: delimitadores de docstring, cercas e
# separadores. Aparecem no topo de arquivos Python do índice legado.
_NOISE = {'"""', "'''", "# ---", "---", "/*", "*/", "{", "}", "*"}


def _first_doc_line(content):
    """Primeira linha que serve de descrição do trecho.

    Usada quando o arquivo não tem resumo em `file_index` — caso do banco
    legado, que não passou pela migration 002.
    """
    for line in content.splitlines():
        s = line.strip()
        if s.startswith(("///", "//", "#", "*")):
            s = s.lstrip("/#* ").strip()
            if s and s not in _NOISE:
                return s
    for line in content.splitlines():
        s = line.strip().strip("\"'").strip()
        if s and s not in _NOISE and not s.startswith("import"):
            return s
    return ""


def format_results(results, mode="map"):
    """Renderiza o resultado no modo pedido.

    O padrão é `map`: caminho + faixa de linhas + uma linha de descrição, sem
    nenhum corpo de código. É o modo mais barato e quase sempre suficiente —
    com a faixa em mãos o agente lê exatamente a janela que precisa.
    """
    if not results:
        return "[RAG vazio, buscando local]\n"

    if mode == "json":
        return json.dumps(results, ensure_ascii=False, indent=2) + "\n"

    out = []
    for i, r in enumerate(results, 1):
        # O banco legado não tem faixa de linhas; sem ela, mostramos só o
        # caminho em vez de um ":None-None" que não serve para nada.
        loc = r["file_path"]
        if r.get("start_line") and r.get("end_line"):
            loc += f":{r['start_line']}-{r['end_line']}"
        out.append(f"{i}  {r['score']:.2f}  {loc}")

        desc = r.get("summary") or _first_doc_line(r["content"])
        # A regra "uma classe por arquivo" faz o tipo principal repetir o
        # nome do arquivo quase sempre; imprimir os dois é token jogado fora.
        stem = os.path.splitext(os.path.basename(r["file_path"]))[0]
        label = r.get("main_type") or ""
        if label == stem:
            label = ""
        if label and desc:
            out.append(f"          {label} · {desc[:78]}")
        elif desc:
            out.append(f"          {desc[:88]}")

        spans = [f"{a}-{b}" for a, b in r.get("also_at", []) if a and b]
        if spans:
            out.append(f"          (+ tambem em {', '.join(spans)})")

        if mode == "snippet":
            # Corta em fronteira de LINHA: um trecho de código partido no
            # meio de um identificador não ajuda ninguém a decidir se vale
            # abrir o arquivo.
            body, size = [], 0
            for ln in r["content"].splitlines():
                if body and size + len(ln) > SNIPPET_CHARS:
                    body.append("…")
                    break
                body.append(ln)
                size += len(ln) + 1
            out.append("".join(f"          | {ln}\n" for ln in body))
        elif mode == "full":
            out.append("".join(f"          | {ln}\n" for ln in r["content"].splitlines()))
    return "\n".join(out) + "\n"


def format_map(rows):
    if not rows:
        return "[RAG vazio, buscando local]\n"
    out = []
    current = None
    for r in rows:
        if r["module"] != current:
            current = r["module"]
            out.append(f"\n{current or '(sem modulo)'}")
        name = os.path.basename(r["file_path"])
        desc = (r["summary"] or "")[:78]
        out.append(f"  {name:<38} {r['n_lines']:>5}L  {desc}")
    return "\n".join(out) + "\n"


def main():
    ap = argparse.ArgumentParser(
        description="Busca RAG hibrida (densa + lexical, fundidas com RRF)")
    ap.add_argument("query", nargs="?", help="consulta")
    ap.add_argument("top_k", nargs="?", type=int, default=5)
    ap.add_argument("--snippet", action="store_true", help="mostra 300 chars do trecho")
    ap.add_argument("--full", action="store_true", help="mostra o trecho inteiro")
    ap.add_argument("--json", action="store_true", help="saida estruturada")
    ap.add_argument("--module", help="restringe a um modulo (ex: TigreAI)")
    ap.add_argument("--path", help="restringe a caminhos contendo este texto")
    ap.add_argument("--ext", help="restringe a uma extensao (ex: .swift)")
    ap.add_argument("--map", dest="map_term", nargs="?", const="",
                    help="inventario de arquivos em vez de busca por trecho")
    ap.add_argument("--dense-only", action="store_true",
                    help="desliga o lado lexical (para comparacao)")
    args = ap.parse_args()

    if args.map_term is not None:
        print(format_map(map_files(term=args.map_term or None, module=args.module)))
        return

    if not args.query:
        ap.error("informe a consulta, ou use --map")

    mode = ("json" if args.json else "full" if args.full
            else "snippet" if args.snippet else "map")
    results = rag_search(args.query, args.top_k, module=args.module,
                         path=args.path, ext=args.ext, dense_only=args.dense_only)
    print(format_results(results, mode=mode))


if __name__ == "__main__":
    main()