load_to_postgres.py

"""Load plan and enriched content JSON files into PostgreSQL tables."""
from __future__ import annotations

import argparse
import json
import logging
import os
from pathlib import Path
from typing import Dict, List

import psycopg

LOGGER = logging.getLogger(__name__)

SCHEMA_SQL = """
CREATE EXTENSION IF NOT EXISTS pgcrypto;

CREATE TABLE IF NOT EXISTS learning_days (
    day INT PRIMARY KEY,
    title TEXT NOT NULL,
    goal TEXT,
    estimated_minutes INT,
    topics JSONB NOT NULL
);

CREATE TABLE IF NOT EXISTS day_content (
    day INT PRIMARY KEY REFERENCES learning_days(day) ON DELETE CASCADE,
    overview TEXT,
    key_points JSONB,
    flashcards JSONB,
    practice JSONB,
    reflection_prompt TEXT
);

CREATE TABLE IF NOT EXISTS books (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    slug TEXT NOT NULL UNIQUE,
    title TEXT NOT NULL,
    description TEXT,
    cover_url TEXT,
    default_days INT,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS book_topics (
    id BIGSERIAL PRIMARY KEY,
    book_id UUID NOT NULL REFERENCES books(id) ON DELETE CASCADE,
    chapter_index INT NOT NULL,
    topic_index INT NOT NULL,
    payload JSONB NOT NULL,
    UNIQUE (book_id, chapter_index, topic_index)
);

CREATE TABLE IF NOT EXISTS book_content (
    id BIGSERIAL PRIMARY KEY,
    book_id UUID NOT NULL REFERENCES books(id) ON DELETE CASCADE,
    day INT NOT NULL,
    payload JSONB NOT NULL,
    UNIQUE (book_id, day)
);

CREATE TABLE IF NOT EXISTS users (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    email TEXT NOT NULL UNIQUE,
    display_name TEXT,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS study_plans (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    book_id UUID NOT NULL REFERENCES books(id) ON DELETE CASCADE,
    total_days INT NOT NULL,
    minutes_per_day INT,
    focus TEXT,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    is_template BOOLEAN NOT NULL DEFAULT FALSE,
    template_key TEXT,
    template_parent_id UUID REFERENCES study_plans(id)
);

CREATE TABLE IF NOT EXISTS plan_days (
    id BIGSERIAL PRIMARY KEY,
    plan_id UUID NOT NULL REFERENCES study_plans(id) ON DELETE CASCADE,
    day_number INT NOT NULL,
    payload JSONB NOT NULL,
    UNIQUE (plan_id, day_number)
);

CREATE TABLE IF NOT EXISTS plan_day_content (
    id BIGSERIAL PRIMARY KEY,
    plan_day_id BIGINT NOT NULL UNIQUE REFERENCES plan_days(id) ON DELETE CASCADE,
    content JSONB NOT NULL
);

CREATE UNIQUE INDEX IF NOT EXISTS idx_plan_day_content_plan_day_id
    ON plan_day_content(plan_day_id);
ALTER TABLE study_plans
    ADD COLUMN IF NOT EXISTS is_template BOOLEAN NOT NULL DEFAULT FALSE;
ALTER TABLE study_plans
    ADD COLUMN IF NOT EXISTS template_key TEXT;
ALTER TABLE study_plans
    ADD COLUMN IF NOT EXISTS template_parent_id UUID REFERENCES study_plans(id);
CREATE UNIQUE INDEX IF NOT EXISTS idx_study_plans_template_key_unique
    ON study_plans(template_key)
    WHERE is_template AND template_key IS NOT NULL;
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Persist plan and daily content into PostgreSQL")
    parser.add_argument("--plan", default="artifacts/plan.json", help="Path to plan JSON file")
    parser.add_argument(
        "--content",
        default="artifacts/daily_content.json",
        help="Path to enriched daily content JSON file",
    )
    parser.add_argument(
        "--syllabus",
        default="artifacts/syllabus.json",
        help="Path to syllabus JSON file for book topic seeding",
    )
    parser.add_argument(
        "--database-url",
        default=os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/learning"),
        help="PostgreSQL connection string",
    )
    parser.add_argument("--book-slug", help="Slug identifier for the book (enables book seeding)")
    parser.add_argument("--book-title", help="Override title for the book row")
    parser.add_argument("--book-description", help="Optional marketing description")
    parser.add_argument("--book-cover-url", help="Optional cover image URL")
    parser.add_argument("--default-days", type=int, help="Default days suggested for the book")
    return parser.parse_args()


def configure_logging() -> None:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")


def load_json(path: Path) -> Dict:
    if not path.exists():
        raise FileNotFoundError(f"File not found: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def ensure_schema(conn: psycopg.Connection) -> None:
    with conn.cursor() as cur:
        cur.execute(SCHEMA_SQL)
    conn.commit()


def upsert_book(cur: psycopg.Cursor, args: argparse.Namespace, plan: Dict) -> str | None:
    if not args.book_slug:
        return None
    title = args.book_title or plan.get("book_title") or args.book_slug.replace("-", " ").title()
    description = args.book_description
    cover_url = args.book_cover_url
    default_days = args.default_days or len(plan.get("days", [])) or None
    cur.execute(
        """
        INSERT INTO books (slug, title, description, cover_url, default_days)
        VALUES (%s, %s, %s, %s, %s)
        ON CONFLICT (slug)
        DO UPDATE SET
            title = EXCLUDED.title,
            description = EXCLUDED.description,
            cover_url = EXCLUDED.cover_url,
            default_days = EXCLUDED.default_days
        RETURNING id;
        """,
        (args.book_slug, title, description, cover_url, default_days),
    )
    book_id = cur.fetchone()[0]
    return book_id


def seed_book_topics(cur: psycopg.Cursor, book_id: str, syllabus_path: Path) -> None:
    if not book_id or not syllabus_path.exists():
        return
    syllabus = load_json(syllabus_path)
    chapters = syllabus.get("chapters", [])
    for chapter_idx, chapter in enumerate(chapters, start=1):
        topics = chapter.get("topics", [])
        for topic_idx, topic in enumerate(topics, start=1):
            payload = json.dumps(
                {
                    "title": chapter.get("title"),
                    "topic": topic,
                },
                ensure_ascii=False,
            )
            cur.execute(
                """
                INSERT INTO book_topics (book_id, chapter_index, topic_index, payload)
                VALUES (%s, %s, %s, %s::jsonb)
                ON CONFLICT (book_id, chapter_index, topic_index)
                DO UPDATE SET payload = EXCLUDED.payload;
                """,
                (book_id, chapter_idx, topic_idx, payload),
            )


def seed_book_content(cur: psycopg.Cursor, book_id: str, plan: Dict, enriched_index: Dict[int, Dict]) -> None:
    if not book_id:
        return
    for day in plan.get("days", []):
        day_number = day.get("day")
        payload = json.dumps(
            {
                "plan": day,
                "content": enriched_index.get(day_number),
            },
            ensure_ascii=False,
        )
        cur.execute(
            """
            INSERT INTO book_content (book_id, day, payload)
            VALUES (%s, %s, %s::jsonb)
            ON CONFLICT (book_id, day)
            DO UPDATE SET payload = EXCLUDED.payload;
            """,
            (book_id, day_number, payload),
        )


def upsert_learning_day(cur: psycopg.Cursor, day: Dict) -> None:
    cur.execute(
        """
        INSERT INTO learning_days (day, title, goal, estimated_minutes, topics)
        VALUES (%s, %s, %s, %s, %s::jsonb)
        ON CONFLICT (day)
        DO UPDATE SET
            title = EXCLUDED.title,
            goal = EXCLUDED.goal,
            estimated_minutes = EXCLUDED.estimated_minutes,
            topics = EXCLUDED.topics;
        """,
        (
            day["day"],
            day.get("title"),
            day.get("goal"),
            day.get("estimated_minutes"),
            json.dumps(day.get("topics", []), ensure_ascii=False),
        ),
    )


def upsert_day_content(cur: psycopg.Cursor, day: int, content: Dict | None) -> None:
    payload = content or {}
    cur.execute(
        """
        INSERT INTO day_content (day, overview, key_points, flashcards, practice, reflection_prompt)
        VALUES (%s, %s, %s::jsonb, %s::jsonb, %s::jsonb, %s)
        ON CONFLICT (day)
        DO UPDATE SET
            overview = EXCLUDED.overview,
            key_points = EXCLUDED.key_points,
            flashcards = EXCLUDED.flashcards,
            practice = EXCLUDED.practice,
            reflection_prompt = EXCLUDED.reflection_prompt;
        """,
        (
            day,
            payload.get("overview"),
            json.dumps(payload.get("key_points", []), ensure_ascii=False),
            json.dumps(payload.get("flashcards", []), ensure_ascii=False),
            json.dumps(payload.get("practice", {}), ensure_ascii=False),
            payload.get("reflection_prompt"),
        ),
    )


def main() -> None:
    configure_logging()
    args = parse_args()

    plan = load_json(Path(args.plan))
    enriched = load_json(Path(args.content)) if Path(args.content).exists() else {"days": []}
    enriched_index = {entry.get("day"): entry for entry in enriched.get("days", [])}

    LOGGER.info("Connecting to %s", args.database_url)
    with psycopg.connect(args.database_url) as conn:
        ensure_schema(conn)
        with conn.cursor() as cur:
            book_id = upsert_book(cur, args, plan)
            seed_book_topics(cur, book_id, Path(args.syllabus))
            seed_book_content(cur, book_id, plan, enriched_index)
            for day in plan.get("days", []):
                upsert_learning_day(cur, day)
                upsert_day_content(cur, day.get("day"), enriched_index.get(day.get("day")))
        conn.commit()
    LOGGER.info("Load complete: %s days persisted", len(plan.get("days", [])))


if __name__ == "__main__":
    main()