#!/usr/bin/env python3

from pathlib import Path

from index_common import (
    VAULT_ROOT,
    embed_text,
    get_conn,
    index_lock,
    is_excluded,
    normalize_rel,
    sha256_text,
    vector_literal,
)

PRIMARY_DIRS = ['01_Projects', '02_Areas']
MIN_TEXT_LEN = 50


def upsert_primary(cur, rel: str, text: str, embedding: list[float]) -> None:
    cur.execute(
        '''
        INSERT INTO memory_primary (id, source, content, content_hash, embedding)
        VALUES (%s, %s, %s, %s, %s::vector)
        ON CONFLICT (id) DO UPDATE SET
          source = EXCLUDED.source,
          content = EXCLUDED.content,
          content_hash = EXCLUDED.content_hash,
          embedding = EXCLUDED.embedding,
          updated_at = now()
        ''',
        (rel, rel, text, sha256_text(text), vector_literal(embedding)),
    )


def run() -> None:
    with index_lock():
        with get_conn() as conn:
            with conn.cursor() as cur:
                valid_ids: set[str] = set()
                secure_ids: set[str] = set()
                failed_ids: list[str] = []

                print(f'[INFO] 开始全量扫描: {VAULT_ROOT}')
                for dir_name in PRIMARY_DIRS:
                    target = VAULT_ROOT / dir_name
                    if not target.exists():
                        continue

                    for md_file in target.rglob('*.md'):
                        rel = normalize_rel(md_file)
                        reason = is_excluded(md_file)
                        if reason:
                            secure_ids.add(rel)
                            cur.execute(
                                '''
                                INSERT INTO memory_secure_audit (id, source, risk)
                                VALUES (%s, %s, %s)
                                ON CONFLICT (id) DO UPDATE SET
                                  risk = EXCLUDED.risk,
                                  updated_at = now()
                                ''',
                                (rel, rel, reason),
                            )
                            cur.execute('DELETE FROM memory_primary WHERE id=%s', (rel,))
                            print(f'[QUARANTINED] {rel}')
                            continue

                        text = md_file.read_text(encoding='utf-8', errors='ignore')
                        if len(text.strip()) < MIN_TEXT_LEN:
                            cur.execute('DELETE FROM memory_primary WHERE id=%s', (rel,))
                            cur.execute('DELETE FROM memory_secure_audit WHERE id=%s', (rel,))
                            print(f'[PRUNED] {rel}')
                            continue

                        try:
                            embedding = embed_text(text)
                        except Exception as error:
                            failed_ids.append(rel)
                            print(f'[EMBED-ERROR] {rel} :: {error}')
                            continue
                        upsert_primary(cur, rel, text, embedding)
                        cur.execute('DELETE FROM memory_secure_audit WHERE id=%s', (rel,))
                        valid_ids.add(rel)
                        print(f'[UPSERTED] {rel}')

                cur.execute('SELECT id FROM memory_primary')
                db_primary_ids = {row[0] for row in cur.fetchall()}
                stale_primary_ids = sorted(db_primary_ids - valid_ids - set(failed_ids))
                for stale_id in stale_primary_ids:
                    cur.execute('DELETE FROM memory_primary WHERE id=%s', (stale_id,))

                cur.execute('SELECT id FROM memory_secure_audit')
                db_secure_ids = {row[0] for row in cur.fetchall()}
                stale_secure_ids = sorted(db_secure_ids - secure_ids)
                for stale_id in stale_secure_ids:
                    cur.execute('DELETE FROM memory_secure_audit WHERE id=%s', (stale_id,))

            conn.commit()

        print(
            f'[DONE] 全量索引完成 primary={len(valid_ids)} secure={len(secure_ids)} '
            f'clean_primary={len(stale_primary_ids)} clean_secure={len(stale_secure_ids)} '
            f'failed={len(failed_ids)}'
        )


if __name__ == '__main__':
    run()