from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass

from pydantic import ValidationError

from backend.models.events import EventKind
from backend.services.generated_document.generated_document_metadata import build_initial_document_metadata
from backend.services.tool_source_artifacts.image_assets import (
    SourceImageAsset,
    build_generation_image_assets_metadata,
    build_image_assets_map,
)

from . import generation_steps, internal_utils, quality_gate, rendering
from .document_models import DocumentSpec
from .parsing import DocxDslContentError, DocxDslSyntaxError, format_document_spec_for_prompt, parse_docx_spec
from .source_artifacts import build_docx_image_assets_user_message


class DocxImageAssetReferenceError(ValueError):
    pass


_RETRYABLE_ERRORS = (
    DocxDslContentError,
    DocxDslSyntaxError,
    DocxImageAssetReferenceError,
    ValidationError,
    quality_gate.DocumentQualityFailure,
)
_REVISION_RETRYABLE_ERRORS = (
    DocxDslContentError,
    DocxDslSyntaxError,
    DocxImageAssetReferenceError,
    ValidationError,
)


@dataclass(frozen=True)
class GeneratedDocxArtifacts:
    filename: str
    artifacts: list[dict]
    raw_yaml: str


@dataclass
class _DocxGenerationState:
    project: object
    llm_id: str
    user_prompt: str
    image_assets: list[SourceImageAsset]
    user_message: dict[str, object] | None
    raw_yaml: str
    emit_progress_event: Callable[[EventKind], None] | None
    repair_progress_emitted: bool = False


@dataclass
class _DocxRevisionState:
    project: object
    llm_id: str
    source_document_yaml: str
    revision_instructions: str
    image_assets: list[SourceImageAsset]
    user_message: dict[str, object] | None
    raw_yaml: str
    emit_progress_event: Callable[[EventKind], None] | None
    repair_progress_emitted: bool = False


def generate_docx_artifacts(
    project,
    llm_id: str,
    user_prompt: str,
    user_message: dict[str, object] | None = None,
    image_assets: list[SourceImageAsset] | None = None,
    requested_filename: str | None = None,
    emit_progress_event: Callable[[EventKind], None] | None = None,
    conversation_history: list[dict] | None = None,
) -> GeneratedDocxArtifacts:
    state = _initialize_generation_state(
        project=project,
        llm_id=llm_id,
        user_prompt=user_prompt,
        user_message=user_message,
        image_assets=image_assets or [],
        emit_progress_event=emit_progress_event,
        conversation_history=conversation_history,
    )
    spec = _generate_validated_document_spec(state)
    if emit_progress_event is not None:
        emit_progress_event(EventKind.DOCX_RENDERING_FILE)
    return _build_generated_docx_artifacts(
        spec=spec,
        image_assets=state.image_assets,
        user_prompt=user_prompt,
        requested_filename=requested_filename,
        document_metadata=build_initial_document_metadata(),
    )


def revise_docx_artifacts(
    project,
    llm_id: str,
    source_document_yaml: str,
    revision_instructions: str,
    source_image_assets: list[SourceImageAsset] | None = None,
    user_message: dict[str, object] | None = None,
    source_filename: str | None = None,
    requested_filename: str | None = None,
    document_metadata: dict | None = None,
    emit_progress_event: Callable[[EventKind], None] | None = None,
    conversation_history: list[dict] | None = None,
) -> GeneratedDocxArtifacts:
    """Revise retained Agent Hub DOCX YAML and package a newly rendered Proposed artifact."""

    source_spec = parse_docx_spec(source_document_yaml)
    normalized_source_document_yaml = format_document_spec_for_prompt(source_spec)
    image_assets = source_image_assets or []
    revision_user_message = user_message
    if revision_user_message is None and image_assets:
        revision_user_message = build_docx_image_assets_user_message(revision_instructions, image_assets)
    if emit_progress_event is not None:
        emit_progress_event(EventKind.DOCX_GENERATING_DOCUMENT)
    raw_yaml = generation_steps.revise_document_yaml(
        project=project,
        llm_id=llm_id,
        source_document_yaml=normalized_source_document_yaml,
        revision_instructions=revision_instructions,
        user_message=revision_user_message,
        conversation_history=conversation_history,
    )
    state = _DocxRevisionState(
        project=project,
        llm_id=llm_id,
        source_document_yaml=normalized_source_document_yaml,
        revision_instructions=revision_instructions,
        image_assets=image_assets,
        user_message=revision_user_message,
        raw_yaml=raw_yaml,
        emit_progress_event=emit_progress_event,
    )
    spec = _generate_validated_revised_document_spec(state)
    if emit_progress_event is not None:
        emit_progress_event(EventKind.DOCX_RENDERING_FILE)
    return _build_generated_docx_artifacts(
        spec=spec,
        image_assets=state.image_assets,
        user_prompt=revision_instructions,
        requested_filename=requested_filename or source_filename,
        document_metadata=document_metadata,
    )


def _initialize_generation_state(
    *,
    project,
    llm_id: str,
    user_prompt: str,
    user_message: dict[str, object] | None,
    image_assets: list[SourceImageAsset],
    emit_progress_event: Callable[[EventKind], None] | None,
    conversation_history: list[dict] | None = None,
) -> _DocxGenerationState:
    if emit_progress_event is not None:
        emit_progress_event(EventKind.DOCX_GENERATING_DOCUMENT)
    raw_yaml = generation_steps.generate_document_yaml(
        project,
        llm_id,
        user_prompt,
        user_message=user_message,
        conversation_history=conversation_history,
    )
    return _DocxGenerationState(
        project=project,
        llm_id=llm_id,
        user_prompt=user_prompt,
        image_assets=image_assets,
        user_message=user_message,
        raw_yaml=raw_yaml,
        emit_progress_event=emit_progress_event,
    )


def _generate_validated_document_spec(state: _DocxGenerationState) -> DocumentSpec:
    for attempt in range(1, internal_utils.MAX_DOCX_GENERATION_ATTEMPTS + 1):
        quality_feedback = None
        try:
            spec = parse_docx_spec(state.raw_yaml)
            _validate_image_asset_references(spec, state.image_assets)
            issues = quality_gate.collect_document_quality_issues(spec)
            if issues:
                quality_feedback = quality_gate.format_quality_gate_findings(issues)
                raise quality_gate.DocumentQualityFailure(issues)
            return spec
        except Exception as exc:
            if not isinstance(exc, _RETRYABLE_ERRORS):
                raise
            if attempt == internal_utils.MAX_DOCX_GENERATION_ATTEMPTS:
                raise
            _emit_repair_progress_once(state)
            state.raw_yaml = generation_steps.repair_document_yaml(
                project=state.project,
                llm_id=state.llm_id,
                user_prompt=state.user_prompt,
                user_message=state.user_message,
                current_document_yaml=state.raw_yaml,
                validation_error=internal_utils.format_retryable_generation_error(exc),
                quality_feedback=quality_feedback,
                attempt=attempt + 1,
                max_attempts=internal_utils.MAX_DOCX_GENERATION_ATTEMPTS,
            )
    raise ValueError("DOCX generation did not produce a valid document specification")


def _generate_validated_revised_document_spec(state: _DocxRevisionState) -> DocumentSpec:
    for attempt in range(1, internal_utils.MAX_DOCX_GENERATION_ATTEMPTS + 1):
        try:
            spec = parse_docx_spec(state.raw_yaml)
            _validate_image_asset_references(spec, state.image_assets)
        except Exception as exc:
            if not isinstance(exc, _REVISION_RETRYABLE_ERRORS):
                raise
            if attempt == internal_utils.MAX_DOCX_GENERATION_ATTEMPTS:
                raise
            if isinstance(exc, DocxDslSyntaxError):
                _repair_revision_syntax_failure(state, exc, attempt)
            else:
                _repair_revised_document(
                    state,
                    validation_error=internal_utils.format_retryable_generation_error(exc),
                    quality_feedback=None,
                    attempt=attempt + 1,
                )
            continue

        issues = quality_gate.collect_document_quality_issues(spec)
        if not issues:
            return spec
        quality_failure = quality_gate.DocumentQualityFailure(issues)
        if attempt == internal_utils.MAX_DOCX_GENERATION_ATTEMPTS:
            raise quality_failure
        _repair_revised_document(
            state,
            validation_error=str(quality_failure),
            quality_feedback=quality_gate.format_quality_gate_findings(issues),
            attempt=attempt + 1,
        )

    raise ValueError("DOCX revision did not produce a valid document specification")


def _emit_repair_progress_once(state: _DocxGenerationState) -> None:
    if state.repair_progress_emitted:
        return
    if state.emit_progress_event is not None:
        state.emit_progress_event(EventKind.DOCX_REPAIRING_DOCUMENT)
    state.repair_progress_emitted = True


def _repair_revision_syntax_failure(state: _DocxRevisionState, exc: DocxDslSyntaxError, attempt: int) -> None:
    _emit_revision_repair_progress_once(state)
    state.raw_yaml = generation_steps.repair_revision_syntax_only(
        project=state.project,
        llm_id=state.llm_id,
        source_document_yaml=state.source_document_yaml,
        revision_instructions=state.revision_instructions,
        user_message=state.user_message,
        current_document_yaml=state.raw_yaml,
        syntax_error=internal_utils.format_retryable_generation_error(exc),
        attempt=attempt + 1,
        max_attempts=internal_utils.MAX_DOCX_GENERATION_ATTEMPTS,
    )


def _repair_revised_document(
    state: _DocxRevisionState,
    *,
    validation_error: str,
    quality_feedback: str | None,
    attempt: int,
) -> None:
    _emit_revision_repair_progress_once(state)
    state.raw_yaml = generation_steps.repair_revision_document_yaml(
        project=state.project,
        llm_id=state.llm_id,
        source_document_yaml=state.source_document_yaml,
        revision_instructions=state.revision_instructions,
        user_message=state.user_message,
        current_document_yaml=state.raw_yaml,
        validation_error=validation_error,
        quality_feedback=quality_feedback,
        attempt=attempt,
        max_attempts=internal_utils.MAX_DOCX_GENERATION_ATTEMPTS,
    )


def _emit_revision_repair_progress_once(state: _DocxRevisionState) -> None:
    if state.repair_progress_emitted:
        return
    if state.emit_progress_event is not None:
        state.emit_progress_event(EventKind.DOCX_REPAIRING_DOCUMENT)
    state.repair_progress_emitted = True


def _build_generated_docx_artifacts(
    *,
    spec: DocumentSpec,
    image_assets: list[SourceImageAsset],
    user_prompt: str,
    requested_filename: str | None,
    document_metadata: dict | None,
) -> GeneratedDocxArtifacts:
    retained_image_assets = _retain_used_image_assets(spec, image_assets)
    docx_bytes = rendering.render_document_to_bytes(spec, build_image_assets_map(retained_image_assets))
    normalized_yaml = format_document_spec_for_prompt(spec)
    filename = internal_utils.build_docx_filename(
        user_prompt=user_prompt,
        requested_filename=requested_filename,
        document_title=spec.document.title,
    )
    artifacts = internal_utils.build_inline_docx_artifacts(
        title=spec.document.title,
        filename=filename,
        docx_bytes=docx_bytes,
        document_yaml=normalized_yaml,
        image_assets=build_generation_image_assets_metadata(retained_image_assets),
        document_metadata=document_metadata,
    )
    return GeneratedDocxArtifacts(
        filename=filename,
        artifacts=artifacts,
        raw_yaml=normalized_yaml,
    )


def _validate_image_asset_references(
    spec: DocumentSpec,
    image_assets: list[SourceImageAsset],
) -> None:
    available_asset_ids = {asset.asset_id for asset in image_assets}
    for block in spec.blocks:
        if block.type == "image" and block.asset_id not in available_asset_ids:
            raise DocxImageAssetReferenceError(
                f"Image block '{block.id}' references unknown image asset '{block.asset_id}'"
            )


def _retain_used_image_assets(
    spec: DocumentSpec,
    image_assets: list[SourceImageAsset],
) -> list[SourceImageAsset]:
    used_asset_ids = {block.asset_id for block in spec.blocks if block.type == "image"}
    return [asset for asset in image_assets if asset.asset_id in used_asset_ids]
