from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from zoneinfo import ZoneInfo

from supply_infra.aigc.client import AigcClient, MAX_VIDEOS_PER_CRAWLER_PLAN
from supply_infra.aigc.plan_map import (
    AigcPlanPair,
    assignment_summary,
    chunk_list,
    distribute_evenly,
    list_unique_plan_pairs,
)
from supply_infra.config import get_infra_settings
from supply_infra.db.repositories.demand_grade_repo import DemandGradeRepository
from supply_infra.db.session import get_session
from supply_infra.services.video_discovery_service import (
    PublishableCandidate,
    get_video_discovery_service,
)

logger = logging.getLogger(__name__)

# 仅发布主推荐，不包含 rejected / pending_evaluation。
_PUBLISHABLE_BUCKETS = ("primary",)


@dataclass
class PublishBatchResult:
    plan_label: str
    produce_plan_id: str
    publish_plan_id: str
    aweme_ids: list[str]
    candidate_ids: list[int]
    crawler_plan_id: str | None = None
    crawler_plan_name: str | None = None
    bind_success: bool = False
    bind_error: str | None = None
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "plan_label": self.plan_label,
            "produce_plan_id": self.produce_plan_id,
            "publish_plan_id": self.publish_plan_id,
            "video_count": len(self.aweme_ids),
            "aweme_ids": self.aweme_ids,
            "candidate_ids": self.candidate_ids,
            "crawler_plan_id": self.crawler_plan_id,
            "crawler_plan_name": self.crawler_plan_name,
            "bind_success": self.bind_success,
            "bind_error": self.bind_error,
            "error": self.error,
        }


def _resolve_biz_dt(biz_dt: str | None) -> str | None:
    if biz_dt:
        text = str(biz_dt).strip()
        if len(text) == 8 and text.isdigit():
            return text
        raise ValueError(f"biz_dt 格式无效，应为 YYYYMMDD: {biz_dt!r}")

    with get_session() as session:
        latest = DemandGradeRepository(session).get_latest_biz_dt()
    return str(latest) if latest else None


def _group_candidates_by_plan(
    candidates: list[PublishableCandidate],
    plan_pairs: list[AigcPlanPair],
) -> list[tuple[AigcPlanPair, list[PublishableCandidate]]]:
    buckets = distribute_evenly(candidates, len(plan_pairs))
    grouped: list[tuple[AigcPlanPair, list[PublishableCandidate]]] = []
    for plan_pair, bucket in zip(plan_pairs, buckets, strict=False):
        if bucket:
            grouped.append((plan_pair, bucket))
    return grouped


def publish_videos_from_discovery(
    *,
    biz_dt: str | None = None,
    run_id: str | None = None,
    skip_published: bool = True,
    limit: int | None = None,
    any_biz_dt: bool = False,
) -> dict[str, Any]:
    """
    从 video_discovery_candidate 读取视频，按轮询均匀分配到各 AIGC 计划对。

    仅处理 decision_bucket 为 primary 且 aweme_id 非空的候选，不区分品类。
    any_biz_dt=True 时不按业务日过滤，用于定时轮询扫全部未发布候选。
    """
    settings = get_infra_settings()
    resolved_biz_dt = None if any_biz_dt else _resolve_biz_dt(biz_dt)
    plan_pairs = list_unique_plan_pairs()
    if not plan_pairs:
        raise RuntimeError("AIGC 计划映射为空，无法分发")

    candidates = get_video_discovery_service().list_publishable_candidates(
        run_id=run_id,
        biz_dt=resolved_biz_dt,
        skip_published=skip_published,
        limit=limit,
    )

    if not candidates:
        return {
            "success": True,
            "message": "没有待发布的候选视频",
            "biz_dt": resolved_biz_dt,
            "any_biz_dt": any_biz_dt,
            "run_id": run_id,
            "plan_count": len(plan_pairs),
            "candidate_count": 0,
            "decision_buckets": list(_PUBLISHABLE_BUCKETS),
            "batches": [],
        }

    grouped = _group_candidates_by_plan(candidates, plan_pairs)
    distribution_preview = assignment_summary(
        [plan for plan, _ in grouped],
        [bucket for _, bucket in grouped],
    )

    client = AigcClient()
    timezone = ZoneInfo(settings.scheduler_timezone)
    timestamp = datetime.now(timezone).strftime("%Y%m%d%H%M%S")
    batch_results: list[PublishBatchResult] = []

    for plan_pair, plan_candidates in grouped:
        candidate_batches = chunk_list(plan_candidates, MAX_VIDEOS_PER_CRAWLER_PLAN)
        for batch_index, candidate_batch in enumerate(candidate_batches, start=1):
            aweme_batch = [str(item.aweme_id) for item in candidate_batch]
            id_batch = [int(item.id) for item in candidate_batch]
            plan_name = (
                f"【SupplyAgent】{plan_pair.label}-均匀分发-"
                f"{timestamp}-批次{batch_index}"
            )
            create_result = client.create_video_crawler_plan(
                aweme_batch,
                plan_name=plan_name,
            )
            batch = PublishBatchResult(
                plan_label=plan_pair.label,
                produce_plan_id=plan_pair.produce_plan_id,
                publish_plan_id=plan_pair.publish_plan_id,
                aweme_ids=aweme_batch,
                candidate_ids=id_batch,
            )

            if not create_result.get("success"):
                batch.error = str(create_result.get("error") or "创建爬取计划失败")
                batch_results.append(batch)
                continue

            crawler_plan_id = str(create_result.get("crawler_plan_id") or "")
            crawler_plan_name = str(create_result.get("crawler_plan_name") or plan_name)
            batch.crawler_plan_id = crawler_plan_id
            batch.crawler_plan_name = crawler_plan_name

            bind_result = client.bind_crawler_to_produce_plan(
                crawler_plan_id,
                plan_pair.produce_plan_id,
                crawler_plan_name=crawler_plan_name,
            )
            batch.bind_success = bool(bind_result.get("success"))
            if not batch.bind_success:
                batch.bind_error = str(bind_result.get("error") or "绑定生成计划失败")

            if crawler_plan_id:
                plan_label = plan_pair.label
                if not batch.bind_success:
                    plan_label = f"BIND_FAILED:{batch.bind_error}"
                updated = get_video_discovery_service().mark_candidates_aigc_plans(
                    id_batch,
                    crawler_plan_id=crawler_plan_id,
                    produce_plan_id=plan_pair.produce_plan_id,
                    publish_plan_id=plan_pair.publish_plan_id,
                    plan_label=plan_label,
                )
                logger.info(
                    "published batch: plan=%s crawler=%s videos=%d db_updated=%d bind_success=%s",
                    plan_pair.label,
                    crawler_plan_id,
                    len(aweme_batch),
                    updated,
                    batch.bind_success,
                )

            batch_results.append(batch)

    failed_batches = [item for item in batch_results if item.error or not item.bind_success]
    return {
        "success": not failed_batches,
        "biz_dt": resolved_biz_dt,
        "any_biz_dt": any_biz_dt,
        "run_id": run_id,
        "plan_count": len(plan_pairs),
        "candidate_count": len(candidates),
        "decision_buckets": list(_PUBLISHABLE_BUCKETS),
        "distribution": distribution_preview,
        "batch_count": len(batch_results),
        "failed_batch_count": len(failed_batches),
        "batches": [item.to_dict() for item in batch_results],
    }


if __name__ == "__main__":
    import argparse

    from supply_infra.scheduler.cli_result import run_cli

    def _build_parser() -> argparse.ArgumentParser:
        parser = argparse.ArgumentParser(
            description="均匀分发 primary 候选视频到 AIGC 发布计划"
        )
        parser.add_argument("biz_dt", nargs="?", help="业务日 YYYYMMDD，默认取最新")
        parser.add_argument("--run-id", dest="run_id", help="仅处理指定 run_id")
        parser.add_argument("--limit", type=int, help="最多处理候选视频数")
        parser.add_argument(
            "--force",
            action="store_true",
            help="包含已写过 aigc_crawler_plan_id 的候选",
        )
        return parser

    _args = _build_parser().parse_args()
    run_cli(
        lambda: publish_videos_from_discovery(
            biz_dt=_args.biz_dt,
            run_id=_args.run_id,
            skip_published=not _args.force,
            limit=_args.limit,
        ),
        label="publish_videos_from_discovery",
    )
