#!/usr/bin/env python3

import argparse

from index_common import embed_text, get_conn, vector_literal


def sanitize(text: str) -> str:
    return text.replace('```', '` ` `').strip()


def query(text: str, top_k: int = 5, max_chars: int = 2500) -> str:
    embedding = embed_text(text)
    vector = vector_literal(embedding)

    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                '''
                SELECT source, content
                FROM memory_primary
                ORDER BY embedding <=> %s::vector
                LIMIT %s
                ''',
                (vector, top_k),
            )
            rows = cur.fetchall()

    parts: list[str] = []
    total_len = 0
    for source, content in rows:
        snippet = sanitize(content[:600])
        block = f'<retrieved_context source="{source}">\n{snippet}\n</retrieved_context>'
        if total_len + len(block) > max_chars:
            break
        parts.append(block)
        total_len += len(block)

    return '\n\n'.join(parts)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument('query', nargs='+')
    parser.add_argument('--top-k', type=int, default=5)
    parser.add_argument('--max-chars', type=int, default=2500)
    args = parser.parse_args()

    text = ' '.join(args.query).strip()
    if not text:
        return
    print(query(text, top_k=args.top_k, max_chars=args.max_chars))


if __name__ == '__main__':
    main()