from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

from backend.services.artifact_models import ArtifactReference, PreparedToolSourceArtifact
from backend.services.artifact_read_service.inline_images import (
    extract_inline_image_candidates,
    filter_sendable_inline_image_parts,
)


@dataclass(frozen=True)
class SourceImageAsset:
    """Document-local image binding used across prompting, validation, rendering, and revision."""

    asset_id: str
    source_reference: ArtifactReference
    image_part_index: int
    mime_type: str
    data_base64: str
    artifact_name: str = ""

    def as_generation_metadata(self) -> dict[str, Any]:
        """Serialize the stable binding needed to re-resolve this image during a later revision."""

        return {
            "messageId": self.source_reference.message_id,
            "artifactsId": self.source_reference.artifacts_id,
            "artifactIndex": self.source_reference.artifact_index,
            "imagePartIndex": self.image_part_index,
            "imageMimeType": self.mime_type,
            **({"artifactName": self.artifact_name} if self.artifact_name else {}),
        }


def build_image_assets_from_source_artifacts(
    source_artifacts: list[PreparedToolSourceArtifact],
) -> list[SourceImageAsset]:
    """Flatten selected image artifacts into an ordered document-local asset registry."""

    assets: list[SourceImageAsset] = []
    for source_artifact in source_artifacts:
        sendable_index = 0
        for part in source_artifact.parts:
            if part.get("type") != "IMAGE_INLINE":
                continue
            inline_image = part.get("inlineImage")
            image_mime_type = part.get("imageMimeType")
            if not isinstance(inline_image, str) or not isinstance(image_mime_type, str):
                continue
            assets.append(
                SourceImageAsset(
                    asset_id=f"img_{len(assets) + 1}",
                    source_reference=source_artifact.reference,
                    image_part_index=sendable_index,
                    mime_type=image_mime_type,
                    data_base64=inline_image,
                    artifact_name=source_artifact.artifact_name,
                )
            )
            sendable_index += 1
    return assets


def build_image_assets_map(image_assets: list[SourceImageAsset]) -> dict[str, SourceImageAsset]:
    return {asset.asset_id: asset for asset in image_assets}


def build_generation_image_assets_metadata(
    image_assets: list[SourceImageAsset],
) -> dict[str, dict[str, Any]]:
    return {asset.asset_id: asset.as_generation_metadata() for asset in image_assets}


def merge_image_assets(
    existing_assets: list[SourceImageAsset],
    additional_assets: list[SourceImageAsset],
) -> list[SourceImageAsset]:
    """Append new assets while preserving the identifiers already used by a generated document."""

    if not additional_assets:
        return list(existing_assets)

    next_asset_number = max((_parse_asset_number(asset.asset_id) for asset in existing_assets), default=0)
    merged_assets = list(existing_assets)
    for asset in additional_assets:
        next_asset_number += 1
        merged_assets.append(
            SourceImageAsset(
                asset_id=f"img_{next_asset_number}",
                source_reference=asset.source_reference,
                image_part_index=asset.image_part_index,
                mime_type=asset.mime_type,
                data_base64=asset.data_base64,
                artifact_name=asset.artifact_name,
            )
        )
    return merged_assets


def resolve_image_asset_from_artifact(
    *,
    asset_id: str,
    metadata: dict[str, Any],
    source_artifact: dict[str, Any],
) -> SourceImageAsset:
    """Rebuild one image binding from retained metadata and its original persisted artifact."""

    image_parts = extract_inline_image_candidates(source_artifact)
    sendable_parts, _skipped_parts = filter_sendable_inline_image_parts(image_parts)

    image_part_index = metadata.get("imagePartIndex")
    if not isinstance(image_part_index, int) or image_part_index < 0:
        raise ValueError(f"Image asset '{asset_id}' is missing a valid imagePartIndex")
    if image_part_index >= len(sendable_parts):
        raise ValueError(f"Image asset '{asset_id}' references an image part that is no longer available")

    message_id = metadata.get("messageId")
    artifacts_id = metadata.get("artifactsId")
    artifact_index = metadata.get("artifactIndex")
    image_mime_type = metadata.get("imageMimeType")
    if not isinstance(message_id, str) or not isinstance(artifacts_id, str) or not isinstance(artifact_index, int):
        raise ValueError(f"Image asset '{asset_id}' is missing source artifact reference metadata")
    if not isinstance(image_mime_type, str) or not image_mime_type.strip():
        raise ValueError(f"Image asset '{asset_id}' is missing image MIME metadata")

    part = sendable_parts[image_part_index]
    return SourceImageAsset(
        asset_id=asset_id,
        source_reference=ArtifactReference(
            message_id=message_id,
            artifacts_id=artifacts_id,
            artifact_index=artifact_index,
        ),
        image_part_index=image_part_index,
        mime_type=part["mimeType"],
        data_base64=part["dataBase64"],
        artifact_name=str(metadata.get("artifactName") or ""),
    )


def _parse_asset_number(asset_id: str) -> int:
    match = re.fullmatch(r"img_(\d+)", asset_id.strip())
    return int(match.group(1)) if match else 0
