from __future__ import annotations

from backend.services.artifact_models import ArtifactReference, ArtifactReferenceSpec
from backend.services.artifact_read_service import get_artifact_from_metadata
from backend.services.artifact_read_service.service import (
    CURRENT_TURN_ARTIFACTS_BY_MESSAGE_ID_CONTEXT_KEY,
)
from backend.services.generated_document.generated_document_metadata import (
    GENERATED_DOCUMENT_FORMAT_PPTX,
    LINEAGE_ROLE_CURRENT,
    LINEAGE_ROLE_HISTORY,
    LINEAGE_ROLE_PROPOSED,
    REVISION_STATUS_ACCEPTED,
    REVISION_STATUS_PENDING,
    GeneratedDocumentFormat,
    extract_generated_document_filename_from_artifact,
    extract_generated_document_image_assets_from_artifact,
    extract_generated_document_source_yaml_from_artifact,
    get_document_metadata,
)
from backend.services.generated_document.revision_models import GeneratedDocumentRevisionSource
from backend.services.tool_source_artifacts.image_assets import SourceImageAsset, resolve_image_asset_from_artifact


class RevisionSourceResolutionError(Exception):
    pass


def resolve_generated_document_revision_source(
    *,
    store,
    conversation_id: str,
    reference: ArtifactReferenceSpec,
    expected_format: GeneratedDocumentFormat,
    runtime_context: dict | None = None,
) -> GeneratedDocumentRevisionSource:
    """Resolve a revisable generated-document artifact into the editable source expected by revision tools."""

    if not conversation_id:
        raise RevisionSourceResolutionError("Conversation id is required to resolve the revision source")

    conversation = store.get_conversation(conversation_id)
    if not conversation:
        raise RevisionSourceResolutionError("Conversation not found")

    messages = list(conversation.messages or [])
    artifacts_meta = _resolve_message_artifacts_metadata(
        messages=messages,
        message_id=reference.message_id,
        runtime_context=runtime_context,
    )

    if artifacts_meta is None:
        raise RevisionSourceResolutionError("Revision source message was not found in the current conversation")

    artifact = get_artifact_from_metadata(
        artifacts_meta=artifacts_meta,
        artifacts_id=reference.artifacts_id,
        artifact_index=reference.artifact_index,
    )
    if artifact is None:
        raise RevisionSourceResolutionError("Revision source artifact reference was not found")

    document_metadata = get_document_metadata(artifact)
    if document_metadata is None:
        raise RevisionSourceResolutionError("The selected artifact is missing generated document metadata")

    lineage_id = document_metadata.get("lineageId")
    if not isinstance(lineage_id, str) or not lineage_id:
        raise RevisionSourceResolutionError("The selected artifact is missing generated document lineage metadata")

    lineage_role = document_metadata.get("lineageRole")
    if lineage_role == LINEAGE_ROLE_HISTORY:
        raise RevisionSourceResolutionError("The selected artifact content is no longer retained")
    if lineage_role not in {LINEAGE_ROLE_CURRENT, LINEAGE_ROLE_PROPOSED}:
        raise RevisionSourceResolutionError("The selected artifact is missing generated document lineage role metadata")

    revision_status = document_metadata.get("revisionStatus")
    if (lineage_role, revision_status) not in {
        (LINEAGE_ROLE_CURRENT, REVISION_STATUS_ACCEPTED),
        (LINEAGE_ROLE_PROPOSED, REVISION_STATUS_PENDING),
    }:
        raise RevisionSourceResolutionError("The selected artifact has inconsistent generated document state")

    source_yaml = extract_generated_document_source_yaml_from_artifact(artifact, expected_format)
    if source_yaml is None:
        format_label = "PPTX" if expected_format == GENERATED_DOCUMENT_FORMAT_PPTX else "DOCX"
        raise RevisionSourceResolutionError(
            f"The selected artifact is not an Agent Hub generated {format_label} with editable source"
        )

    image_assets = _resolve_revision_image_assets(
        messages=messages,
        image_assets_metadata=extract_generated_document_image_assets_from_artifact(artifact, expected_format) or {},
        runtime_context=runtime_context,
    )

    source = ArtifactReference(
        message_id=reference.message_id,
        artifacts_id=reference.artifacts_id,
        artifact_index=reference.artifact_index,
    )
    return GeneratedDocumentRevisionSource(
        source_yaml=source_yaml,
        document_format=expected_format,
        filename=extract_generated_document_filename_from_artifact(artifact, expected_format),
        lineage_id=lineage_id,
        lineage_role=lineage_role,
        source=source,
        image_assets=image_assets,
    )


def _resolve_revision_image_assets(
    *,
    messages: list,
    image_assets_metadata: dict[str, dict],
    runtime_context: dict | None,
) -> list[SourceImageAsset]:
    """Re-resolve each persisted image binding against the source messages still present in the conversation."""

    resolved_assets: list[SourceImageAsset] = []
    for asset_id, metadata in image_assets_metadata.items():
        message_id = metadata.get("messageId")
        artifacts_id = metadata.get("artifactsId")
        artifact_index = metadata.get("artifactIndex")
        if not isinstance(message_id, str) or not isinstance(artifacts_id, str) or not isinstance(artifact_index, int):
            raise RevisionSourceResolutionError(f"Image asset '{asset_id}' is missing source artifact metadata")

        source_artifacts_meta = _resolve_message_artifacts_metadata(
            messages=messages,
            message_id=message_id,
            runtime_context=runtime_context,
        )
        if source_artifacts_meta is None:
            raise RevisionSourceResolutionError(f"Image asset '{asset_id}' source message was not found")

        source_artifact = get_artifact_from_metadata(
            artifacts_meta=source_artifacts_meta,
            artifacts_id=artifacts_id,
            artifact_index=artifact_index,
        )
        if source_artifact is None:
            raise RevisionSourceResolutionError(f"Image asset '{asset_id}' source artifact was not found")
        try:
            resolved_assets.append(
                resolve_image_asset_from_artifact(
                    asset_id=asset_id,
                    metadata=metadata,
                    source_artifact=source_artifact,
                )
            )
        except ValueError as exc:
            raise RevisionSourceResolutionError(str(exc)) from exc
    return resolved_assets


def _resolve_message_artifacts_metadata(
    *,
    messages: list,
    message_id: str,
    runtime_context: dict | None,
) -> dict | None:
    source_message = next((message for message in messages if getattr(message, "id", None) == message_id), None)
    artifacts_meta = getattr(source_message, "artifacts_metadata", None) if source_message is not None else None
    if artifacts_meta is not None:
        return artifacts_meta

    if not isinstance(runtime_context, dict):
        return None

    # revise_docx/revise_pptx resolve artifacts directly, so they need the same same-turn fallback as read_artifact.
    current_turn_artifacts = runtime_context.get(CURRENT_TURN_ARTIFACTS_BY_MESSAGE_ID_CONTEXT_KEY)
    if not isinstance(current_turn_artifacts, dict):
        return None

    resolved_artifacts_meta = current_turn_artifacts.get(message_id)
    return resolved_artifacts_meta if isinstance(resolved_artifacts_meta, dict) else None
