content-find-agent-new/aigc_push_registry/registry.py

from __future__ import annotations

import json
from dataclasses import dataclass
from datetime import datetime, timezone
from decimal import Decimal
from pathlib import Path
from typing import Any

from content_agent.integrations.database_runtime import ContentSupplyDbConfig


SCHEMA_VERSION = "content_agent.v1"


CREATE_TABLES_SQL = """
CREATE TABLE IF NOT EXISTS content_agent_aigc_push_records (
  id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
  schema_version VARCHAR(80) NOT NULL DEFAULT 'content_agent.v1',
  push_record_id VARCHAR(160) NOT NULL,
  run_id VARCHAR(80) NOT NULL,
  policy_run_id VARCHAR(80) NULL,
  platform VARCHAR(32) NOT NULL,
  demand_content_id BIGINT UNSIGNED NULL,
  demand_name VARCHAR(512) NULL,
  demand_dt VARCHAR(32) NULL,
  hive_merge_leve2 VARCHAR(160) NULL,
  hive_gap_dt VARCHAR(32) NULL,
  evidence_platform VARCHAR(80) NULL,
  run_label VARCHAR(240) NULL,
  run_status VARCHAR(32) NULL,
  pushed_at DATETIME(3) NULL,
  push_status VARCHAR(32) NOT NULL DEFAULT 'pushed',
  source_kind VARCHAR(80) NOT NULL DEFAULT 'manual',
  source_ref VARCHAR(512) NULL,
  uploaded_by VARCHAR(120) NULL,
  notes TEXT NULL,
  raw_payload JSON NULL,
  created_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
  updated_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
  PRIMARY KEY (id),
  UNIQUE KEY uk_ca_aigc_push_record (push_record_id),
  KEY idx_ca_aigc_push_run (run_id),
  KEY idx_ca_aigc_push_demand (demand_content_id),
  KEY idx_ca_aigc_push_platform_time (platform, pushed_at),
  KEY idx_ca_aigc_push_status (push_status)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;

CREATE TABLE IF NOT EXISTS content_agent_aigc_push_plans (
  id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
  schema_version VARCHAR(80) NOT NULL DEFAULT 'content_agent.v1',
  push_record_id VARCHAR(160) NOT NULL,
  plan_part VARCHAR(32) NOT NULL,
  produce_plan_id VARCHAR(120) NOT NULL,
  crawler_plan_id VARCHAR(120) NOT NULL,
  content_count INT UNSIGNED NOT NULL DEFAULT 0,
  pushed_at DATETIME(3) NULL,
  raw_payload JSON NULL,
  created_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
  updated_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
  PRIMARY KEY (id),
  UNIQUE KEY uk_ca_aigc_push_plan (push_record_id, plan_part),
  KEY idx_ca_aigc_push_plan_crawler (crawler_plan_id),
  KEY idx_ca_aigc_push_plan_produce (produce_plan_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;

CREATE TABLE IF NOT EXISTS content_agent_aigc_push_items (
  id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
  schema_version VARCHAR(80) NOT NULL DEFAULT 'content_agent.v1',
  push_record_id VARCHAR(160) NOT NULL,
  plan_part VARCHAR(32) NOT NULL,
  item_index INT UNSIGNED NOT NULL,
  platform_content_id VARCHAR(160) NOT NULL,
  decision_id VARCHAR(80) NULL,
  content_title TEXT NULL,
  author_display_name VARCHAR(256) NULL,
  raw_payload JSON NULL,
  created_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
  updated_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
  PRIMARY KEY (id),
  UNIQUE KEY uk_ca_aigc_push_item (push_record_id, platform_content_id),
  KEY idx_ca_aigc_push_item_part (push_record_id, plan_part),
  KEY idx_ca_aigc_push_item_content (platform_content_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
"""


@dataclass(frozen=True)
class AigcPushRegistry:
    config: ContentSupplyDbConfig

    @classmethod
    def from_env(cls, env_file: str | Path | None = ".env") -> "AigcPushRegistry":
        return cls(ContentSupplyDbConfig.from_env(env_file=env_file))

    def connect(self) -> Any:
        return self.config.connect()

    def init_db(self) -> None:
        with self.connect() as conn:
            with conn.cursor() as cur:
                for statement in _split_sql(CREATE_TABLES_SQL):
                    cur.execute(statement)
            conn.commit()

    def upsert_record(self, record: dict[str, Any]) -> None:
        prepared = normalize_manual_record(record)
        with self.connect() as conn:
            with conn.cursor() as cur:
                _upsert(cur, "content_agent_aigc_push_records", _record_row(prepared), ("push_record_id",))
                cur.execute(
                    "DELETE FROM content_agent_aigc_push_plans WHERE push_record_id = %s",
                    (prepared["push_record_id"],),
                )
                cur.execute(
                    "DELETE FROM content_agent_aigc_push_items WHERE push_record_id = %s",
                    (prepared["push_record_id"],),
                )
                for plan in prepared["plans"]:
                    _upsert(
                        cur,
                        "content_agent_aigc_push_plans",
                        _plan_row(prepared, plan),
                        ("push_record_id", "plan_part"),
                    )
                    for index, item in enumerate(plan.get("items") or [], start=1):
                        _upsert(
                            cur,
                            "content_agent_aigc_push_items",
                            _item_row(prepared, plan, item, index),
                            ("push_record_id", "platform_content_id"),
                        )
            conn.commit()

    def list_records(
        self,
        *,
        platform: str | None = None,
        demand_content_id: int | None = None,
        run_id: str | None = None,
        limit: int = 200,
    ) -> dict[str, Any]:
        where: list[str] = ["r.push_status = 'pushed'"]
        params: list[Any] = []
        if platform:
            where.append("r.platform = %s")
            params.append(platform)
        if demand_content_id is not None:
            where.append("r.demand_content_id = %s")
            params.append(demand_content_id)
        if run_id:
            where.append("r.run_id = %s")
            params.append(run_id)
        where_sql = "WHERE " + " AND ".join(where)
        params.append(max(1, min(int(limit), 500)))
        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"""
                    SELECT r.*,
                           COALESCE(SUM(p.content_count), 0) AS pushed_content_count,
                           COUNT(DISTINCT p.crawler_plan_id) AS crawler_plan_count
                    FROM content_agent_aigc_push_records r
                    LEFT JOIN content_agent_aigc_push_plans p
                      ON p.push_record_id = r.push_record_id
                    {where_sql}
                    GROUP BY r.id
                    ORDER BY r.pushed_at DESC, r.id DESC
                    LIMIT %s
                    """,
                    tuple(params),
                )
                records = [_decode_row(row) for row in cur.fetchall()]
                ids = [row["push_record_id"] for row in records]
                plans = _fetch_grouped(cur, "content_agent_aigc_push_plans", ids)
                items = _fetch_grouped(cur, "content_agent_aigc_push_items", ids)
        for record in records:
            push_id = record["push_record_id"]
            record["plans"] = plans.get(push_id, [])
            record["items"] = items.get(push_id, [])
        return {
            "items": records,
            "total": len(records),
            "data_origin": "production_db",
        }

    def build_record_from_run(
        self,
        *,
        run_id: str,
        platform: str,
        pushed_at: str,
        plans: list[dict[str, Any]],
        source_kind: str = "manual",
        source_ref: str | None = None,
        notes: str | None = None,
        uploaded_by: str | None = None,
    ) -> dict[str, Any]:
        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT r.run_id,
                           (SELECT p.policy_run_id
                            FROM content_agent_policy_runs p
                            WHERE p.run_id = r.run_id
                            ORDER BY p.id DESC LIMIT 1) AS policy_run_id,
                           r.platform, r.status AS run_status,
                           r.demand_content_id, d.name AS demand_name, d.dt AS demand_dt,
                           d.merge_leve2,
                           JSON_UNQUOTE(JSON_EXTRACT(d.ext_data,'$.run_label')) AS run_label,
                           JSON_UNQUOTE(JSON_EXTRACT(d.ext_data,'$.evidence_pack.demand_scope.merge_leve2')) AS hive_merge_leve2,
                           JSON_UNQUOTE(JSON_EXTRACT(d.ext_data,'$.evidence_pack.demand_scope.gap_dt')) AS hive_gap_dt,
                           JSON_UNQUOTE(JSON_EXTRACT(d.ext_data,'$.evidence_pack.demand_scope.platform')) AS evidence_platform
                    FROM content_agent_runs r
                    LEFT JOIN demand_content d ON d.id = r.demand_content_id
                    WHERE r.run_id = %s
                    LIMIT 1
                    """,
                    (run_id,),
                )
                run = cur.fetchone()
                if not run:
                    raise ValueError(f"run not found: {run_id}")
                cur.execute(
                    """
                    SELECT di.platform_content_id, rd.decision_id, di.description AS content_title,
                           di.author_display_name
                    FROM content_agent_rule_decisions rd
                    JOIN content_agent_discovered_content_items di
                      ON rd.run_id = di.run_id AND rd.decision_target_id = di.platform_content_id
                    WHERE rd.run_id = %s
                      AND di.platform = %s
                      AND rd.decision_action = 'ADD_TO_CONTENT_POOL'
                    ORDER BY di.id
                    """,
                    (run_id, platform),
                )
                pooled = cur.fetchall()
        cursor = 0
        planned: list[dict[str, Any]] = []
        for plan in plans:
            count = int(plan.get("content_count") or 0)
            if not count:
                count = len(plan.get("video_ids") or plan.get("content_ids") or [])
            plan_items = pooled[cursor: cursor + count] if count else []
            cursor += count
            planned.append(
                {
                    **plan,
                    "items": [
                        {
                            "platform_content_id": row["platform_content_id"],
                            "decision_id": row.get("decision_id"),
                            "content_title": row.get("content_title"),
                            "author_display_name": row.get("author_display_name"),
                        }
                        for row in plan_items
                    ],
                }
            )
        pushed_at_dt = parse_datetime(pushed_at)
        return normalize_manual_record(
            {
                "push_record_id": f"aigc_push_{run_id}_{platform}_{pushed_at_dt.strftime('%Y%m%d%H%M%S')}",
                "run_id": run_id,
                "policy_run_id": run.get("policy_run_id"),
                "platform": platform,
                "demand_content_id": run.get("demand_content_id"),
                "demand_name": run.get("demand_name"),
                "demand_dt": run.get("demand_dt"),
                "hive_merge_leve2": run.get("hive_merge_leve2") or run.get("merge_leve2"),
                "hive_gap_dt": run.get("hive_gap_dt"),
                "evidence_platform": run.get("evidence_platform"),
                "run_label": run.get("run_label"),
                "run_status": run.get("run_status"),
                "pushed_at": pushed_at_dt.isoformat(),
                "push_status": "pushed",
                "source_kind": source_kind,
                "source_ref": source_ref,
                "uploaded_by": uploaded_by,
                "notes": notes,
                "plans": planned,
            }
        )


def import_live_logs(registry: AigcPushRegistry, log_dir: Path) -> list[dict[str, Any]]:
    live_by_run: dict[str, list[dict[str, Any]]] = {}
    for path in sorted(log_dir.glob("aigc_push_*.jsonl")):
        for line in path.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            record = json.loads(line)
            if record.get("mode") != "live" or record.get("status") != "ok":
                continue
            record["_source_file"] = str(path)
            live_by_run.setdefault(str(record["run_id"]), []).append(record)
    imported: list[dict[str, Any]] = []
    for run_id, rows in sorted(live_by_run.items()):
        rows.sort(key=lambda row: str(row.get("half") or ""))
        platform = str(rows[0]["platform"])
        plans = [
            {
                "plan_part": str(row.get("half") or index),
                "produce_plan_id": str(row["produce_plan_id"]),
                "crawler_plan_id": str(row["crawler_plan_id"]),
                "content_count": int(row.get("content_count") or row.get("aweme_count") or 0),
                "pushed_at": row.get("ts"),
                "raw_payload": {"source_file": row.get("_source_file")},
            }
            for index, row in enumerate(rows, start=1)
        ]
        pushed_at = rows[0].get("ts") or datetime.now(timezone.utc).isoformat()
        record = registry.build_record_from_run(
            run_id=run_id,
            platform=platform,
            pushed_at=pushed_at,
            plans=plans,
            source_kind="e2e_log_import",
            source_ref=";".join(sorted({str(row.get("_source_file")) for row in rows})),
            notes="Imported from live AIGC push logs; dry-run records ignored.",
        )
        registry.upsert_record(record)
        imported.append(record)
    return imported


def normalize_manual_record(record: dict[str, Any]) -> dict[str, Any]:
    plans = record.get("plans") or []
    if not plans:
        raise ValueError("manual AIGC push record must include plans")
    push_record_id = str(record.get("push_record_id") or "").strip()
    if not push_record_id:
        raise ValueError("push_record_id is required")
    result = {
        **record,
        "schema_version": record.get("schema_version") or SCHEMA_VERSION,
        "push_record_id": push_record_id,
        "push_status": record.get("push_status") or "pushed",
        "source_kind": record.get("source_kind") or "manual",
        "pushed_at": parse_datetime(record["pushed_at"]).isoformat() if record.get("pushed_at") else None,
        "plans": [],
    }
    if result["push_status"] != "pushed":
        raise ValueError("only real pushed AIGC records are stored; dry-run records are ignored")
    for plan in plans:
        part = str(plan.get("plan_part") or plan.get("half") or "").strip()
        if not part:
            raise ValueError("each plan needs plan_part/half")
        crawler_plan_id = str(plan.get("crawler_plan_id") or "").strip()
        if not crawler_plan_id:
            raise ValueError("crawler_plan_id is required for real pushed records")
        produce_plan_id = str(plan.get("produce_plan_id") or "").strip()
        if not produce_plan_id:
            raise ValueError("produce_plan_id is required")
        raw_items = plan.get("items")
        if raw_items is None:
            ids = plan.get("video_ids") or plan.get("content_ids") or []
            raw_items = [{"platform_content_id": item_id} for item_id in ids]
        result["plans"].append(
            {
                **plan,
                "plan_part": part,
                "crawler_plan_id": crawler_plan_id,
                "produce_plan_id": produce_plan_id,
                "pushed_at": parse_datetime(plan["pushed_at"]).isoformat() if plan.get("pushed_at") else result["pushed_at"],
                "items": raw_items,
                "content_count": int(plan.get("content_count") or len(raw_items)),
            }
        )
    return result


def parse_datetime(value: str | datetime) -> datetime:
    if isinstance(value, datetime):
        return value
    text = str(value).strip()
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    if " " in text and "T" not in text:
        text = text.replace(" ", "T")
    parsed = datetime.fromisoformat(text)
    if parsed.tzinfo is not None:
        parsed = parsed.astimezone(timezone.utc).replace(tzinfo=None)
    return parsed


def _record_row(record: dict[str, Any]) -> dict[str, Any]:
    keys = [
        "schema_version",
        "push_record_id",
        "run_id",
        "policy_run_id",
        "platform",
        "demand_content_id",
        "demand_name",
        "demand_dt",
        "hive_merge_leve2",
        "hive_gap_dt",
        "evidence_platform",
        "run_label",
        "run_status",
        "pushed_at",
        "push_status",
        "source_kind",
        "source_ref",
        "uploaded_by",
        "notes",
        "raw_payload",
    ]
    row = {key: record.get(key) for key in keys}
    row["pushed_at"] = parse_datetime(record["pushed_at"]) if record.get("pushed_at") else None
    return row


def _plan_row(record: dict[str, Any], plan: dict[str, Any]) -> dict[str, Any]:
    return {
        "schema_version": SCHEMA_VERSION,
        "push_record_id": record["push_record_id"],
        "plan_part": plan["plan_part"],
        "produce_plan_id": plan["produce_plan_id"],
        "crawler_plan_id": plan["crawler_plan_id"],
        "content_count": plan["content_count"],
        "pushed_at": parse_datetime(plan["pushed_at"]) if plan.get("pushed_at") else None,
        "raw_payload": plan.get("raw_payload"),
    }


def _item_row(record: dict[str, Any], plan: dict[str, Any], item: dict[str, Any], index: int) -> dict[str, Any]:
    return {
        "schema_version": SCHEMA_VERSION,
        "push_record_id": record["push_record_id"],
        "plan_part": plan["plan_part"],
        "item_index": index,
        "platform_content_id": str(item.get("platform_content_id") or item.get("content_id") or item.get("aweme_id")),
        "decision_id": item.get("decision_id"),
        "content_title": item.get("content_title") or item.get("title"),
        "author_display_name": item.get("author_display_name"),
        "raw_payload": item.get("raw_payload"),
    }


def _upsert(cur: Any, table: str, row: dict[str, Any], key_columns: tuple[str, ...]) -> None:
    clean = {key: value for key, value in row.items() if value is not None}
    columns = list(clean)
    placeholders = ", ".join(["%s"] * len(columns))
    column_sql = ", ".join(f"`{column}`" for column in columns)
    updates = ", ".join(
        f"`{column}` = VALUES(`{column}`)"
        for column in columns
        if column not in key_columns
    )
    values = [_db_value(value) for value in clean.values()]
    sql = f"INSERT INTO `{table}` ({column_sql}) VALUES ({placeholders})"
    if updates:
        sql += f" ON DUPLICATE KEY UPDATE {updates}"
    cur.execute(sql, values)


def _fetch_grouped(cur: Any, table: str, push_ids: list[str]) -> dict[str, list[dict[str, Any]]]:
    if not push_ids:
        return {}
    cur.execute(
        f"SELECT * FROM `{table}` WHERE push_record_id IN %s ORDER BY push_record_id, id",
        (tuple(push_ids),),
    )
    result: dict[str, list[dict[str, Any]]] = {}
    for row in cur.fetchall():
        decoded = _decode_row(row)
        result.setdefault(str(decoded["push_record_id"]), []).append(decoded)
    return result


def _decode_row(row: dict[str, Any]) -> dict[str, Any]:
    decoded = dict(row)
    for key, value in list(decoded.items()):
        if isinstance(value, (datetime,)):
            decoded[key] = value.isoformat(sep=" ")
        elif isinstance(value, bytes):
            decoded[key] = value.decode("utf-8")
        elif isinstance(value, Decimal):
            decoded[key] = int(value) if value == value.to_integral_value() else float(value)
        elif key in {"raw_payload"} and isinstance(value, str):
            try:
                decoded[key] = json.loads(value)
            except json.JSONDecodeError:
                pass
    return decoded


def _db_value(value: Any) -> Any:
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False)
    if isinstance(value, datetime):
        return value.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
    return value


def _split_sql(sql: str) -> list[str]:
    return [part.strip() for part in sql.split(";") if part.strip()]