scripts/preindex_abbrevs.py

"""
Pre-index all (ย่อ) abbreviation markers.

Supports:
  - NVIDIA NIM (free DeepSeek, 40 RPM)  → --provider nim
  - DeepSeek API directly               → --provider deepseek

Usage:
    # Full run (default: NVIDIA NIM)
    python scripts/preindex_abbrevs.py

    # Resume with DeepSeek API
    python scripts/preindex_abbrevs.py --resume --provider deepseek

    # Test: first 10 pages
    python scripts/preindex_abbrevs.py --limit 10

Scans all pages in tipitaka_mcu.db for (ย่อ) markers.
For each page, calls LLM to expand ALL abbreviations in one request.
Stores JSON array in reference_markers (type='abbrev', marker_id='_default').
Supports resume — skips pages already in DB.

Log file: scripts/preindex_abbrevs.log
"""
import asyncio
import json
import logging
import re
import time
import argparse
import sqlite3
from pathlib import Path

from openai import AsyncOpenAI
from dotenv import load_dotenv
import os

# ── Paths ──
SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
DB_PATH = PROJECT_ROOT / "tipitaka_mcu.db"
LOG_PATH = SCRIPT_DIR / "preindex_abbrevs.log"
ENV_PATH = PROJECT_ROOT.parent / "tipitaka_context" / "nvidia" / ".env"

# ══════════════════════════════════════════════════════════════
# Config (loaded from .env)
# ══════════════════════════════════════════════════════════════
load_dotenv(ENV_PATH)

# NVIDIA NIM
NIM_API_KEY = os.getenv("TIPITAKA_API_KEY") or os.getenv("NVIDIA_API_KEY") or ""
NIM_BASE_URL = os.getenv("TIPITAKA_BASE_URL", "https://integrate.api.nvidia.com/v1")
NIM_MODEL = os.getenv("TIPITAKA_MODEL", "deepseek-ai/deepseek-v4-flash")

# DeepSeek API
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY") or os.getenv("DEEPSEEK_API_KEY") or ""
DEEPSEEK_BASE_URL = os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
DEEPSEEK_MODEL = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")

RATE_INTERVAL = float(os.getenv("TIPITAKA_RATE_INTERVAL", "1.6"))  # 40 RPM

MAX_RETRIES = 4
RETRY_DELAYS = [10, 20, 40, 80]

# ── Logging ──
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(LOG_PATH, encoding="utf-8"),
    ],
)
log = logging.getLogger(__name__)


# ══════════════════════════════════════════════════════════════
# Prompts
# ══════════════════════════════════════════════════════════════
SYSTEM_PROMPT = (
    "คุณคือผู้เชี่ยวชาญพระไตรปิฎก มจร.\n"
    "ตอบเป็น JSON array เท่านั้น โดยแต่ละรายการเป็น string"
)

USER_PROMPT_TEMPLATE = """ขยายความ (ย่อ) แต่ละแห่งในเนื้อหาด้านล่าง

ให้ตอบเป็น JSON array:
[
  "ข้อความที่ถูกย่อแห่งที่ 1...",
  "ข้อความที่ถูกย่อแห่งที่ 2...",
  ...
]
โดยเรียงลำดับตามที่ (ย่อ) ปรากฏในเนื้อหา
แต่ละรายการต้องเป็น string สั้นๆ อธิบายสิ่งที่ถูกย่อไว้
ตอบเฉพาะ JSON array เท่านั้น

เนื้อหา:
{content}"""


def build_prompt(content_text: str) -> str:
    return USER_PROMPT_TEMPLATE.format(content=content_text[:8000])


# ══════════════════════════════════════════════════════════════
# Parsing
# ══════════════════════════════════════════════════════════════
def parse_response(raw: str, expected_count: int) -> list[str]:
    """Parse JSON array from LLM response into list of expansion strings."""
    raw = re.sub(r'<think>.*?</think>', '', raw, flags=re.DOTALL).strip()
    raw = re.sub(r'^```(?:json)?\s*', '', raw)
    raw = re.sub(r'\s*```$', '', raw)
    raw = raw.strip()

    match = re.search(r'\[[\s\S]*\]', raw)
    if not match:
        raise ValueError(f"No JSON array found: {raw[:300]}")

    data = json.loads(match.group(0))
    if not isinstance(data, list):
        raise ValueError(f"Not an array: {type(data)}")

    result = []
    for item in data:
        if isinstance(item, str):
            result.append(item)
        elif isinstance(item, dict):
            # Try known keys, then concatenate all values
            for key in ("expansion", "expanded", "content", "text",
                        "explanation", "detail", "description", "meaning",
                        "abbrev", "ขยายความ", "ข้อความ", "ความหมาย", "answer"):
                val = item.get(key, "")
                if isinstance(val, str) and len(val.strip()) > 10:
                    result.append(val.strip())
                    break
            else:
                parts = [v for v in item.values()
                         if isinstance(v, str) and len(v.strip()) > 10]
                if parts:
                    result.append(" | ".join(parts))

    return result


# ══════════════════════════════════════════════════════════════
# Core: expand one page
# ══════════════════════════════════════════════════════════════
async def expand_page(
    client: AsyncOpenAI,
    model: str,
    vol: int, page: int, content_text: str,
    conn: sqlite3.Connection,
) -> bool:
    """Expand all (ย่อ) on one page. Stores result in DB. Returns True on success."""
    abbrev_count = content_text.count("(ย่อ)")
    if abbrev_count == 0:
        return False

    prompt = build_prompt(content_text)

    for attempt in range(MAX_RETRIES):
        try:
            resp = await client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                temperature=0.3,
                max_tokens=2000,
                timeout=120,
            )

            raw = resp.choices[0].message.content or ""

            expansions = parse_response(raw, abbrev_count)

            # Filter meaningful expansions
            expansions = [e for e in expansions if len(e.strip()) > 15]

            if not expansions:
                log.warning(f"    Empty! retrying... raw={raw[:150]}")
                if attempt < MAX_RETRIES - 1:
                    await asyncio.sleep(RETRY_DELAYS[attempt])
                    continue
                return False

            # Store in DB
            stored = json.dumps(expansions, ensure_ascii=False)
            conn.execute("""
                INSERT OR REPLACE INTO reference_markers
                    (volume_num, page_num, marker_id, type, content)
                VALUES (?, ?, '_default', 'abbrev', ?)
            """, (vol, page, stored))
            conn.commit()

            total_chars = sum(len(e) for e in expansions)
            log.info(f"  ✅ Vol {vol} P{page}: {len(expansions)} abbrevs, {total_chars}c")
            return True

        except json.JSONDecodeError as e:
            log.warning(f"  ⚠ Vol {vol} P{page} JSON error (attempt {attempt+1}): {e}")
            if attempt < MAX_RETRIES - 1:
                await asyncio.sleep(RETRY_DELAYS[attempt])
        except Exception as e:
            err_str = str(e)
            is_429 = "429" in err_str or "Too Many Requests" in err_str
            log.warning(f"  ⚠ Vol {vol} P{page} Error (attempt {attempt+1}): {e}")
            if attempt < MAX_RETRIES - 1:
                delay = RETRY_DELAYS[attempt] * (2 if is_429 else 1)
                await asyncio.sleep(delay)

    log.error(f"  ❌ Vol {vol} P{page}: Failed after {MAX_RETRIES} attempts")
    return False


# ══════════════════════════════════════════════════════════════
# Main
# ══════════════════════════════════════════════════════════════
async def main():
    parser = argparse.ArgumentParser(
        description="Pre-index (ย่อ) abbreviations via LLM"
    )
    parser.add_argument("--resume", action="store_true", help="Skip already-indexed pages")
    parser.add_argument("--force", action="store_true", help="Re-index even if already in DB")
    parser.add_argument("--limit", type=int, default=0, help="Max pages (0 = all)")
    parser.add_argument(
        "--provider", choices=["nim", "deepseek"], default="nim",
        help="LLM provider: nim (NVIDIA NIM, default) or deepseek (DeepSeek API)"
    )
    args = parser.parse_args()

    # ── Select provider config ──
    if args.provider == "nim":
        api_key = NIM_API_KEY
        base_url = NIM_BASE_URL
        model = NIM_MODEL
        provider_name = "NVIDIA NIM"
    else:
        api_key = DEEPSEEK_API_KEY or os.getenv("DEEPSEEK_API_KEY") or ""
        base_url = DEEPSEEK_BASE_URL
        model = DEEPSEEK_MODEL
        provider_name = "DeepSeek API"

    if not api_key:
        log.error(f"❌ No API key for provider '{args.provider}'. Check env file: {ENV_PATH}")
        if args.provider == "nim":
            log.error("   Set TIPITAKA_API_KEY or NVIDIA_API_KEY")
        else:
            log.error("   Set DEEPSEEK_API_KEY")
        return

    if not DB_PATH.exists():
        log.error(f"❌ DB not found: {DB_PATH}")
        return

    log.info(f"🚀 preindex_abbrevs | provider={provider_name} model={model}")
    log.info(f"   DB: {DB_PATH}")
    log.info(f"   Key: {api_key[:8]}...")
    log.info(f"   Rate: {RATE_INTERVAL}s")

    # ── Load pages ──
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    c = conn.cursor()

    c.execute("SELECT COUNT(*) as cnt FROM reference_markers WHERE type='abbrev'")
    log.info(f"   Existing: {c.fetchone()['cnt']} abbrev entries")

    c.execute("""
        SELECT v.volume_number, p.page_number, p.content_text
        FROM pages p
        JOIN volumes v ON p.volume_id = v.id
        WHERE p.content_text LIKE '%(ย่อ)%'
        ORDER BY v.volume_number, p.page_number
    """)
    rows = c.fetchall()
    log.info(f"   Pages: {len(rows)}")

    if args.limit > 0:
        rows = rows[:args.limit]
        log.info(f"   Limit: {args.limit}")

    if args.resume:
        c.execute("SELECT DISTINCT volume_num, page_num FROM reference_markers WHERE type='abbrev'")
        indexed = {(r["volume_num"], r["page_num"]) for r in c.fetchall()}
        rows = [r for r in rows if (r["volume_number"], r["page_number"]) not in indexed]
        log.info(f"   Resume: {len(indexed)} done, {len(rows)} remaining")
        if not rows:
            log.info("✅ All done!")
            conn.close()
            return
    elif args.force:
        c.execute("DELETE FROM reference_markers WHERE type='abbrev'")
        conn.commit()
        log.info("   Force: cleared all")

    # ── Setup client ──
    client = AsyncOpenAI(
        base_url=base_url,
        api_key=api_key,
        max_retries=0,
    )

    total = len(rows)
    success = fail = 0
    start = time.time()

    log.info(f"\n{'='*60}")
    log.info(f"Processing {total} pages...")
    log.info(f"{'='*60}\n")

    for i, row in enumerate(rows):
        vol = row["volume_number"]
        page = row["page_number"]
        text = row["content_text"] or ""

        # Rate limit
        if i > 0:
            elapsed = time.time() - start
            expected = i * RATE_INTERVAL
            wait = max(0, expected - elapsed)
            if wait > 0:
                await asyncio.sleep(wait)

        ok = await expand_page(client, model, vol, page, text, conn)

        if ok:
            success += 1
        else:
            fail += 1

        # Progress every 50 pages
        if (i + 1) % 50 == 0:
            elapsed = time.time() - start
            rate = (i + 1) / elapsed * 60
            remain = total - i - 1
            eta = remain / max(rate, 0.1) * 60
            log.info(f"📊 [{i+1}/{total}] {rate:.0f} pg/min | ETA: {eta/60:.1f}h")

    # ── Summary ──
    elapsed = time.time() - start
    log.info(f"\n{'='*60}")
    log.info(f"✅ Done!")
    log.info(f"   Success: {success}  Failed: {fail}")
    log.info(f"   Time: {elapsed:.0f}s ({elapsed/max(success,1):.1f}s/page)")
    log.info(f"   Log: {LOG_PATH}")

    conn.close()


if __name__ == "__main__":
    asyncio.run(main())