writer-work-flow/packages/db/ww_db/models.py

"""ORM 模型——全部 MVP 表（ARCH §3.1）。

snake_case；无向量列；P2 表 timeline/decisions 不建（见 DEV_PLAN T0.2）。
"""

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Any

from sqlalchemy import (
    BigInteger,
    CheckConstraint,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    LargeBinary,
    Text,
    UniqueConstraint,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column

from ww_db.base import Base, CreatedAt, TimestampedMixin, UuidPk

# ---- 运营/系统表 ----


class User(UuidPk, CreatedAt, Base):
    __tablename__ = "users"
    email: Mapped[str] = mapped_column(Text, unique=True, nullable=False)
    display_name: Mapped[str | None] = mapped_column(Text)


# ---- 创作数据表 ----


class Project(UuidPk, TimestampedMixin, Base):
    __tablename__ = "projects"
    owner_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id"), nullable=False, index=True)
    title: Mapped[str] = mapped_column(Text, nullable=False)
    genre: Mapped[str | None] = mapped_column(Text)
    logline: Mapped[str | None] = mapped_column(Text)
    premise: Mapped[str | None] = mapped_column(Text)
    theme: Mapped[str | None] = mapped_column(Text)
    selling_points: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    structure: Mapped[str | None] = mapped_column(Text)


class Character(UuidPk, CreatedAt, Base):
    __tablename__ = "characters"
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    name: Mapped[str] = mapped_column(Text, nullable=False)
    role: Mapped[str | None] = mapped_column(Text)
    traits: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    appearance: Mapped[str | None] = mapped_column(Text)
    motive: Mapped[str | None] = mapped_column(Text)
    backstory: Mapped[str | None] = mapped_column(Text)
    arc: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    speech_tics: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    tags: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    relations: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    first_chapter: Mapped[int | None] = mapped_column(Integer)
    latest_state: Mapped[str | None] = mapped_column(Text)


class WorldEntity(UuidPk, CreatedAt, Base):
    __tablename__ = "world_entities"
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    type: Mapped[str] = mapped_column(Text, nullable=False)
    name: Mapped[str] = mapped_column(Text, nullable=False)
    rules: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    first_chapter: Mapped[int | None] = mapped_column(Integer)
    latest_state: Mapped[str | None] = mapped_column(Text)


class Outline(UuidPk, Base):
    __tablename__ = "outline"
    __table_args__ = (UniqueConstraint("project_id", "chapter_no"),)
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    volume: Mapped[int] = mapped_column(Integer, nullable=False)
    chapter_no: Mapped[int] = mapped_column(Integer, nullable=False)
    beats: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    foreshadow_windows: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))


class ChapterDigest(UuidPk, CreatedAt, Base):
    __tablename__ = "chapter_digests"
    __table_args__ = (
        Index(
            "ix_chapter_digests_project_chapter",
            "project_id",
            text("chapter_no DESC"),
        ),
    )
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    chapter_no: Mapped[int] = mapped_column(Integer, nullable=False)
    facts: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)


class Foreshadow(UuidPk, Base):
    __tablename__ = "foreshadow"
    __table_args__ = (
        UniqueConstraint("project_id", "code"),
        CheckConstraint(
            "status IN ('OPEN', 'PARTIAL', 'CLOSED', 'OVERDUE')",
            name="ck_foreshadow_status",
        ),
    )
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    code: Mapped[str] = mapped_column(Text, nullable=False)
    title: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("'OPEN'"))
    planted_at: Mapped[int | None] = mapped_column(Integer)
    content: Mapped[str | None] = mapped_column(Text)
    expected_close_from: Mapped[int | None] = mapped_column(Integer)
    expected_close_to: Mapped[int | None] = mapped_column(Integer)
    importance: Mapped[str | None] = mapped_column(Text)
    links: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    progress: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=text("now()"),
        onupdate=text("now()"),
    )


class StyleFingerprint(UuidPk, CreatedAt, Base):
    __tablename__ = "style_fingerprint"
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    dimensions_json: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
    evidence_json: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
    version: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("1"))


class Rule(UuidPk, CreatedAt, Base):
    __tablename__ = "rules"
    __table_args__ = (
        Index(
            "ix_rules_global",
            "id",
            postgresql_where=text("project_id IS NULL"),
        ),
    )
    project_id: Mapped[uuid.UUID | None] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), index=True
    )
    level: Mapped[str] = mapped_column(Text, nullable=False)
    content: Mapped[str] = mapped_column(Text, nullable=False)


class Chapter(UuidPk, CreatedAt, Base):
    __tablename__ = "chapters"
    __table_args__ = (
        UniqueConstraint("project_id", "chapter_no", "version"),
        CheckConstraint(
            "status IN ('draft', 'accepted')",
            name="ck_chapters_status",
        ),
    )
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    volume: Mapped[int] = mapped_column(Integer, nullable=False)
    chapter_no: Mapped[int] = mapped_column(Integer, nullable=False)
    content: Mapped[str | None] = mapped_column(Text)
    status: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("'draft'"))
    version: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("1"))


class ChapterInjection(UuidPk, TimestampedMixin, Base):
    """本章注入覆盖（B0 可控版）：作者对确定性自动选择的微调，每章一行。

    与 `outline.beats` 分表存放（复用 outline 行会污染 beats/prompt，已否决）。
    pinned/excluded 为 JSONB list[{kind,name}]；recent_n 覆盖近况回看章数（NULL=用默认）。
    """

    __tablename__ = "chapter_injection"
    __table_args__ = (UniqueConstraint("project_id", "chapter_no"),)
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    chapter_no: Mapped[int] = mapped_column(Integer, nullable=False)
    pinned: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    excluded: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    recent_n: Mapped[int | None] = mapped_column(Integer)


class ChapterReview(UuidPk, CreatedAt, Base):
    __tablename__ = "chapter_reviews"
    __table_args__ = (
        Index(
            "ix_chapter_reviews_project_chapter",
            "project_id",
            "chapter_no",
            text("created_at DESC"),
        ),
    )
    project_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
    )
    chapter_no: Mapped[int] = mapped_column(Integer, nullable=False)
    chapter_version: Mapped[int | None] = mapped_column(Integer)
    conflicts: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    foreshadow_sug: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    style: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    pace: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    health_score: Mapped[int | None] = mapped_column(Integer)
    decisions: Mapped[dict[str, Any] | None] = mapped_column(JSONB)


# ---- 运营/系统表（续）----


class ProviderCredential(UuidPk, CreatedAt, Base):
    __tablename__ = "provider_credentials"
    __table_args__ = (UniqueConstraint("owner_id", "project_id", "provider"),)
    owner_id: Mapped[uuid.UUID] = mapped_column(
        ForeignKey("users.id", ondelete="CASCADE"), nullable=False
    )
    project_id: Mapped[uuid.UUID | None] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE")
    )
    provider: Mapped[str] = mapped_column(Text, nullable=False)
    auth_type: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("'api_key'"))
    api_key_enc: Mapped[bytes | None] = mapped_column(LargeBinary, nullable=True)
    oauth_enc: Mapped[bytes | None] = mapped_column(LargeBinary, nullable=True)


class TierRouting(UuidPk, Base):
    __tablename__ = "tier_routing"
    __table_args__ = (UniqueConstraint("project_id", "tier"),)
    project_id: Mapped[uuid.UUID | None] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE")
    )
    tier: Mapped[str] = mapped_column(Text, nullable=False)
    provider: Mapped[str] = mapped_column(Text, nullable=False)
    model: Mapped[str] = mapped_column(Text, nullable=False)
    fallback: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))


class UsageLedger(UuidPk, CreatedAt, Base):
    __tablename__ = "usage_ledger"
    owner_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id"), nullable=False, index=True)
    project_id: Mapped[uuid.UUID | None] = mapped_column(ForeignKey("projects.id"), index=True)
    provider: Mapped[str] = mapped_column(Text, nullable=False)
    model: Mapped[str] = mapped_column(Text, nullable=False)
    input_tokens: Mapped[int] = mapped_column(Integer, nullable=False)
    output_tokens: Mapped[int] = mapped_column(Integer, nullable=False)
    cache_read: Mapped[int] = mapped_column(Integer, server_default=text("0"))
    cost_minor: Mapped[int] = mapped_column(BigInteger, nullable=False)
    currency: Mapped[str] = mapped_column(Text, nullable=False)


class Skill(UuidPk, CreatedAt, Base):
    __tablename__ = "skills"
    scope: Mapped[str] = mapped_column(Text, nullable=False)
    owner_id: Mapped[uuid.UUID | None] = mapped_column(ForeignKey("users.id"))
    name: Mapped[str] = mapped_column(Text, nullable=False)
    description: Mapped[str | None] = mapped_column(Text)
    tier: Mapped[str] = mapped_column(Text, nullable=False)
    system_prompt: Mapped[str] = mapped_column(Text, nullable=False)
    input_schema: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    output_schema: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    reads: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    writes: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))
    genre: Mapped[str | None] = mapped_column(Text)
    examples: Mapped[list[Any]] = mapped_column(JSONB, server_default=text("'[]'"))


class PromptTemplate(UuidPk, CreatedAt, Base):
    """作者保存/复用的提示词模板（单用户本地版；无分享/市场）。

    owner_id 为单用户原型 stub；tool_key 可选关联某生成器（NULL=通用）。
    """

    __tablename__ = "prompt_templates"
    owner_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id"), nullable=False, index=True)
    title: Mapped[str] = mapped_column(Text, nullable=False)
    body: Mapped[str] = mapped_column(Text, nullable=False)
    category: Mapped[str | None] = mapped_column(Text)
    tool_key: Mapped[str | None] = mapped_column(Text)


class Job(UuidPk, TimestampedMixin, Base):
    __tablename__ = "jobs"
    __table_args__ = (
        Index(
            "ix_jobs_status_running",
            "status",
            postgresql_where=text("status = 'running'"),
        ),
    )
    project_id: Mapped[uuid.UUID | None] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE")
    )
    kind: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(Text, nullable=False, server_default=text("'queued'"))
    progress: Mapped[int] = mapped_column(Integer, server_default=text("0"))
    result: Mapped[dict[str, Any] | None] = mapped_column(JSONB)
    error: Mapped[str | None] = mapped_column(Text)