from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from backend.services.artifact_read_service import get_artifact_from_metadata
from backend.services.generated_document.generated_document_metadata import (
    LINEAGE_ROLE_CURRENT,
    LINEAGE_ROLE_HISTORY,
    LINEAGE_ROLE_PROPOSED,
    REVISION_STATUS_ACCEPTED,
    REVISION_STATUS_PENDING,
    REVISION_STATUS_REJECTED,
    REVISION_STATUS_SUPERSEDED,
    copy_artifacts_metadata,
    get_document_metadata,
    iter_document_artifacts,
    make_history_marker,
    refresh_artifacts_group_metadata,
)


@dataclass(frozen=True)
class RevisionReviewError(Exception):
    message: str
    status_code: int


def review_generated_document_revision(
    *,
    store,
    conversation_id: str,
    message_id: str,
    artifacts_id: str,
    artifact_index: int,
    review_status: str,
) -> list[dict[str, Any]]:
    if review_status not in {REVISION_STATUS_ACCEPTED, REVISION_STATUS_REJECTED}:
        raise RevisionReviewError("Invalid review status", 400)

    conversation = store.get_conversation(conversation_id)
    if not conversation:
        raise RevisionReviewError("Conversation not found", 404)

    messages = list(conversation.messages or [])
    target_message = next((message for message in messages if getattr(message, "id", None) == message_id), None)
    if target_message is None:
        raise RevisionReviewError("Message not found", 404)

    original_artifacts_by_message_id = {
        getattr(message, "id", ""): message.artifacts_metadata or {} for message in messages
    }
    artifacts_by_message_id = {
        message_id: copy_artifacts_metadata(artifacts_metadata)
        for message_id, artifacts_metadata in original_artifacts_by_message_id.items()
    }
    target_artifacts_metadata = artifacts_by_message_id[message_id]
    target_artifact = get_artifact_from_metadata(
        artifacts_meta=target_artifacts_metadata,
        artifacts_id=artifacts_id,
        artifact_index=artifact_index,
    )
    if target_artifact is None:
        raise RevisionReviewError("Artifact reference not found", 404)

    target_document_metadata = get_document_metadata(target_artifact)
    if target_document_metadata is None or not isinstance(target_document_metadata.get("source"), dict):
        raise RevisionReviewError("Artifact is not a generated document revision proposal", 400)

    lineage_id = target_document_metadata.get("lineageId")
    if not isinstance(lineage_id, str) or not lineage_id:
        raise RevisionReviewError("Generated document lineage is missing", 400)

    current_status = target_document_metadata.get("revisionStatus")
    current_lineage_role = target_document_metadata.get("lineageRole")
    same_accepted_outcome = (
        review_status == REVISION_STATUS_ACCEPTED
        and current_status == REVISION_STATUS_ACCEPTED
        and current_lineage_role in {LINEAGE_ROLE_CURRENT, LINEAGE_ROLE_HISTORY}
    )
    same_rejected_outcome = (
        review_status == REVISION_STATUS_REJECTED
        and current_status == REVISION_STATUS_REJECTED
        and current_lineage_role == LINEAGE_ROLE_HISTORY
    )
    if same_accepted_outcome or same_rejected_outcome:
        return _lineage_message_snapshots(messages, artifacts_by_message_id, lineage_id)
    if current_status != REVISION_STATUS_PENDING:
        raise RevisionReviewError("The revision proposal was already resolved with a different outcome", 409)
    if target_document_metadata.get("lineageRole") != LINEAGE_ROLE_PROPOSED:
        raise RevisionReviewError("Only the Proposed generated document artifact can be reviewed", 409)

    if review_status == REVISION_STATUS_REJECTED:
        make_history_marker(target_artifact, revision_status=REVISION_STATUS_REJECTED)
        refresh_artifacts_group_metadata(target_artifacts_metadata, artifacts_id)
    else:
        _accept_revision_proposal(
            artifacts_by_message_id=artifacts_by_message_id,
            target_message_id=message_id,
            target_artifacts_id=artifacts_id,
            target_artifact_index=artifact_index,
            lineage_id=lineage_id,
        )

    changed_artifacts_by_message_id: dict[str, dict[str, Any]] = {}
    for message in messages:
        msg_id = getattr(message, "id", "")
        original = message.artifacts_metadata or {}
        updated = artifacts_by_message_id.get(msg_id) or {}
        if updated != original:
            changed_artifacts_by_message_id[msg_id] = updated

    if changed_artifacts_by_message_id:
        store.update_message_artifacts_batch(conversation_id, changed_artifacts_by_message_id)

    return _lineage_message_snapshots(messages, artifacts_by_message_id, lineage_id)


def _accept_revision_proposal(
    *,
    artifacts_by_message_id: dict[str, dict[str, Any]],
    target_message_id: str,
    target_artifacts_id: str,
    target_artifact_index: int,
    lineage_id: str,
) -> None:
    target_count = 0
    previous_current_count = 0
    for msg_id, artifacts_metadata in artifacts_by_message_id.items():
        for artifacts_id, artifact_index, artifact, document_metadata in iter_document_artifacts(artifacts_metadata):
            if document_metadata.get("lineageId") != lineage_id:
                continue

            is_target = (
                msg_id == target_message_id
                and artifacts_id == target_artifacts_id
                and artifact_index == target_artifact_index
            )
            if is_target:
                document_metadata["lineageRole"] = LINEAGE_ROLE_CURRENT
                document_metadata["revisionStatus"] = REVISION_STATUS_ACCEPTED
                target_count += 1
                continue

            lineage_role = document_metadata.get("lineageRole")
            if lineage_role == LINEAGE_ROLE_CURRENT:
                previous_current_count += 1
                make_history_marker(artifact, revision_status=REVISION_STATUS_ACCEPTED)
                refresh_artifacts_group_metadata(artifacts_metadata, artifacts_id)
            elif lineage_role == LINEAGE_ROLE_PROPOSED:
                make_history_marker(artifact, revision_status=REVISION_STATUS_SUPERSEDED)
                refresh_artifacts_group_metadata(artifacts_metadata, artifacts_id)

    if target_count != 1 or previous_current_count != 1:
        raise RevisionReviewError("Generated document lineage has an invalid Current artifact state", 409)


def _lineage_message_snapshots(
    messages: list[Any],
    artifacts_by_message_id: dict[str, dict[str, Any]],
    lineage_id: str,
) -> list[dict[str, Any]]:
    snapshots: list[dict[str, Any]] = []
    for message in messages:
        message_id = getattr(message, "id", "")
        artifacts_metadata = artifacts_by_message_id.get(message_id) or {}
        if not any(
            document_metadata.get("lineageId") == lineage_id
            for _artifacts_id, _artifact_index, _artifact, document_metadata in iter_document_artifacts(
                artifacts_metadata
            )
        ):
            continue
        snapshots.append({"messageId": message_id, "artifactsMetadata": artifacts_metadata})
    return snapshots
