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_PROPOSED,
    REVISION_STATUS_ACCEPTED,
    REVISION_STATUS_PENDING,
    REVISION_STATUS_SUPERSEDED,
    copy_artifacts_metadata,
    detect_generated_document_format,
    extract_generated_document_source_yaml_from_artifact,
    get_document_metadata,
    iter_document_artifacts,
    make_history_marker,
    message_artifacts_metadata,
    refresh_artifacts_group_metadata,
)


class StaleGeneratedDocumentRevisionError(Exception):
    pass


@dataclass(frozen=True)
class GeneratedDocumentAppendTransition:
    existing_artifact_updates: dict[str, dict[str, Any]]


@dataclass(frozen=True)
class _ProposedArtifact:
    artifact: dict[str, Any]
    document_metadata: dict[str, Any]


@dataclass(frozen=True)
class _ProposalSource:
    lineage_id: str
    message_id: str
    artifacts_id: str
    artifact_index: int
    expected_lineage_role: Any
    expected_revision_status: Any


@dataclass(frozen=True)
class _ExistingProposal:
    artifacts_metadata: dict[str, Any]
    artifacts_id: str
    artifact: dict[str, Any]


@dataclass(frozen=True)
class _ProposalTransitionPlan:
    new_document_metadata: dict[str, Any]
    existing_proposals: tuple[_ExistingProposal, ...]


def append_messages_with_generated_document_transition(
    *,
    store: Any,
    conversation_id: str,
    existing_messages: list[Any],
    messages: list[Any],
    assistant_message: Any,
) -> None:
    transition = prepare_generated_document_append_transition(
        messages=existing_messages,
        new_artifacts_metadata=assistant_message.artifacts,
    )
    if transition is None or not transition.existing_artifact_updates:
        store.append_messages(conversation_id, messages)
        return

    store.append_messages_with_artifact_updates(
        conversation_id,
        messages,
        transition.existing_artifact_updates,
    )


def prepare_generated_document_append_transition(
    *,
    messages: list[Any],
    new_artifacts_metadata: dict[str, Any] | None,
) -> GeneratedDocumentAppendTransition | None:
    proposed_artifacts = _collect_proposed_artifacts(new_artifacts_metadata)
    if not proposed_artifacts:
        return None

    _validate_unique_proposal_lineages(proposed_artifacts)

    original_artifacts_by_message_id = {
        getattr(message, "id", ""): message_artifacts_metadata(message) or {} for message in messages
    }
    working_artifacts_by_message_id = {
        message_id: copy_artifacts_metadata(artifacts_metadata)
        for message_id, artifacts_metadata in original_artifacts_by_message_id.items()
    }

    transition_plans = [
        _validate_proposal_and_plan_transition(proposal, working_artifacts_by_message_id)
        for proposal in proposed_artifacts
    ]
    _apply_proposal_transition_plans(transition_plans)

    existing_artifact_updates = {
        message_id: updated
        for message_id, updated in working_artifacts_by_message_id.items()
        if updated != original_artifacts_by_message_id[message_id]
    }
    return GeneratedDocumentAppendTransition(existing_artifact_updates=existing_artifact_updates)


def _collect_proposed_artifacts(new_artifacts_metadata: dict[str, Any] | None) -> list[_ProposedArtifact]:
    return [
        _ProposedArtifact(artifact=artifact, document_metadata=document_metadata)
        for _artifacts_id, _artifact_index, artifact, document_metadata in iter_document_artifacts(
            new_artifacts_metadata
        )
        if document_metadata.get("lineageRole") == LINEAGE_ROLE_PROPOSED
    ]


def _validate_unique_proposal_lineages(proposed_artifacts: list[_ProposedArtifact]) -> None:
    lineages = [proposal.document_metadata.get("lineageId") for proposal in proposed_artifacts]
    if len(lineages) != len(set(lineages)):
        raise StaleGeneratedDocumentRevisionError(
            "Only one Proposed generated document artifact can be created per lineage in a turn"
        )


def _parse_proposal_source(proposal: _ProposedArtifact) -> _ProposalSource:
    document_metadata = proposal.document_metadata
    lineage_id = document_metadata.get("lineageId")
    source_reference = document_metadata.get("source")
    expected_source_state = document_metadata.get("_sourceState")
    if not isinstance(lineage_id, str) or not lineage_id:
        raise StaleGeneratedDocumentRevisionError("The generated proposal is missing lineage metadata")
    if not isinstance(source_reference, dict) or not isinstance(expected_source_state, dict):
        raise StaleGeneratedDocumentRevisionError("The generated proposal is missing source commit metadata")

    source_message_id = source_reference.get("messageId")
    source_artifacts_id = source_reference.get("artifactsId")
    source_artifact_index = source_reference.get("artifactIndex")
    if (
        not isinstance(source_message_id, str)
        or not isinstance(source_artifacts_id, str)
        or not isinstance(source_artifact_index, int)
    ):
        raise StaleGeneratedDocumentRevisionError("The generated proposal has an invalid source reference")

    return _ProposalSource(
        lineage_id=lineage_id,
        message_id=source_message_id,
        artifacts_id=source_artifacts_id,
        artifact_index=source_artifact_index,
        expected_lineage_role=expected_source_state.get("lineageRole"),
        expected_revision_status=expected_source_state.get("revisionStatus"),
    )


def _validate_proposal_and_plan_transition(
    proposal: _ProposedArtifact,
    artifacts_by_message_id: dict[str, dict[str, Any]],
) -> _ProposalTransitionPlan:
    source = _parse_proposal_source(proposal)
    source_artifact = get_artifact_from_metadata(
        artifacts_meta=artifacts_by_message_id.get(source.message_id),
        artifacts_id=source.artifacts_id,
        artifact_index=source.artifact_index,
    )
    source_metadata = get_document_metadata(source_artifact)
    if source_metadata is None:
        raise StaleGeneratedDocumentRevisionError("The generated proposal source no longer exists")
    if source_metadata.get("lineageId") != source.lineage_id:
        raise StaleGeneratedDocumentRevisionError("The generated proposal source lineage changed")
    if (
        source_metadata.get("lineageRole") != source.expected_lineage_role
        or source_metadata.get("revisionStatus") != source.expected_revision_status
    ):
        raise StaleGeneratedDocumentRevisionError("The generated proposal source changed during generation")
    proposal_format = detect_generated_document_format(proposal.artifact)
    if proposal_format is None:
        raise StaleGeneratedDocumentRevisionError(
            "The generated proposal is missing a supported editable document source"
        )
    if extract_generated_document_source_yaml_from_artifact(source_artifact, proposal_format) is None:
        raise StaleGeneratedDocumentRevisionError(
            "The generated proposal source content is no longer retained in the proposal format"
        )

    current_count = 0
    existing_proposals = []
    for artifacts_metadata in artifacts_by_message_id.values():
        for artifacts_id, _artifact_index, artifact, document_metadata in iter_document_artifacts(artifacts_metadata):
            if document_metadata.get("lineageId") != source.lineage_id:
                continue
            lineage_role = document_metadata.get("lineageRole")
            if lineage_role == LINEAGE_ROLE_CURRENT:
                if document_metadata.get("revisionStatus") != REVISION_STATUS_ACCEPTED:
                    raise StaleGeneratedDocumentRevisionError("The Current generated document state is invalid")
                current_count += 1
            elif lineage_role == LINEAGE_ROLE_PROPOSED:
                if document_metadata.get("revisionStatus") != REVISION_STATUS_PENDING:
                    raise StaleGeneratedDocumentRevisionError("The Proposed generated document state is invalid")
                existing_proposals.append(
                    _ExistingProposal(
                        artifacts_metadata=artifacts_metadata,
                        artifacts_id=artifacts_id,
                        artifact=artifact,
                    )
                )

    if current_count != 1:
        raise StaleGeneratedDocumentRevisionError("The generated document lineage must have exactly one Current")

    return _ProposalTransitionPlan(
        new_document_metadata=proposal.document_metadata,
        existing_proposals=tuple(existing_proposals),
    )


def _apply_proposal_transition_plans(transition_plans: list[_ProposalTransitionPlan]) -> None:
    for plan in transition_plans:
        for existing_proposal in plan.existing_proposals:
            make_history_marker(existing_proposal.artifact, revision_status=REVISION_STATUS_SUPERSEDED)
            refresh_artifacts_group_metadata(existing_proposal.artifacts_metadata, existing_proposal.artifacts_id)

    for plan in transition_plans:
        plan.new_document_metadata.pop("_sourceState", None)
