from __future__ import annotations

from datetime import datetime

from sqlalchemy import BigInteger, Index, Integer, String, Text, UniqueConstraint, func
from sqlalchemy.orm import Mapped, mapped_column

from supply_infra.db.base import Base


class DemandGradePlan(Base):
    """统筹规划 Agent 产出的单日全量分级计划。"""

    __tablename__ = "demand_grade_plan"
    __table_args__ = (Index("idx_demand_grade_plan_biz_dt", "biz_dt"),)

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    plan_id: Mapped[str] = mapped_column(String(36), nullable=False, unique=True)
    biz_dt: Mapped[str] = mapped_column(String(8), nullable=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="planned")
    total_hanging_nodes: Mapped[int] = mapped_column(Integer, nullable=False)
    group_count: Mapped[int] = mapped_column(Integer, nullable=False)
    coverage_complete: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    plan_json: Mapped[str] = mapped_column(Text, nullable=False)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now(), onupdate=func.now())


class DemandGradePlanGroup(Base):
    """日计划中的节点组任务，由并发分级 worker 领取执行。"""

    __tablename__ = "demand_grade_plan_group"
    __table_args__ = (
        UniqueConstraint("plan_id", "group_no", name="uk_demand_grade_plan_group"),
        Index("idx_demand_grade_plan_group_claim", "plan_id", "status", "group_no"),
        Index("idx_demand_grade_plan_group_biz_dt", "biz_dt"),
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    plan_id: Mapped[str] = mapped_column(String(36), nullable=False)
    biz_dt: Mapped[str] = mapped_column(String(8), nullable=False)
    group_no: Mapped[int] = mapped_column(Integer, nullable=False)
    group_key: Mapped[str] = mapped_column(String(128), nullable=False)
    category_ids: Mapped[str] = mapped_column(Text, nullable=False)
    planning_reason: Mapped[str] = mapped_column(Text, nullable=False)
    shared_traits: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="pending")
    attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    started_at: Mapped[datetime | None] = mapped_column(nullable=True)
    finished_at: Mapped[datetime | None] = mapped_column(nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now(), onupdate=func.now())


class DemandGradePlanGroupItem(Base):
    """计划组内待执行分级的具体需求池记录。"""

    __tablename__ = "demand_grade_plan_group_item"
    __table_args__ = (
        UniqueConstraint("group_id", "pool_id", name="uk_demand_grade_plan_group_item"),
        Index("idx_demand_grade_plan_group_item_group_status", "group_id", "status", "sort_order"),
        Index("idx_demand_grade_plan_group_item_biz_dt", "biz_dt"),
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    group_id: Mapped[int] = mapped_column(BigInteger, nullable=False)
    biz_dt: Mapped[str] = mapped_column(String(8), nullable=False)
    pool_id: Mapped[int] = mapped_column(BigInteger, nullable=False)
    demand_name: Mapped[str] = mapped_column(String(512), nullable=False)
    sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="pending")
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    finished_at: Mapped[datetime | None] = mapped_column(nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now(), onupdate=func.now())