app/schedule_store.py

"""Persist day plans, block feedback, and learned schedule priors.

JSON files under DATA_ROOT/schedule/; feedback is append-only JSONL.
"""

from __future__ import annotations

import json
from datetime import date, datetime, timezone
from typing import Any, Literal
from uuid import uuid4

from pydantic import BaseModel, ConfigDict, Field, field_validator

from app.fsutil import append_jsonl, atomic_write_text, file_lock, read_json, read_jsonl
from app.paths import Paths
from app.schedule_math import (
    PRIOR_DEFAULT_MIN,
    is_strong_feedback,
    minutes_between,
    plan_health,
    priors_markdown,
    recompute_priors,
    validate_blocks,
)
from app.schedule_templates import KIND_LABELS, list_templates

TaskKind = Literal[
    "earn_ship",
    "admin_spain",
    "body_care",
    "move_out",
    "boundary",
    "food_out",
    "stabilize",
    "explore",
    "restore_fun",
    "sleep_window",
    "other",
]
Priority = Literal["P0", "P1", "P2"]
Intent = Literal["duty", "explore", "restore_fun", "measure"]
BlockStatus = Literal["planned", "done", "partial", "skipped", "moved", "cancelled"]
Did = Literal["done", "partial", "skipped"]
WouldRepeat = Literal["yes", "no", "maybe"]
SkipReason = Literal["time", "fear", "fse", "locks", "boring", "urge", "other"]
PlanSource = Literal["cursor", "openrouter", "user", "rules", "chat"]


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


class DayReview(BaseModel):
    """End-of-day review for chat export / import."""

    model_config = ConfigDict(extra="ignore")

    comment: str = ""
    emotions: list[str] = Field(default_factory=list)
    fse_events: str = ""
    what_moved: str = ""
    what_avoided: str = ""
    tomorrow_change: str = ""


class ScheduledBlock(BaseModel):
    """One timed block on a day plan."""

    model_config = ConfigDict(extra="ignore")

    id: str = Field(default_factory=lambda: str(uuid4()))
    date: str
    start: str
    end: str
    title: str = Field(min_length=1, max_length=200)
    kind: TaskKind = "other"
    intent: Intent = "duty"
    priority: Priority = "P2"
    planned_min: int = Field(ge=1, le=24 * 60)
    status: BlockStatus = "planned"
    source: PlanSource | str = "user"
    locked: bool = False
    notes: str = ""
    version_added: int = 1

    @field_validator("start", "end")
    @classmethod
    def _hhmm(cls, value: str) -> str:
        parts = value.split(":")
        if len(parts) != 2:
            raise ValueError("time must be HH:MM")
        h, m = int(parts[0]), int(parts[1])
        if not (0 <= h <= 23 and 0 <= m <= 59):
            raise ValueError("invalid time")
        return f"{h:02d}:{m:02d}"


class BlockFeedback(BaseModel):
    """Thick feedback for one block."""

    model_config = ConfigDict(extra="ignore")

    block_id: str
    date: str
    did: Did
    actual_min: int | None = Field(default=None, ge=0, le=24 * 60)
    quality: int | None = Field(default=None, ge=1, le=5)
    fun: int | None = Field(default=None, ge=1, le=5)
    energy_after: int | None = Field(default=None, ge=-2, le=2)
    money_amount: float | None = None
    money_currency: str = "TZS"
    would_repeat: WouldRepeat | None = None
    skip_reason: SkipReason | None = None
    note: str = ""
    emotions: list[str] = Field(default_factory=list)
    fse_event: str = ""
    intensity: int | None = Field(default=None, ge=1, le=10)
    strong: bool = False
    ts: datetime = Field(default_factory=utc_now)

    def with_strong_flag(self) -> "BlockFeedback":
        payload = self.model_dump()
        strong = is_strong_feedback(payload)
        return self.model_copy(update={"strong": strong})


class DayPlan(BaseModel):
    """Versioned day schedule."""

    model_config = ConfigDict(extra="ignore")

    date: str
    version: int = 1
    source: PlanSource | str = "user"
    blocks: list[ScheduledBlock] = Field(default_factory=list)
    capacity_hint: float = Field(default=1.0, ge=0.0, le=1.0)
    notes: str = ""
    title: str = ""
    intention: str = ""
    constraints: list[str] = Field(default_factory=list)
    day_review: DayReview = Field(default_factory=DayReview)
    warnings: list[str] = Field(default_factory=list)
    updated_at: datetime = Field(default_factory=utc_now)
    parent_version: int | None = None


class BlockCreate(BaseModel):
    """Payload for adding one block."""

    model_config = ConfigDict(extra="forbid")

    start: str
    end: str
    title: str = Field(min_length=1, max_length=200)
    kind: TaskKind = "other"
    intent: Intent = "duty"
    priority: Priority = "P2"
    planned_min: int | None = None
    locked: bool | None = None
    notes: str = ""


class BlockPatch(BaseModel):
    """Partial block update."""

    model_config = ConfigDict(extra="forbid")

    start: str | None = None
    end: str | None = None
    title: str | None = Field(default=None, min_length=1, max_length=200)
    kind: TaskKind | None = None
    intent: Intent | None = None
    priority: Priority | None = None
    planned_min: int | None = Field(default=None, ge=1, le=24 * 60)
    status: BlockStatus | None = None
    locked: bool | None = None
    notes: str | None = None


class DayPlanPut(BaseModel):
    """Replace/create a day plan."""

    model_config = ConfigDict(extra="ignore")

    blocks: list[ScheduledBlock]
    source: PlanSource | str = "user"
    capacity_hint: float | None = Field(default=None, ge=0.0, le=1.0)
    notes: str = ""
    title: str | None = None
    intention: str | None = None
    constraints: list[str] | None = None
    day_review: DayReview | None = None
    force_p0_move: bool = False


class PlanMetaPatch(BaseModel):
    """Day-level chat meta without touching blocks."""

    model_config = ConfigDict(extra="forbid")

    title: str | None = None
    intention: str | None = None
    constraints: list[str] | None = None
    day_review: DayReview | None = None


class BlockCheckBody(BaseModel):
    """Fast checkbox status without full feedback."""

    model_config = ConfigDict(extra="forbid")

    status: Literal["done", "partial", "skipped", "planned"]
    skip_reason: SkipReason | None = None


class ScheduleStore:
    """Read and mutate schedule plans and priors."""

    def __init__(self, paths: Paths, *, shrink_k: float = 3.0, max_blocks: int = 7) -> None:
        self.paths = paths
        self.shrink_k = shrink_k
        self.max_blocks = max_blocks
        self.paths.schedule_dir.mkdir(parents=True, exist_ok=True)

    def empty_plan(self, day: str) -> DayPlan:
        return DayPlan(date=day, blocks=[], source="user")

    def get_plan(self, day: str) -> DayPlan:
        path = self.paths.plan_path(day)
        raw = read_json(path)
        if raw is None:
            return self.empty_plan(day)
        return DayPlan.model_validate(raw)

    def save_plan(
        self,
        day: str,
        put: DayPlanPut,
        *,
        bump_version: bool = True,
    ) -> DayPlan:
        current = self.get_plan(day)
        blocks = []
        for block in put.blocks:
            data = block.model_dump()
            data["date"] = day
            if not data.get("planned_min"):
                data["planned_min"] = max(1, minutes_between(data["start"], data["end"]))
            if data.get("priority") == "P0" and put.source in (
                "openrouter",
                "rules",
                "cursor",
                "chat",
            ):
                data["locked"] = True if data.get("locked") is None else data["locked"]
            blocks.append(ScheduledBlock.model_validate(data))

        errors, warnings = validate_blocks(
            [b.model_dump() for b in blocks],
            max_blocks=self.max_blocks,
            previous_p0=[b.model_dump() for b in current.blocks],
            allow_p0_move=put.force_p0_move or put.source == "user",
            must_include_explore_or_restore=True,
            capacity_hint=put.capacity_hint
            if put.capacity_hint is not None
            else current.capacity_hint,
            hard_explore=False,
        )
        if errors:
            raise ValueError("; ".join(errors))

        version = current.version + 1 if bump_version and current.blocks else max(1, current.version)
        if not current.blocks and not bump_version:
            version = 1
        day_review = (
            put.day_review
            if put.day_review is not None
            else current.day_review
        )
        plan = DayPlan(
            date=day,
            version=version,
            source=put.source,
            blocks=blocks,
            capacity_hint=put.capacity_hint if put.capacity_hint is not None else current.capacity_hint,
            notes=put.notes if put.notes is not None else current.notes,
            title=put.title if put.title is not None else current.title,
            intention=put.intention if put.intention is not None else current.intention,
            constraints=(
                put.constraints if put.constraints is not None else current.constraints
            ),
            day_review=day_review,
            warnings=warnings,
            updated_at=utc_now(),
            parent_version=current.version if current.blocks else None,
        )
        self._write_plan(plan)
        return plan

    def patch_meta(self, day: str, patch: PlanMetaPatch) -> DayPlan:
        plan = self.get_plan(day)
        data = plan.model_dump()
        if patch.title is not None:
            data["title"] = patch.title
        if patch.intention is not None:
            data["intention"] = patch.intention
        if patch.constraints is not None:
            data["constraints"] = patch.constraints
        if patch.day_review is not None:
            data["day_review"] = patch.day_review.model_dump()
        data["updated_at"] = utc_now().isoformat()
        updated = DayPlan.model_validate(data)
        self._write_plan(updated)
        return updated

    def put_preserving_meta(
        self,
        day: str,
        blocks: list[ScheduledBlock],
        *,
        source: str = "user",
        force_p0_move: bool = True,
        bump_version: bool = False,
        notes: str | None = None,
    ) -> DayPlan:
        plan = self.get_plan(day)
        return self.save_plan(
            day,
            DayPlanPut(
                blocks=blocks,
                source=source,
                capacity_hint=plan.capacity_hint,
                notes=notes if notes is not None else plan.notes,
                title=plan.title,
                intention=plan.intention,
                constraints=plan.constraints,
                day_review=plan.day_review,
                force_p0_move=force_p0_move,
            ),
            bump_version=bump_version,
        )

    def _write_plan(self, plan: DayPlan) -> None:
        path = self.paths.plan_path(plan.date)
        payload = json.dumps(plan.model_dump(mode="json"), ensure_ascii=False, indent=2)
        with file_lock(path):
            atomic_write_text(path, payload + "\n")

    def add_block(self, day: str, body: BlockCreate) -> DayPlan:
        plan = self.get_plan(day)
        planned = body.planned_min or max(1, minutes_between(body.start, body.end))
        locked = body.locked if body.locked is not None else body.priority == "P0"
        block = ScheduledBlock(
            date=day,
            start=body.start,
            end=body.end,
            title=body.title,
            kind=body.kind,
            intent=body.intent,
            priority=body.priority,
            planned_min=planned,
            locked=locked,
            notes=body.notes,
            source="user",
            version_added=plan.version,
        )
        blocks = list(plan.blocks) + [block]
        return self.put_preserving_meta(day, blocks, bump_version=False)

    def patch_block(self, day: str, block_id: str, patch: BlockPatch) -> DayPlan:
        plan = self.get_plan(day)
        found = False
        blocks: list[ScheduledBlock] = []
        for block in plan.blocks:
            if block.id != block_id:
                blocks.append(block)
                continue
            found = True
            data = block.model_dump()
            data.update(patch.model_dump(exclude_unset=True))
            if patch.start is not None or patch.end is not None:
                data["planned_min"] = patch.planned_min or max(
                    1, minutes_between(data["start"], data["end"])
                )
            blocks.append(ScheduledBlock.model_validate(data))
        if not found:
            raise KeyError(block_id)
        return self.put_preserving_meta(day, blocks, bump_version=False)

    def delete_block(self, day: str, block_id: str) -> DayPlan:
        plan = self.get_plan(day)
        blocks = [b for b in plan.blocks if b.id != block_id]
        if len(blocks) == len(plan.blocks):
            raise KeyError(block_id)
        return self.put_preserving_meta(day, blocks, bump_version=False)

    def list_feedback(self, day: str | None = None) -> list[dict[str, Any]]:
        rows = read_jsonl(self.paths.schedule_feedback)
        if day is None:
            return rows
        return [r for r in rows if r.get("date") == day]

    def feedback_for_plan(self, day: str) -> list[dict[str, Any]]:
        return self.list_feedback(day)

    def submit_feedback(self, day: str, fb: BlockFeedback) -> tuple[BlockFeedback, DayPlan]:
        plan = self.get_plan(day)
        block = next((b for b in plan.blocks if b.id == fb.block_id), None)
        if block is None:
            raise KeyError(fb.block_id)
        stored = fb.model_copy(update={"date": day}).with_strong_flag()
        append_jsonl(self.paths.schedule_feedback, stored.model_dump(mode="json"))

        status_map = {"done": "done", "partial": "partial", "skipped": "skipped"}
        patched = self.patch_block(
            day,
            fb.block_id,
            BlockPatch(status=status_map[fb.did]),  # type: ignore[arg-type]
        )
        self.recompute_and_save_priors()
        return stored, patched

    def load_priors(self) -> dict[str, dict[str, Any]]:
        raw = read_json(self.paths.schedule_priors)
        if raw and isinstance(raw.get("kinds"), dict):
            return raw["kinds"]
        # Seed defaults
        kinds = {
            k: {
                "kind": k,
                "n": 0,
                "mean_actual": None,
                "d_hat": float(v),
                "mean_quality": None,
                "mean_fun": None,
                "mean_energy": None,
                "p_done": 0.5,
                "mean_slip": None,
                "repeat_score": 0.0,
            }
            for k, v in PRIOR_DEFAULT_MIN.items()
        }
        return kinds

    def recompute_and_save_priors(self) -> dict[str, dict[str, Any]]:
        feedback = self.list_feedback()
        blocks_by_id: dict[str, dict[str, Any]] = {}
        # Load blocks from feedback dates
        dates = {str(f.get("date")) for f in feedback if f.get("date")}
        for day in dates:
            for block in self.get_plan(day).blocks:
                blocks_by_id[block.id] = block.model_dump()
        kinds = recompute_priors(feedback, blocks_by_id, shrink_k=self.shrink_k)
        payload = {
            "updated_at": utc_now().isoformat(),
            "kinds": kinds,
            "kind_labels": KIND_LABELS,
        }
        text = json.dumps(payload, ensure_ascii=False, indent=2) + "\n"
        with file_lock(self.paths.schedule_priors):
            atomic_write_text(self.paths.schedule_priors, text)
        return kinds

    def get_meta(self) -> dict[str, Any]:
        return read_json(self.paths.schedule_meta) or {}

    def set_meta(self, **kwargs: Any) -> dict[str, Any]:
        meta = self.get_meta()
        meta.update(kwargs)
        text = json.dumps(meta, ensure_ascii=False, indent=2) + "\n"
        with file_lock(self.paths.schedule_meta):
            atomic_write_text(self.paths.schedule_meta, text)
        return meta

    def health(self, day: str) -> dict[str, Any]:
        plan = self.get_plan(day)
        return plan_health(
            [b.model_dump() for b in plan.blocks],
            self.feedback_for_plan(day),
        )

    def plan_with_feedback(self, day: str) -> dict[str, Any]:
        plan = self.get_plan(day)
        feedback = self.feedback_for_plan(day)
        by_block = {str(f.get("block_id")): f for f in feedback}
        blocks = []
        for block in plan.blocks:
            item = block.model_dump(mode="json")
            item["feedback"] = by_block.get(block.id)
            item["label"] = KIND_LABELS.get(block.kind, block.kind)
            blocks.append(item)
        data = plan.model_dump(mode="json")
        data["blocks"] = blocks
        data["health"] = self.health(day)
        return data

    def templates(self) -> list[dict[str, Any]]:
        return list_templates()

    def priors_table(self) -> str:
        return priors_markdown(self.load_priors())