content-find-agent-new/content_agent/integrations/oss_archive.py

from __future__ import annotations

import os
import time
from datetime import datetime, timedelta, timezone
from threading import Lock
from typing import Any, Callable
from urllib.parse import urlparse

from content_agent.integrations import oss_upload, timeout_config
from content_agent.integrations.bounded_pool import DaemonThreadPoolExecutor, run_bounded
from content_agent.interfaces import RuntimeFileStore


# 单次 OSS 上传尝试上限(原 3600,即便生效也能挂 1 小时)。env 可覆盖,被硬上限钳制。
DEFAULT_OSS_ATTEMPT_TIMEOUT_SECONDS = timeout_config.total_timeout("oss")
# 主线程等单条 OSS worker 的看门狗:略大于单次尝试上限,只当兜底。
OSS_WORKER_RESULT_TIMEOUT_SECONDS = timeout_config.total_timeout("oss") + 60.0
DEFAULT_OSS_ARCHIVE_WINDOW_SECONDS = 24 * 60 * 60.0
DEFAULT_OSS_RETRY_DELAY_SECONDS = 15 * 60.0
DEFAULT_OSS_ARCHIVE_MAX_WORKERS = 5
DEFAULT_OSS_URL_FALLBACK_MAX_ATTEMPTS = 3
DEFAULT_OSS_MEDIA_UPLOAD_BUDGET_SECONDS = 360.0


class AsyncArchiveDispatcher:
    def __init__(
        self,
        runtime: RuntimeFileStore,
        run_id: str,
        *,
        upload_fn: Callable[..., dict[str, Any]] = oss_upload.upload_video_from_env,
        now_fn: Callable[[], datetime] | None = None,
        attempt_timeout_seconds: float = DEFAULT_OSS_ATTEMPT_TIMEOUT_SECONDS,
        retry_delay_seconds: float = DEFAULT_OSS_RETRY_DELAY_SECONDS,
        max_workers: int | None = None,
    ) -> None:
        self.runtime = runtime
        self.run_id = run_id
        self.upload_fn = upload_fn
        self.now_fn = now_fn
        self.attempt_timeout_seconds = attempt_timeout_seconds
        self.retry_delay_seconds = retry_delay_seconds
        self._max_workers = _resolve_max_workers(max_workers)
        # daemon 池:卡死的上传 worker 不阻止进程退出(配合 OSS read 短超时,worker 必有限时间自终)。
        self._executor = DaemonThreadPoolExecutor(
            max_workers=self._max_workers, thread_name_prefix="oss-archive"
        )
        self._lock = Lock()
        self._write_lock = Lock()
        self._completed: list[dict[str, Any]] = []
        self._runtime_writes_enabled = False

    def submit_records(self, records: list[dict[str, Any]]) -> None:
        now = _utc_now(self.now_fn)
        for record in records:
            if not _is_due(record, now):
                continue
            future = self._executor.submit(
                _archive_one,
                record,
                now,
                upload_fn=self.upload_fn,
                attempt_timeout_seconds=self.attempt_timeout_seconds,
                retry_delay_seconds=self.retry_delay_seconds,
            )
            future.add_done_callback(self._store_completed)

    def archive_records(self, records: list[dict[str, Any]]) -> list[dict[str, Any]]:
        return archive_due_records(
            records,
            upload_fn=self.upload_fn,
            now_fn=self.now_fn,
            attempt_timeout_seconds=self.attempt_timeout_seconds,
            retry_delay_seconds=self.retry_delay_seconds,
            max_workers=self._max_workers,
        )

    def drain_completed(self) -> list[dict[str, Any]]:
        with self._lock:
            completed = list(self._completed)
            self._completed.clear()
        return completed

    def enable_runtime_writes(self) -> None:
        with self._lock:
            self._runtime_writes_enabled = True
            completed = list(self._completed)
            self._completed.clear()
        self._write_records(completed)

    def shutdown(self, *, wait: bool = False) -> None:
        # wait=True 是"排空全部归档"语义 → 不取消队列;wait=False 是"放弃卡住的"→ 取消未启动任务。
        self._executor.shutdown(wait=wait, cancel_futures=not wait)

    def _store_completed(self, future: Any) -> None:
        try:
            record = future.result()
        except Exception:
            return
        write_now = False
        with self._lock:
            if self._runtime_writes_enabled:
                write_now = True
            else:
                self._completed.append(record)
        if write_now:
            self._write_records([record])

    def _write_records(self, records: list[dict[str, Any]]) -> None:
        if not records:
            return
        with self._write_lock:
            self.runtime.append_jsonl(self.run_id, "content_media_records.jsonl", records)


def mark_archive_pending(
    records: list[dict[str, Any]],
    *,
    now_fn: Callable[[], datetime] | None = None,
    archive_window_seconds: float = DEFAULT_OSS_ARCHIVE_WINDOW_SECONDS,
) -> list[dict[str, Any]]:
    now = _utc_now(now_fn)
    return [
        _mark_one_pending(record, now, archive_window_seconds=archive_window_seconds)
        for record in records
    ]


def archive_due_records(
    records: list[dict[str, Any]],
    *,
    upload_fn: Callable[..., dict[str, Any]] = oss_upload.upload_video_from_env,
    now_fn: Callable[[], datetime] | None = None,
    attempt_timeout_seconds: float = DEFAULT_OSS_ATTEMPT_TIMEOUT_SECONDS,
    retry_delay_seconds: float = DEFAULT_OSS_RETRY_DELAY_SECONDS,
    max_workers: int | None = None,
) -> list[dict[str, Any]]:
    now = _utc_now(now_fn)
    archived = list(records)
    due_records: list[tuple[int, dict[str, Any]]] = []
    for index, record in enumerate(records):
        if not _is_due(record, now):
            continue
        due_records.append((index, record))
    if not due_records:
        return archived

    worker_count = _resolve_max_workers(max_workers)
    if worker_count <= 1 or len(due_records) == 1:
        for index, record in due_records:
            archived[index] = _archive_one(
                record,
                now,
                upload_fn=upload_fn,
                attempt_timeout_seconds=attempt_timeout_seconds,
                retry_delay_seconds=retry_delay_seconds,
            )
        return archived

    due_only = [record for _, record in due_records]

    def _work(record: dict[str, Any]) -> dict[str, Any]:
        return _archive_one(
            record,
            now,
            upload_fn=upload_fn,
            attempt_timeout_seconds=attempt_timeout_seconds,
            retry_delay_seconds=retry_delay_seconds,
        )

    def _on_timeout(record: dict[str, Any], _offset: int) -> dict[str, Any]:
        raw_payload = dict(record.get("raw_payload") or {})
        attempt_count = int(raw_payload.get("oss_archive_attempt_count") or 0) + 1
        return _with_failed_archive(record, raw_payload, now, attempt_count, "oss_worker_timeout")

    results = run_bounded(
        due_only,
        _work,
        max_workers=worker_count,
        per_future_timeout=OSS_WORKER_RESULT_TIMEOUT_SECONDS,
        on_timeout=_on_timeout,
        thread_name_prefix="oss-archive",
    )
    for (index, _record), result in zip(due_records, results):
        archived[index] = result
    return archived


def archive_record(
    record: dict[str, Any],
    *,
    upload_fn: Callable[..., dict[str, Any]] = oss_upload.upload_video_from_env,
    now_fn: Callable[[], datetime] | None = None,
    attempt_timeout_seconds: float = DEFAULT_OSS_ATTEMPT_TIMEOUT_SECONDS,
    retry_delay_seconds: float = DEFAULT_OSS_RETRY_DELAY_SECONDS,
) -> dict[str, Any]:
    """Archive one due media record using the same status contract as batch archive."""
    now = _utc_now(now_fn)
    if not _is_due(record, now):
        return record
    return _archive_one(
        record,
        now,
        upload_fn=upload_fn,
        attempt_timeout_seconds=attempt_timeout_seconds,
        retry_delay_seconds=retry_delay_seconds,
    )


def archive_pending_for_run(
    runtime: RuntimeFileStore,
    run_id: str,
    *,
    upload_fn: Callable[..., dict[str, Any]] = oss_upload.upload_video_from_env,
    now_fn: Callable[[], datetime] | None = None,
    attempt_timeout_seconds: float = DEFAULT_OSS_ATTEMPT_TIMEOUT_SECONDS,
    retry_delay_seconds: float = DEFAULT_OSS_RETRY_DELAY_SECONDS,
    max_workers: int | None = None,
) -> list[dict[str, Any]]:
    records = runtime.read_jsonl(run_id, "content_media_records.jsonl")
    archived = archive_due_records(
        records,
        upload_fn=upload_fn,
        now_fn=now_fn,
        attempt_timeout_seconds=attempt_timeout_seconds,
        retry_delay_seconds=retry_delay_seconds,
        max_workers=max_workers,
    )
    runtime.append_jsonl(run_id, "content_media_records.jsonl", archived)
    return archived


def _mark_one_pending(
    record: dict[str, Any],
    now: datetime,
    *,
    archive_window_seconds: float,
) -> dict[str, Any]:
    if not record.get("play_url"):
        return record
    if record.get("content_media_status") == "oss_uploaded":
        return record
    raw_payload = dict(record.get("raw_payload") or {})
    deadline = raw_payload.get("oss_archive_deadline_at")
    if not deadline:
        deadline = _iso(now + timedelta(seconds=archive_window_seconds))
    updated_raw = {
        **raw_payload,
        "content_media_status": "oss_upload_pending",
        "oss_url": record.get("oss_url"),
        "local_path": record.get("local_path"),
        "oss_archive_status": "pending",
        "oss_archive_attempt_count": int(raw_payload.get("oss_archive_attempt_count") or 0),
        "oss_archive_next_retry_at": raw_payload.get("oss_archive_next_retry_at") or _iso(now),
        "oss_archive_deadline_at": deadline,
        "oss_archive_updated_at": _iso(now),
    }
    return {
        **record,
        "content_media_status": "oss_upload_pending",
        "oss_url": record.get("oss_url"),
        "local_path": record.get("local_path"),
        "raw_payload": _compact(updated_raw),
    }


def _archive_one(
    record: dict[str, Any],
    now: datetime,
    *,
    upload_fn: Callable[..., dict[str, Any]],
    attempt_timeout_seconds: float,
    retry_delay_seconds: float,
) -> dict[str, Any]:
    raw_payload = dict(record.get("raw_payload") or {})
    previous_attempt_count = int(raw_payload.get("oss_archive_attempt_count") or 0)
    deadline = _parse_time(raw_payload.get("oss_archive_deadline_at"))
    if deadline and now >= deadline:
        return _with_failed_archive(record, raw_payload, now, previous_attempt_count + 1, "oss_upload_failed")
    media_started_at = time.monotonic()
    candidates = _upload_candidates(record, raw_payload)
    max_attempts = _resolve_url_fallback_max_attempts()
    media_budget_seconds = _resolve_media_upload_budget_seconds()
    upload_result: dict[str, Any] = {"status": "failed", "failure_type": "missing_src_url"}
    selected_candidate: dict[str, Any] | None = None
    oss_url_attempts: list[dict[str, Any]] = []
    for candidate in candidates[:max_attempts]:
        remaining_seconds = media_budget_seconds - (time.monotonic() - media_started_at)
        if remaining_seconds <= 0:
            upload_result = {
                "status": "failed",
                "failure_type": "oss_upload_budget_exceeded",
                "attempt_timeout_seconds": attempt_timeout_seconds,
                "response_absent": True,
            }
            break
        candidate_started_at = time.monotonic()
        upload_result = upload_fn(
            candidate.get("url") or "",
            timeout_seconds=min(attempt_timeout_seconds, remaining_seconds),
        )
        oss_url_attempts.append(
            _oss_url_attempt_summary(
                candidate,
                upload_result,
                duration_ms=_elapsed_ms(candidate_started_at, time.monotonic()),
            )
        )
        if upload_result.get("status") == "ok":
            selected_candidate = candidate
            break
    if not oss_url_attempts and upload_result.get("failure_type") == "missing_src_url":
        oss_url_attempts.append(
            {
                "candidate_index": 0,
                "host": "",
                "path": "",
                "status": "failed",
                "failure_type": "missing_src_url",
                "duration_ms": 0,
            }
        )
    attempt_count = previous_attempt_count + len(oss_url_attempts)
    timing_metrics = {
        "oss_upload_duration_ms": _elapsed_ms(media_started_at, time.monotonic()),
        "oss_upload_timeout_seconds": attempt_timeout_seconds,
        "oss_media_upload_budget_seconds": media_budget_seconds,
        "oss_archive_attempt_started_at": _iso(now),
    }
    if upload_result.get("status") == "ok":
        selected_candidate = selected_candidate or {}
        updated_raw = {
            **raw_payload,
            "content_media_status": "oss_uploaded",
            "oss_url": upload_result.get("oss_url"),
            "local_path": None,
            "oss_archive_status": "uploaded",
            "oss_archive_attempt_count": attempt_count,
            "oss_archive_next_retry_at": None,
            "oss_archive_updated_at": _iso(now),
            "oss_object_key": upload_result.get("oss_object_key"),
            "save_oss_timestamp": upload_result.get("save_oss_timestamp"),
            "oss_timing_metrics": timing_metrics,
            "oss_payload_mode": upload_result.get("oss_payload_mode"),
            "oss_url_attempts": oss_url_attempts,
            "oss_archive_attempted_candidate_count": len(oss_url_attempts),
            "oss_archive_selected_video_url_host": selected_candidate.get("host"),
            "oss_archive_selected_video_url_path": selected_candidate.get("path"),
            "oss_archive_selected_video_url_index": selected_candidate.get("candidate_index"),
        }
        return {
            **record,
            "content_media_status": "oss_uploaded",
            "oss_url": upload_result.get("oss_url"),
            "local_path": None,
            "raw_payload": _compact(updated_raw),
        }
    failure = upload_result.get("failure_type") or "oss_upload_failed"
    if deadline and now >= deadline:
        return _with_failed_archive(record, raw_payload, now, attempt_count, failure, upload_result)
    updated_raw = {
        **raw_payload,
        "content_media_status": "oss_upload_pending",
        "oss_url": None,
        "local_path": record.get("local_path"),
        "failure_reason": failure,
        "oss_archive_status": "pending",
        "oss_archive_attempt_count": attempt_count,
        "oss_archive_next_retry_at": _iso(now + timedelta(seconds=retry_delay_seconds)),
        "oss_archive_last_error": failure,
        "oss_archive_last_exception_type": upload_result.get("exception_type"),
        "oss_archive_last_error_message": upload_result.get("error_message"),
        "oss_archive_last_http_status_code": upload_result.get("http_status_code"),
        "oss_payload_mode": upload_result.get("oss_payload_mode"),
        "oss_response_summary": upload_result.get("oss_response_summary"),
        "oss_endpoint_host": upload_result.get("endpoint_host"),
        "oss_read_timeout_seconds": upload_result.get("read_timeout_seconds"),
        "oss_attempt_timeout_seconds": upload_result.get("attempt_timeout_seconds"),
        "oss_response_absent": upload_result.get("response_absent"),
        "oss_archive_updated_at": _iso(now),
        "upload_failure_reason": failure,
        "oss_timing_metrics": timing_metrics,
        "oss_url_attempts": oss_url_attempts,
        "oss_archive_attempted_candidate_count": len(oss_url_attempts),
    }
    return {
        **record,
        "content_media_status": "oss_upload_pending",
        "oss_url": None,
        "local_path": record.get("local_path"),
        "failure_reason": failure,
        "raw_payload": _compact(updated_raw),
    }


def _with_failed_archive(
    record: dict[str, Any],
    raw_payload: dict[str, Any],
    now: datetime,
    attempt_count: int,
    failure: str,
    upload_result: dict[str, Any] | None = None,
) -> dict[str, Any]:
    upload_result = upload_result or {}
    updated_raw = {
        **raw_payload,
        "content_media_status": "oss_upload_failed",
        "oss_url": None,
        "local_path": record.get("local_path"),
        "failure_reason": failure,
        "oss_archive_status": "failed",
        "oss_archive_attempt_count": attempt_count,
        "oss_archive_next_retry_at": None,
        "oss_archive_last_error": failure,
        "oss_archive_last_exception_type": upload_result.get("exception_type"),
        "oss_archive_last_error_message": upload_result.get("error_message"),
        "oss_archive_last_http_status_code": upload_result.get("http_status_code"),
        "oss_payload_mode": upload_result.get("oss_payload_mode"),
        "oss_response_summary": upload_result.get("oss_response_summary"),
        "oss_archive_updated_at": _iso(now),
        "upload_failure_reason": failure,
    }
    return {
        **record,
        "content_media_status": "oss_upload_failed",
        "oss_url": None,
        "local_path": record.get("local_path"),
        "failure_reason": failure,
        "raw_payload": _compact(updated_raw),
    }


def _upload_candidates(record: dict[str, Any], raw_payload: dict[str, Any]) -> list[dict[str, Any]]:
    candidates = raw_payload.get("video_url_candidates")
    if not isinstance(candidates, list):
        candidates = []
    result: list[dict[str, Any]] = []
    seen: set[str] = set()
    play_url = str(record.get("play_url") or "")
    if play_url:
        matched = next(
            (
                candidate
                for candidate in candidates
                if isinstance(candidate, dict) and str(candidate.get("url") or "") == play_url
            ),
            None,
        )
        first = dict(matched or {})
        first.setdefault("url", play_url)
        first.setdefault("host", urlparse(play_url).netloc)
        first.setdefault("path", "play_url")
        first.setdefault("source", "play_url")
        first.setdefault("candidate_index", 0)
        result.append(first)
        seen.add(play_url)
    for candidate in candidates:
        if not isinstance(candidate, dict):
            continue
        url = str(candidate.get("url") or "")
        if not url or url in seen:
            continue
        seen.add(url)
        item = dict(candidate)
        item.setdefault("host", urlparse(url).netloc)
        item.setdefault("path", "")
        item.setdefault("candidate_index", len(result))
        result.append(item)
    return result


def _oss_url_attempt_summary(
    candidate: dict[str, Any],
    upload_result: dict[str, Any],
    *,
    duration_ms: int,
) -> dict[str, Any]:
    return _compact(
        {
            "candidate_index": candidate.get("candidate_index"),
            "host": candidate.get("host"),
            "path": candidate.get("path"),
            "status": upload_result.get("status"),
            "failure_type": upload_result.get("failure_type"),
            "exception_type": upload_result.get("exception_type"),
            "duration_ms": duration_ms,
        }
    )


def _is_due(record: dict[str, Any], now: datetime) -> bool:
    if record.get("content_media_status") == "oss_uploaded":
        return False
    if not record.get("play_url"):
        return False
    raw_payload = record.get("raw_payload") if isinstance(record.get("raw_payload"), dict) else {}
    status = raw_payload.get("oss_archive_status")
    if record.get("content_media_status") != "oss_upload_pending" and status != "pending":
        return False
    if status not in {"pending", None}:
        return False
    due_at = _parse_time(raw_payload.get("oss_archive_next_retry_at"))
    return due_at is None or now >= due_at


def _utc_now(now_fn: Callable[[], datetime] | None) -> datetime:
    now = now_fn() if now_fn else datetime.now(timezone.utc)
    if now.tzinfo is None:
        return now.replace(tzinfo=timezone.utc)
    return now.astimezone(timezone.utc)


def _parse_time(value: Any) -> datetime | None:
    if not value:
        return None
    try:
        parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


def _iso(value: datetime) -> str:
    return value.astimezone(timezone.utc).isoformat()


def _compact(payload: dict[str, Any]) -> dict[str, Any]:
    return dict(payload)


def _elapsed_ms(start: float, end: float) -> int:
    return max(0, int((end - start) * 1000))


def _resolve_max_workers(value: int | None = None) -> int:
    if value is not None and value > 0:
        return value
    raw = os.environ.get("CONTENT_AGENT_OSS_ARCHIVE_MAX_WORKERS")
    try:
        parsed = int(raw) if raw else DEFAULT_OSS_ARCHIVE_MAX_WORKERS
    except ValueError:
        return DEFAULT_OSS_ARCHIVE_MAX_WORKERS
    return parsed if parsed > 0 else DEFAULT_OSS_ARCHIVE_MAX_WORKERS


def _resolve_url_fallback_max_attempts() -> int:
    raw = os.environ.get("CONTENT_AGENT_OSS_URL_FALLBACK_MAX_ATTEMPTS")
    try:
        value = int(raw) if raw else DEFAULT_OSS_URL_FALLBACK_MAX_ATTEMPTS
    except (TypeError, ValueError):
        value = DEFAULT_OSS_URL_FALLBACK_MAX_ATTEMPTS
    return max(1, min(value, DEFAULT_OSS_URL_FALLBACK_MAX_ATTEMPTS))


def _resolve_media_upload_budget_seconds() -> float:
    raw = os.environ.get("CONTENT_AGENT_OSS_MEDIA_UPLOAD_BUDGET_SECONDS")
    try:
        value = float(raw) if raw else DEFAULT_OSS_MEDIA_UPLOAD_BUDGET_SECONDS
    except (TypeError, ValueError):
        value = DEFAULT_OSS_MEDIA_UPLOAD_BUDGET_SECONDS
    return max(1.0, min(value, DEFAULT_OSS_MEDIA_UPLOAD_BUDGET_SECONDS))