scripts/setup.py

"""
scripts/setup.py
────────────────
One-command pipeline: download → preprocess → validate → ingest → index

Usage
-----
    python scripts/setup.py                     # auto-detect Kaggle or use seed
    python scripts/setup.py --source kaggle     # force Kaggle download
    python scripts/setup.py --source seed       # use bundled seed data
    python scripts/setup.py --skip-index        # skip Qdrant indexing (ingest only)
    python scripts/setup.py --rebuild           # drop DB + Qdrant and start fresh
    python scripts/setup.py --stats-only        # print stats on existing DB
"""

from __future__ import annotations

import argparse
import asyncio
import sys
import time
from pathlib import Path

# Make root importable
sys.path.insert(0, str(Path(__file__).parent.parent))

from rich.console import Console
from rich.panel import Panel
from rich.rule import Rule
from rich import box

console = Console()


def banner() -> None:
    console.print(
        Panel.fit(
            "[bold cyan]PharmaAI[/bold cyan] — Data Setup Pipeline\n"
            "[dim]Medicine Alternative Recommendation System[/dim]",
            box=box.DOUBLE_EDGE,
            border_style="cyan",
            padding=(1, 4),
        )
    )


def step(n: int, total: int, label: str) -> None:
    console.print(f"\n[bold cyan][{n}/{total}][/bold cyan] {label}")


def run(args: argparse.Namespace) -> None:
    start = time.perf_counter()
    banner()

    from core.config import settings
    from core.database import init_db, get_connection
    from data.pipeline.download import DataDownloader
    from data.pipeline.preprocess import Preprocessor
    from data.pipeline.validate import Validator
    from data.pipeline.stats import DataStats

    TOTAL_STEPS = 4 if args.skip_index else 5

    # ── Stats-only mode ────────────────────────────────────────────────────────
    if args.stats_only:
        import pandas as pd
        init_db()
        with get_connection() as conn:
            df = pd.read_sql("SELECT * FROM medicines", conn)
        if df.empty:
            console.print("[red]No medicines in DB. Run setup first.[/red]")
            return
        DataStats.summary(df, "Current Database")
        DataStats.sample(df)
        return

    # ── Rebuild: drop existing data ────────────────────────────────────────────
    if args.rebuild:
        console.print(Rule("[yellow]Rebuild mode — dropping existing data[/yellow]"))
        db_path = Path(settings.DB_PATH)
        if db_path.exists():
            db_path.unlink()
            console.print(f"  [dim]Removed {db_path}[/dim]")
        qdrant_path = Path(settings.QDRANT_PATH)
        if qdrant_path.exists():
            import shutil
            shutil.rmtree(qdrant_path)
            console.print(f"  [dim]Removed {qdrant_path}[/dim]")

    # ── Step 1: Download ───────────────────────────────────────────────────────
    step(1, TOTAL_STEPS, "Downloading / loading raw data")
    downloader = DataDownloader()

    if args.source == "kaggle":
        from data.pipeline.download import RAW_DIR, KAGGLE_FILENAME
        raw_path = downloader._download_kaggle(RAW_DIR / KAGGLE_FILENAME)
    elif args.source == "seed":
        from data.pipeline.seed_generator import generate_csv
        raw_path = Path(settings.DB_PATH).parent / "raw" / "seed_medicines.csv"
        generate_csv(raw_path)
        console.print(f"[green]✓ Seed CSV generated: {raw_path.name}[/green]")
    else:
        raw_path = downloader.get(force=args.rebuild)

    # ── Step 2: Preprocess ────────────────────────────────────────────────────
    step(2, TOTAL_STEPS, "Preprocessing — parse compositions, infer categories, clean text")
    preprocessor = Preprocessor()
    clean_path = preprocessor.run(raw_path)

    # ── Step 3: Validate ──────────────────────────────────────────────────────
    step(3, TOTAL_STEPS, "Validating data quality")
    import pandas as pd

    df = pd.read_csv(clean_path)
    rejected_path = clean_path.parent / "rejected_rows.csv"
    result = Validator().run(df, rejected_path=rejected_path)

    DataStats.summary(result.valid, "Validated Dataset")

    # ── Step 4: Ingest to SQLite ──────────────────────────────────────────────
    step(4, TOTAL_STEPS, "Loading medicines into SQLite")
    init_db()

    with get_connection() as conn:
        conn.execute("DELETE FROM medicines")
        conn.execute("DELETE FROM sqlite_sequence WHERE name = 'medicines'")
        ingest_df = result.valid.copy()
        text_columns = [
            "brand_name", "composition", "salt_name", "manufacturer",
            "strength", "category", "description", "uses",
            "side_effects", "image_url",
        ]
        ingest_df[text_columns] = ingest_df[text_columns].fillna("")
        rows = ingest_df.to_dict(orient="records")
        conn.executemany(
            """
            INSERT INTO medicines
                (brand_name, composition, salt_name, manufacturer, strength, category,
                 description, uses, side_effects, image_url,
                 excellent_review_pct, average_review_pct, poor_review_pct)
            VALUES
                (:brand_name, :composition, :salt_name, :manufacturer, :strength, :category,
                 :description, :uses, :side_effects, :image_url,
                 :excellent_review_pct, :average_review_pct, :poor_review_pct)
            """,
            rows,
        )
        conn.commit()
        count = conn.execute("SELECT COUNT(*) FROM medicines").fetchone()[0]

    console.print(f"[green]✓ {count:,} medicines inserted into SQLite[/green]")

    # ── Step 5: Build vector index ────────────────────────────────────────────
    if not args.skip_index:
        step(5, TOTAL_STEPS, f"Building Qdrant vector index  [{settings.QDRANT_MODE} mode]")
        asyncio.run(_build_index())

    elapsed = time.perf_counter() - start
    result_label = "loaded" if args.skip_index else "indexed"
    console.print(
        Panel.fit(
            f"[bold green]Setup complete![/bold green]  "
            f"{count:,} medicines {result_label}  •  {elapsed:.1f}s",
            border_style="green",
        )
    )


async def _build_index() -> None:
    from services.vector_store import VectorStoreService

    await VectorStoreService.connect()

    if VectorStoreService.is_ready():
        from core.config import settings as s
        if s.QDRANT_MODE == "local":
            console.print("[dim]Existing local index found — use --rebuild to force[/dim]")
            return

    await VectorStoreService.build_from_db()


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="PharmaAI data setup pipeline")
    p.add_argument(
        "--source",
        choices=["auto", "kaggle", "seed"],
        default="auto",
        help="Data source (default: auto — Kaggle if creds available, else seed)",
    )
    p.add_argument("--skip-index",  action="store_true", help="Skip Qdrant indexing")
    p.add_argument("--rebuild",     action="store_true", help="Drop and rebuild everything")
    p.add_argument("--stats-only",  action="store_true", help="Print DB stats and exit")
    return p.parse_args()


if __name__ == "__main__":
    run(parse_args())