#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Item 111: Spaced retrieval scheduling (Memory) — FSRS-6 per-item.

Per-item scheduling persisted to SQLite (spaced_schedule table).
Upgraded from SM-2 to FSRS-6: stability-based scheduling with retrievability-aware
interval growth. Spacing effect: items recalled at low R get larger boost.

Math lives in _fsrs_math.py. In-memory fallback in _spaced_mem_fallback.py.
Old domain-level schedule retained as legacy shim for backward compat.
"""
from __future__ import annotations

import logging
import time
from datetime import datetime, timezone, timedelta
from typing import Optional

# Re-export FSRS-6 math + constants (keeps test imports from this module working)
from ._fsrs_math import (  # noqa: F401
    _DESIRED_RETENTION, _FSRS_INITIAL_STAB, _FSRS_INITIAL_DIFF,
    _FSRS_MIN_STAB, _FSRS_MAX_STAB, _FSRS_MIN_DIFF, _FSRS_MAX_DIFF,
    _GRADE_SUCCESS, _GRADE_FAILURE,
    _retrievability, _next_interval,
    _stability_recall, _stability_lapse, _difficulty_update,
)
from ._spaced_mem_fallback import (  # noqa: F401
    _MEM_SCHEDULE,
    _schedule_item_in_memory,
    _mark_recalled_in_memory,
)

log = logging.getLogger("pulse.brain.spaced_retrieval")

# ---------------------------------------------------------------------------
# Legacy domain-level schedule (retained for backward compat)
# ---------------------------------------------------------------------------
_DOMAIN_QUERY_SCHEDULE: dict = {}
_SPACING_INTERVALS_HOURS = [1, 6, 24, 72]


def _schedule_next_retrieval(domain: str) -> dict:
    """Legacy: compute next retrieval time for a domain (fixed intervals)."""
    try:
        now = datetime.now(timezone.utc)
        entry = _DOMAIN_QUERY_SCHEDULE.get(domain, {})
        prev_interval = entry.get("interval_hours", 0)
        next_interval = _SPACING_INTERVALS_HOURS[0]
        for iv in _SPACING_INTERVALS_HOURS:
            if iv > prev_interval:
                next_interval = iv
                break
        else:
            next_interval = _SPACING_INTERVALS_HOURS[-1]
        next_dt = now + timedelta(hours=next_interval)
        record = {
            "domain": domain, "last_queried": now.isoformat(),
            "interval_hours": next_interval, "next_retrieval": next_dt.isoformat(),
        }
        _DOMAIN_QUERY_SCHEDULE[domain] = record
        return record
    except Exception:
        return {"domain": domain, "interval_hours": 1, "next_retrieval": "", "last_queried": ""}


def get_due_domains() -> list:
    """Return list of domains whose next_retrieval time has passed (legacy)."""
    due = []
    try:
        now = datetime.now(timezone.utc)
        for domain, entry in _DOMAIN_QUERY_SCHEDULE.items():
            nxt = entry.get("next_retrieval", "")
            if not nxt:
                continue
            try:
                nxt_dt = datetime.fromisoformat(nxt.replace("Z", "+00:00"))
                if nxt_dt.tzinfo is None:
                    nxt_dt = nxt_dt.replace(tzinfo=timezone.utc)
                if now >= nxt_dt:
                    due.append(domain)
            except (ValueError, TypeError):
                pass
    except Exception:
        pass
    return due


# ---------------------------------------------------------------------------
# SQLite-backed per-item scheduling
# ---------------------------------------------------------------------------

def _get_db():
    try:
        from pulse.core.brain_db import get_conn
        return get_conn()
    except Exception:
        return None


def _ensure_table(conn) -> None:
    """Create spaced_schedule table and add difficulty column (FSRS-6 upgrade)."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS spaced_schedule (
            item_id          TEXT PRIMARY KEY,
            item_type        TEXT NOT NULL DEFAULT 'lesson',
            interval_hours   REAL NOT NULL DEFAULT 1.0,
            ef               REAL NOT NULL DEFAULT 2.5,
            next_review_ts   REAL NOT NULL,
            last_reviewed_ts REAL,
            review_count     INTEGER NOT NULL DEFAULT 0,
            created_ts       REAL NOT NULL
        )
    """)
    try:
        conn.execute(
            "ALTER TABLE spaced_schedule ADD COLUMN difficulty REAL NOT NULL DEFAULT 5.0"
        )
    except Exception:
        pass  # column already exists
    conn.commit()


def schedule_item(item_id: str, item_type: str = "lesson") -> dict:
    """Register a new item for spaced retrieval (FSRS-6 initial values)."""
    conn = _get_db()
    if conn is None:
        return _schedule_item_in_memory(item_id, item_type)
    try:
        _ensure_table(conn)
        now_ts = time.time()
        next_ts = now_ts + _FSRS_INITIAL_STAB * 3600
        conn.execute("""
            INSERT OR IGNORE INTO spaced_schedule
                (item_id, item_type, interval_hours, ef, next_review_ts,
                 last_reviewed_ts, review_count, created_ts, difficulty)
            VALUES (?, ?, ?, ?, ?, NULL, 0, ?, ?)
        """, (item_id, item_type, _FSRS_INITIAL_STAB, 2.5,
              next_ts, now_ts, _FSRS_INITIAL_DIFF))
        conn.commit()
        return _load_item_record(conn, item_id)
    except Exception as exc:
        log.debug("schedule_item SQLite failed: %s", exc)
        return _schedule_item_in_memory(item_id, item_type)


def mark_recalled(item_id: str, success: bool) -> dict:
    """Update schedule after retrieval attempt (FSRS-6).

    success=True  → stability grows; spacing effect (low R → larger boost)
    success=False → lapse: stability reduces based on difficulty
    """
    conn = _get_db()
    if conn is None:
        return _mark_recalled_in_memory(item_id, success)
    try:
        _ensure_table(conn)
        row = conn.execute(
            "SELECT interval_hours, ef, review_count, last_reviewed_ts, "
            "created_ts, COALESCE(difficulty, 5.0) "
            "FROM spaced_schedule WHERE item_id=?",
            (item_id,)
        ).fetchone()
        if row is None:
            schedule_item(item_id)
            row = conn.execute(
                "SELECT interval_hours, ef, review_count, last_reviewed_ts, "
                "created_ts, COALESCE(difficulty, 5.0) "
                "FROM spaced_schedule WHERE item_id=?",
                (item_id,)
            ).fetchone()
        stability  = float(row[0])
        count      = int(row[2])
        last_ts    = row[3]
        created_ts = float(row[4])
        difficulty = float(row[5])
        now_ts     = time.time()

        ref_ts = float(last_ts) if last_ts is not None else created_ts
        elapsed = max(0.0, (now_ts - ref_ts) / 3600.0)
        R = _retrievability(elapsed, stability)

        grade = _GRADE_SUCCESS if success else _GRADE_FAILURE
        new_stab = _stability_recall(stability, difficulty, R) if success else _stability_lapse(stability, difficulty)
        new_diff = _difficulty_update(difficulty, grade)
        new_ivl  = _next_interval(new_stab)

        conn.execute("""
            UPDATE spaced_schedule
            SET interval_hours=?, next_review_ts=?, last_reviewed_ts=?,
                review_count=?, difficulty=?
            WHERE item_id=?
        """, (new_stab, now_ts + new_ivl * 3600, now_ts,
              count + 1, new_diff, item_id))
        conn.commit()
        return _load_item_record(conn, item_id)
    except Exception as exc:
        log.debug("mark_recalled SQLite failed: %s", exc)
        return _mark_recalled_in_memory(item_id, success)


def get_due_items(item_type: Optional[str] = None) -> list:
    """Return item_ids due for review (next_review_ts <= now)."""
    conn = _get_db()
    if conn is None:
        return []
    try:
        _ensure_table(conn)
        now_ts = time.time()
        if item_type:
            rows = conn.execute(
                "SELECT item_id FROM spaced_schedule "
                "WHERE next_review_ts <= ? AND item_type=?",
                (now_ts, item_type)
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT item_id FROM spaced_schedule WHERE next_review_ts <= ?",
                (now_ts,)
            ).fetchall()
        return [r[0] for r in rows]
    except Exception as exc:
        log.debug("get_due_items failed: %s", exc)
        return []


def _load_item_record(conn, item_id: str) -> dict:
    row = conn.execute(
        "SELECT item_id, item_type, interval_hours, ef, next_review_ts, "
        "last_reviewed_ts, review_count, COALESCE(difficulty, 5.0) "
        "FROM spaced_schedule WHERE item_id=?",
        (item_id,)
    ).fetchone()
    if not row:
        return {}
    return {
        "item_id": row[0], "item_type": row[1],
        "interval_hours": row[2], "ef": row[3],
        "next_review_ts": row[4], "last_reviewed_ts": row[5],
        "review_count": row[6], "difficulty": row[7],
    }
