"""Utilities for coordinating screenshot batches via the Azure Function API."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil
from collections.abc import Callable, Iterable, Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any

from screenshot import ScreenshotOptions

from .client import ScreenshotServiceClient
from .storage import ArtifactStorage, NullStorage, create_azure_storage
from .utils import derive_job_id, ensure_capture_enabled, serialize_options

logger = logging.getLogger("screenshot_batch")


def _copy_file(src: Path, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dest)


@dataclass(frozen=True, kw_only=True, slots=True)
class ScreenshotJobSpec:
    """Description of a screenshot job to submit via the Azure service.

    Attributes:
        job_id: Unique identifier for the job.
        url: Target URL to screenshot.
        options: Screenshot configuration options.
        partition_date: Optional date partition for organization (YYYY-MM-DD format).
        html_snapshot_path: Optional path to HTML snapshot file.
        metadata: Additional job metadata (immutable mapping).
    """

    job_id: str
    url: str
    options: ScreenshotOptions
    partition_date: str | None = None
    html_snapshot_path: Path | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Validate job specification fields."""
        if not self.job_id or not self.job_id.strip():
            raise ValueError("job_id must be a non-empty string")
        if not self.url or not self.url.strip():
            raise ValueError("url must be a non-empty string")

    def to_payload(self, *, default_metadata: dict[str, Any] | None = None) -> dict[str, Any]:
        """Return a JSON-serialisable payload for the Azure Function API."""

        payload = {
            "job_id": self.job_id,
            "url": self.url,
            "options": serialize_options(self.options),
        }
        if self.partition_date:
            payload["partition_date"] = self.partition_date
        if self.html_snapshot_path:
            payload["html_snapshot_path"] = str(self.html_snapshot_path)

        merged_metadata: dict[str, Any] = {}
        if default_metadata:
            merged_metadata.update(default_metadata)
        merged_metadata.update(self.metadata)
        if merged_metadata:
            payload["metadata"] = merged_metadata

        return payload

    @classmethod
    def from_url(
        cls,
        url: str,
        *,
        job_id: str | None = None,
        options: ScreenshotOptions | None = None,
        partition_date: str | None = None,
        metadata: dict[str, Any] | None = None,
        html_snapshot_path: Path | None = None,
    ) -> ScreenshotJobSpec:
        """Convenience helper that derives a job identifier from the URL."""

        normalized_url = (url or "").strip()
        if not normalized_url:
            raise ValueError("url must be a non-empty string")

        resolved_job_id = job_id or derive_job_id(normalized_url, fallback="site")

        resolved_options = options or ScreenshotOptions()
        resolved_options = ensure_capture_enabled(resolved_options)

        metadata_payload = dict(metadata or {})

        return cls(
            job_id=resolved_job_id,
            url=normalized_url,
            options=resolved_options,
            partition_date=partition_date,
            metadata=metadata_payload,
            html_snapshot_path=html_snapshot_path,
        )


@dataclass(frozen=True, kw_only=True, slots=True)
class ScreenshotBatchResult:
    """Outcome of a remote screenshot batch run.

    Attributes:
        batch_id: Batch identifier.
        job_id: Job run identifier.
        status: Run status (completed, failed, timeout, etc.).
        job_total: Total number of jobs in batch.
        job_completed: Number of successfully completed jobs.
        job_failed: Number of failed jobs.
        status_payload: Raw status response from API.
        manifest_payload: Raw manifest response from API.
        local_root: Local directory where results were downloaded (if any).
    """

    batch_id: str
    job_id: str
    status: str
    job_total: int
    job_completed: int
    job_failed: int
    status_payload: dict[str, Any] = field(default_factory=dict)
    manifest_payload: dict[str, Any] = field(default_factory=dict)
    local_root: Path | None = None

    def __post_init__(self) -> None:
        """Validate result fields."""
        if self.job_total < 0:
            raise ValueError(f"job_total must be >= 0, got {self.job_total}")
        if self.job_completed < 0:
            raise ValueError(f"job_completed must be >= 0, got {self.job_completed}")
        if self.job_failed < 0:
            raise ValueError(f"job_failed must be >= 0, got {self.job_failed}")


def is_batch_successful(result: ScreenshotBatchResult) -> bool:
    """Check if a batch run completed successfully without failures.

    Args:
        result: Batch result to check.

    Returns:
        True if status is 'completed' and no jobs failed, False otherwise.
    """
    return result.status == "completed" and result.job_failed == 0


@dataclass(frozen=True, kw_only=True, slots=True)
class JobRunSummary:
    """Summary of an individual job run within a batch.

    Attributes:
        job_id: Job identifier.
        status: Job status (success, failed, skipped, etc.).
        blob_url: Azure blob URL where results are stored (if applicable).
        errors: Sequence of error messages for failed jobs.
        local_path: Local path where results were downloaded (if applicable).
    """

    job_id: str
    status: str
    blob_url: str | None = None
    errors: Sequence[str] = field(default_factory=tuple)
    local_path: Path | None = None


class ScreenshotBatchCoordinator:
    """High-level helper for pushing bulk screenshot jobs to the Azure service.

    This coordinator submits screenshot jobs to an Azure Function service and
    optionally downloads results from Azure blob storage.

    Args:
        base_url: Base URL of the screenshot service API.
        batch_id: Unique identifier for this batch of screenshots.
        store_dir: Local directory for storing downloaded results.
        concurrency: Number of concurrent screenshot jobs to run.
        poll_interval: Seconds between status polls.
        poll_timeout: Maximum seconds to wait for batch completion.
        result_metadata: Default metadata to merge into all jobs.
        download_results: Whether to download artifacts after completion.
        client: Optional pre-configured screenshot service client.
        storage: Optional storage backend for downloads. If None and download_results
            is True, creates Azure storage from environment variables.
    """

    def __init__(
        self,
        *,
        base_url: str,
        batch_id: str,
        store_dir: Path | None = None,
        concurrency: int = 4,
        poll_interval: float = 5.0,
        poll_timeout: float = 900.0,
        result_metadata: dict[str, Any] | None = None,
        download_results: bool = False,
        client: ScreenshotServiceClient | None = None,
        storage: ArtifactStorage | None = None,
    ) -> None:
        if not base_url:
            raise ValueError("base_url is required")
        if not batch_id:
            raise ValueError("batch_id is required")

        self.base_url = base_url
        self.batch_id = batch_id
        self.store_dir = Path(store_dir) if store_dir else None
        self.concurrency = max(1, int(concurrency))
        self.poll_interval = max(0.5, float(poll_interval))
        self.poll_timeout = max(self.poll_interval, float(poll_timeout))
        self.default_metadata = dict(result_metadata or {})
        self.download_results = bool(download_results and self.store_dir)
        self._client = client or ScreenshotServiceClient(base_url)

        # Set up storage backend
        if storage:
            self._storage = storage
        elif self.download_results:
            # Auto-create Azure storage from environment if needed
            azure_storage = create_azure_storage()
            self._storage = azure_storage if azure_storage else NullStorage()
        else:
            self._storage = NullStorage()

        container_from_storage = getattr(self._storage, "container_name", None)
        self._source_container_name = container_from_storage or os.getenv("SCREENSHOT_AZURE_STORAGE_CONTAINER") or None

    async def run_batch(self, jobs: Sequence[ScreenshotJobSpec]) -> ScreenshotBatchResult:
        """Submit a batch run and wait for completion."""

        if not jobs:
            raise ValueError("jobs must be a non-empty sequence")

        payload = {
            "jobs": [job.to_payload(default_metadata=self.default_metadata) for job in jobs],
            "concurrency": self.concurrency,
        }

        logger.info(
            "Submitting screenshot batch: batch=%s jobs=%d concurrency=%d",
            self.batch_id,
            len(jobs),
            self.concurrency,
        )

        handle = await asyncio.to_thread(self._client.start_batch, self.batch_id, payload)
        # Convert UUID to string for internal use
        job_id_str = str(handle.job_id)

        logger.info(
            "Screenshot batch scheduled: batch=%s run=%s",
            handle.batch_id,
            job_id_str,
        )

        status_payload = await asyncio.to_thread(
            self._client.poll_run,
            handle.batch_id,
            job_id_str,
            timeout=self.poll_timeout,
            poll_interval=self.poll_interval,
        )

        manifest_payload = await asyncio.to_thread(
            self._client.get_result,
            handle.batch_id,
            job_id_str,
        )

        local_root: Path | None = None
        if self.download_results and self.store_dir:
            local_root = await self._materialise_manifest(job_id_str, manifest_payload)

        status = status_payload.get("status", "unknown")
        job_total = int(status_payload.get("job_total", len(jobs)))
        job_completed = int(status_payload.get("job_completed", 0))
        job_failed = int(status_payload.get("job_failed", 0))

        logger.info(
            "Screenshot batch finished: batch=%s run=%s status=%s completed=%d failed=%d",
            handle.batch_id,
            job_id_str,
            status,
            job_completed,
            job_failed,
        )

        return ScreenshotBatchResult(
            batch_id=handle.batch_id,
            job_id=job_id_str,
            status=status,
            job_total=job_total,
            job_completed=job_completed,
            job_failed=job_failed,
            status_payload=status_payload,
            manifest_payload=manifest_payload,
            local_root=local_root,
        )

    def run_batch_sync(self, jobs: Sequence[ScreenshotJobSpec]) -> ScreenshotBatchResult:
        """Synchronous wrapper around `run_batch`."""

        return asyncio.run(self.run_batch(jobs))

    async def _materialise_manifest(self, run_id: str, manifest_payload: dict[str, Any]) -> Path | None:
        manifest = _unwrap_manifest(manifest_payload)
        storage_root_raw = manifest.get("storage_path")
        files = manifest.get("files")
        blob_prefix = manifest.get("blob_prefix")

        # Need either files with storage_path, or a blob_prefix to download from
        if (not storage_root_raw or not files) and not blob_prefix:
            return None

        storage_root = Path(storage_root_raw) if storage_root_raw else None
        run_root = self.store_dir / run_id if self.store_dir else None
        if run_root is None:
            return None

        copied_any = False
        if storage_root and files:
            for entry in files:
                rel_path = entry.get("path")
                if not rel_path:
                    continue
                target = storage_root / rel_path
                if not target.exists():
                    downloaded = await asyncio.to_thread(self._ensure_source_file, target, blob_prefix, rel_path)
                    if not downloaded or not target.exists():
                        continue
                dest_path = run_root / rel_path
                if dest_path.exists():
                    continue
                await asyncio.to_thread(_copy_file, target, dest_path)
                copied_any = True

        if not copied_any and blob_prefix:
            copied_any = await asyncio.to_thread(self._download_blob_prefix, blob_prefix, run_root)

        if not copied_any:
            container = self._source_container_name or "unspecified"
            location = blob_prefix or storage_root_raw or "(unknown prefix)"
            raise RuntimeError(
                "Failed to download screenshot artifacts from container "
                f"'{container}' at prefix '{location}' for run {run_id}."
            )

        return run_root

    def _ensure_source_file(self, target: Path, blob_prefix: str | None, rel_path: str) -> bool:
        if target.exists():
            return True

        blob_name = self._build_blob_name(blob_prefix, rel_path, target)
        if not blob_name:
            return False

        return self._download_from_screenshot_storage(blob_name, target)

    def _download_from_screenshot_storage(self, blob_name: str, target: Path) -> bool:
        """Download a file from storage backend."""
        return self._storage.download_file(blob_name, target)

    def _download_blob_prefix(self, blob_prefix: str, run_root: Path) -> bool:
        """Download all blobs with given prefix using storage backend."""
        return self._storage.download_prefix(blob_prefix, run_root)

    def _build_blob_name(self, manifest_prefix: str | None, rel_path: str, target: Path) -> str | None:
        """Build blob name from manifest prefix and relative path."""
        path_part = rel_path.lstrip("/")
        if manifest_prefix:
            prefix = manifest_prefix.rstrip("/")
            if prefix:
                return f"{prefix}/{path_part}"
            return path_part
        return path_part or None


@dataclass(frozen=True, kw_only=True, slots=True)
class ScreenshotRunHandle:
    """Persisted reference to an in-flight or completed screenshot run.

    Attributes:
        job_key: Unique key for state tracking.
        job_id: Batch job identifier.
        run_id: Individual run identifier.
        chunk_index: Index of this chunk in the overall sequence.
    """

    job_key: str
    job_id: str
    run_id: str
    chunk_index: int


@dataclass(frozen=True, kw_only=True, slots=True)
class ScreenshotRunSummary:
    """Summary data for an individual screenshot run executed by the multi-run coordinator.

    Attributes:
        chunk_index: Index of this chunk in the overall sequence.
        specs: Job specifications that were executed in this run.
        result: Batch execution result.
        uploaded_blobs: Mapping of local paths to uploaded blob URLs.
        job_results: Mapping of job IDs to their individual results.
    """

    chunk_index: int
    specs: tuple[ScreenshotJobSpec, ...]
    result: ScreenshotBatchResult
    uploaded_blobs: dict[str, str] = field(default_factory=dict)
    job_results: dict[str, JobRunSummary] = field(default_factory=dict)


class ScreenshotMultiRunCoordinator:
    """High-level batching helper that fans out multiple screenshot API calls.

    This coordinator manages multiple ScreenshotBatchCoordinator instances to
    process large job collections in parallel chunks.

    Args:
        base_url: Base URL of the screenshot service API.
        batch_id: Unique identifier for this batch.
        store_dir: Local directory for storing downloaded results.
        per_run_concurrency: Concurrent jobs per individual run.
        poll_interval: Seconds between status polls.
        poll_timeout: Maximum seconds to wait for batch completion.
        default_metadata: Default metadata merged into all jobs.
        download_results: Whether to download artifacts after completion.
        state_path: Path to persist run state for resumability.
        upload_results: Whether to upload results to storage.
        uploader: Custom uploader callable. If None and upload_results is True,
            creates Azure uploader from environment.
        storage: Optional storage backend. If None, creates from environment.
    """

    _uploader: Callable[[Path, str | None], None] | None

    def __init__(
        self,
        *,
        base_url: str,
        batch_id: str,
        store_dir: Path | None = None,
        per_run_concurrency: int = 4,
        poll_interval: float = 5.0,
        poll_timeout: float = 900.0,
        default_metadata: dict[str, Any] | None = None,
        download_results: bool = False,
        state_path: Path | None = None,
        upload_results: bool = False,
        uploader: Callable[[Path, str | None], None] | None = None,
        storage: ArtifactStorage | None = None,
    ) -> None:
        if not base_url:
            raise ValueError("base_url is required")
        if not batch_id:
            raise ValueError("batch_id is required")

        self.base_url = base_url
        self.batch_id = batch_id
        self.store_dir = Path(store_dir) if store_dir else None
        self.per_run_concurrency = max(1, int(per_run_concurrency))
        self.poll_interval = max(0.5, float(poll_interval))
        self.poll_timeout = max(self.poll_interval, float(poll_timeout))
        self.default_metadata = dict(default_metadata or {})
        self.download_results = bool(download_results and self.store_dir)
        self.state_path = Path(state_path) if state_path else None
        self._handles: dict[str, ScreenshotRunHandle] = {}
        self.upload_results = bool(upload_results)

        # Set up storage backend
        if storage:
            self._storage = storage
        elif upload_results or download_results:
            # Auto-create Azure storage from environment if needed
            azure_storage = create_azure_storage()
            self._storage = azure_storage if azure_storage else NullStorage()
        else:
            self._storage = NullStorage()

        # Set up uploader
        if uploader:
            self._uploader = uploader
        elif upload_results:
            # Create uploader from storage backend
            self._uploader = self._create_uploader_from_storage()
        else:
            self._uploader = None

        if self.download_results and self.store_dir:
            self.store_dir.mkdir(parents=True, exist_ok=True)

    def _create_uploader_from_storage(self) -> Callable[[Path, str | None], None] | None:
        """Create uploader callable from storage backend."""
        storage = self._storage
        if isinstance(storage, NullStorage):
            return None

        def uploader(path: Path, blob_path: str | None = None) -> None:
            if blob_path:
                storage.upload_file(path, blob_path)
            else:
                # Use path as blob name if no explicit blob_path
                storage.upload_file(path, str(path))

        return uploader

    def run_jobs(
        self,
        jobs: Sequence[ScreenshotJobSpec] | Iterable[ScreenshotJobSpec],
        *,
        chunk_size: int = 1,
        max_parallel_runs: int = 1,
        client_factory: Callable[[str], ScreenshotServiceClient] | None = None,
        progress_callback: Callable[[int, int, ScreenshotRunSummary], None] | None = None,
    ) -> list[ScreenshotRunSummary]:
        """Submit jobs in chunks, optionally parallelising across API calls."""

        chunk_size = max(1, int(chunk_size))
        max_parallel_runs = max(1, int(max_parallel_runs))

        job_specs = list(jobs)
        if not job_specs:
            return []

        self._load_state()

        summaries: list[ScreenshotRunSummary] = []
        chunked_specs = list(_chunk(job_specs, chunk_size))

        processed_chunks = 0
        total_chunks = len(chunked_specs)

        with ThreadPoolExecutor(max_workers=max_parallel_runs) as executor:
            future_map = {
                executor.submit(
                    self._run_chunk,
                    chunk_index,
                    tuple(specs),
                    client_factory,
                ): chunk_index
                for chunk_index, specs in enumerate(chunked_specs, start=1)
            }

            for future in as_completed(future_map):
                future_map.pop(future, None)
                result = future.result()
                summaries.append(result)
                processed_chunks += 1
                if progress_callback:
                    try:
                        progress_callback(processed_chunks, total_chunks, result)
                    except Exception:  # pragma: no cover - defensive logging
                        logger.debug("Progress callback failed", exc_info=True)

        self._persist_state()
        summaries.sort(key=lambda summary: summary.chunk_index)
        return summaries

    def _run_chunk(
        self,
        chunk_index: int,
        specs: tuple[ScreenshotJobSpec, ...],
        client_factory: Callable[[str], ScreenshotServiceClient] | None,
    ) -> ScreenshotRunSummary:
        job_key = self._build_chunk_key(chunk_index, specs)
        handle = self._handles.get(job_key)

        client = client_factory(self.base_url) if client_factory else ScreenshotServiceClient(self.base_url)

        coordinator = ScreenshotBatchCoordinator(
            base_url=self.base_url,
            batch_id=self.batch_id,
            store_dir=self._resolve_chunk_store_dir(chunk_index),
            concurrency=min(self.per_run_concurrency, len(specs)),
            poll_interval=self.poll_interval,
            poll_timeout=self.poll_timeout,
            result_metadata=self.default_metadata,
            download_results=self.download_results,
            client=client,
            storage=self._storage,
        )

        uploaded_blobs: dict[str, str] = {}

        if handle is None:
            result = coordinator.run_batch_sync(specs)
            run_id = result.job_id
        else:
            status_payload = client.poll_run(
                self.batch_id,
                handle.run_id,
                timeout=self.poll_timeout,
                poll_interval=self.poll_interval,
            )
            manifest_payload = client.get_result(self.batch_id, handle.run_id)
            local_root: Path | None = None
            if self.download_results and self.store_dir:
                local_root = asyncio.run(coordinator._materialise_manifest(handle.run_id, manifest_payload))
            status = status_payload.get("status", "unknown")
            job_total = int(status_payload.get("job_total", len(specs)))
            job_completed = int(status_payload.get("job_completed", 0))
            job_failed = int(status_payload.get("job_failed", 0))
            result = ScreenshotBatchResult(
                batch_id=self.batch_id,
                job_id=handle.run_id,
                status=status,
                job_total=job_total,
                job_completed=job_completed,
                job_failed=job_failed,
                status_payload=status_payload,
                manifest_payload=manifest_payload,
                local_root=local_root,
            )
            run_id = handle.run_id

        local_root_path = Path(result.local_root) if result.local_root else None
        job_results = self._collect_job_results(specs, result, local_root_path)

        if self.upload_results and self._uploader is not None and local_root_path:
            uploaded_blobs = self._upload_run_artifacts(local_root_path, job_results)
            for job_id, url in uploaded_blobs.items():
                summary = job_results.get(job_id)
                if summary is None:
                    continue
                if summary.status == "succeeded":
                    job_results[job_id] = replace(summary, blob_url=url)

        self._handles[job_key] = ScreenshotRunHandle(
            job_key=job_key,
            job_id=result.job_id,
            run_id=run_id,
            chunk_index=chunk_index,
        )

        return ScreenshotRunSummary(
            chunk_index=chunk_index,
            specs=specs,
            result=result,
            uploaded_blobs=uploaded_blobs,
            job_results=job_results,
        )

    def _resolve_chunk_store_dir(self, chunk_index: int) -> Path | None:
        if not self.store_dir:
            return None
        run_dir = self.store_dir / f"chunk-{chunk_index:05d}"
        run_dir.mkdir(parents=True, exist_ok=True)
        return run_dir

    def _collect_job_results(
        self,
        specs: tuple[ScreenshotJobSpec, ...],
        result: ScreenshotBatchResult,
        local_root: Path | None,
    ) -> dict[str, JobRunSummary]:
        outcomes: dict[str, JobRunSummary] = {}

        default_status = (
            "failed"
            if result.job_failed
            else ("succeeded" if is_batch_successful(result) else result.status or "unknown")
        )

        run_has_one_job = len(specs) == 1

        local_root_path = local_root if local_root and local_root.exists() else None

        if local_root_path:
            for spec in specs:
                job_id = spec.job_id
                if not job_id:
                    continue
                slug = derive_job_id(job_id or spec.url, fallback="job")
                candidates = [
                    local_root_path / job_id,
                    local_root_path / "screens" / job_id,
                    local_root_path / slug,
                    local_root_path / "screens" / slug,
                ]
                job_dir: Path | None = None
                local_path: Path | None = None
                for candidate in candidates:
                    if candidate.exists():
                        job_dir = candidate
                        break
                if job_dir and job_dir.exists():
                    status, errors = _read_job_metadata(job_dir)
                    local_path = job_dir
                else:
                    status, errors = (default_status if run_has_one_job else None), ()
                    local_path = local_root_path if run_has_one_job else None
                status = status or default_status
                outcomes[job_id] = JobRunSummary(
                    job_id=job_id,
                    status=status,
                    errors=errors,
                    local_path=local_path,
                )
        else:
            for spec in specs:
                job_id = spec.job_id
                if not job_id:
                    continue
                outcomes[job_id] = JobRunSummary(job_id=job_id, status=default_status, local_path=None)

        return outcomes

    def _upload_run_artifacts(self, run_root: Path, job_results: dict[str, JobRunSummary]) -> dict[str, str]:
        if not run_root.exists() or not self.upload_results or self._uploader is None:
            return {}

        uploaded: dict[str, str] = {}

        for job_id, summary in job_results.items():
            job_root = summary.local_path if summary.local_path and summary.local_path.exists() else None
            if job_root is None:
                if len(job_results) == 1:
                    job_root = run_root
                else:
                    continue

            prefix = f"{self.batch_id}/{job_id}"
            uploaded_any = False

            for path in job_root.rglob("*"):
                if not path.is_file():
                    continue
                relative = path.relative_to(job_root)
                blob_path = f"{prefix}/{relative.as_posix()}"
                try:
                    self._uploader(path, blob_path)
                except Exception:  # pragma: no cover - defensive logging
                    try:
                        self._uploader(path, None)
                    except Exception:
                        logger.debug("Failed to upload %s via uploader", path, exc_info=True)
                        continue
                uploaded_any = True

            if uploaded_any:
                blob_url = _build_blob_folder_url(self.batch_id, job_id)
                if blob_url:
                    uploaded[job_id] = blob_url

        shutil.rmtree(run_root, ignore_errors=True)
        return uploaded

    def _build_chunk_key(self, chunk_index: int, specs: tuple[ScreenshotJobSpec, ...]) -> str:
        if len(specs) == 1:
            return specs[0].job_id
        return f"chunk-{chunk_index:05d}"

    def _load_state(self) -> None:
        if not self.state_path or self._handles:
            return
        try:
            raw = self.state_path.read_text(encoding="utf-8")
        except FileNotFoundError:
            return
        except OSError:
            logger.debug(
                "Unable to read screenshot coordinator state at %s",
                self.state_path,
                exc_info=True,
            )
            return
        try:
            payload = json.loads(raw) if raw.strip() else {}
        except json.JSONDecodeError:
            logger.warning(
                "Screenshot coordinator state %s is not valid JSON; ignoring",
                self.state_path,
            )
            return
        if not isinstance(payload, dict):
            return
        handles: dict[str, ScreenshotRunHandle] = {}
        for key, entry in payload.items():
            if not isinstance(entry, dict):
                continue
            job_id = entry.get("job_id")
            run_id = entry.get("run_id") or job_id
            chunk_index = entry.get("chunk_index")
            if not job_id or chunk_index is None:
                continue
            handles[key] = ScreenshotRunHandle(
                job_key=key,
                job_id=str(job_id),
                run_id=str(run_id),
                chunk_index=int(chunk_index),
            )
        self._handles = handles

    def _persist_state(self) -> None:
        if not self.state_path:
            return
        payload = {
            key: {
                "job_id": handle.job_id,
                "run_id": handle.run_id,
                "chunk_index": handle.chunk_index,
            }
            for key, handle in self._handles.items()
        }
        self.state_path.parent.mkdir(parents=True, exist_ok=True)
        self.state_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def _chunk(items: Sequence[ScreenshotJobSpec], size: int) -> Iterable[tuple[ScreenshotJobSpec, ...]]:
    chunk: list[ScreenshotJobSpec] = []
    for item in items:
        chunk.append(item)
        if len(chunk) == size:
            yield tuple(chunk)
            chunk = []
    if chunk:
        yield tuple(chunk)


def _read_job_metadata(job_dir: Path) -> tuple[str | None, tuple[str, ...]]:
    metadata_path = job_dir / "metadata" / "external_screenshots.json"
    if not metadata_path.exists():
        fallback = job_dir / "metadata" / "screenshots.json"
        metadata_path = fallback if fallback.exists() else metadata_path
    if not metadata_path.exists():
        return None, ()
    try:
        payload = json.loads(metadata_path.read_text(encoding="utf-8"))
    except Exception:  # pragma: no cover - defensive
        logger.debug("Failed to read metadata at %s", metadata_path, exc_info=True)
        return None, ()

    entries = payload.get("entries") if isinstance(payload, dict) else None
    errors_field = payload.get("errors") if isinstance(payload, dict) else None

    errors: list[str] = []
    if isinstance(errors_field, list):
        errors.extend(str(item) for item in errors_field if item)

    status: str | None = None
    if isinstance(entries, list):
        has_capture = False
        has_failure = False
        for entry in entries:
            if not isinstance(entry, dict):
                continue
            entry_status = str(entry.get("status") or "").lower()
            if entry_status == "captured":
                has_capture = True
            elif entry_status == "failed":
                has_failure = True
                message = entry.get("error") or entry.get("message")
                if message:
                    errors.append(str(message))
        if has_failure:
            status = "failed"
        elif has_capture:
            status = "succeeded"

    if status is None:
        status = "failed" if errors else None

    return status, tuple(errors)


def _build_blob_folder_url(batch_id: str, job_key: str) -> str | None:
    """Build blob folder URL from environment configuration.

    Note: This function reads from environment variables directly.
    Consider using storage backend's URL construction instead.
    """
    import os

    container = os.getenv("SCREENSHOT_AZURE_STORAGE_CONTAINER") or "screenshot-output"
    connection = os.getenv("SCREENSHOT_AZURE_STORAGE_CONNECTION_STRING") or ""
    endpoint = os.getenv("AZURE_STORAGE_BLOB_ENDPOINT") or ""
    account = os.getenv("AZURE_STORAGE_ACCOUNT") or ""

    base_url: str | None = None
    if endpoint:
        base_url = endpoint.rstrip("/")
    elif account:
        base_url = f"https://{account}.blob.core.windows.net"
    elif connection:
        # Parse connection string for account name
        parts = dict(segment.split("=", 1) for segment in connection.split(";") if "=" in segment)
        account_name = parts.get("AccountName", "").strip()
        if account_name:
            base_url = f"https://{account_name}.blob.core.windows.net"

    if not base_url or not container:
        return None

    path = f"{batch_id}/{job_key}".strip("/")
    return f"{base_url}/{container}/{path}"


def _unwrap_manifest(payload: dict[str, Any]) -> dict[str, Any]:
    manifest = payload.get("manifest")
    if isinstance(manifest, dict):
        return manifest
    return payload
