from __future__ import annotations

from backend.services.artifact_models import PreparedToolSourceArtifact
from backend.services.tool_source_artifacts.image_assets import SourceImageAsset


def add_source_artifacts_to_docx_prompt(
    query: str,
    source_artifacts: list[PreparedToolSourceArtifact],
    image_assets: list[SourceImageAsset] | None = None,
) -> str:
    """Append normalized, DOCX-supported source material without exposing storage references."""

    if not source_artifacts:
        return query

    lines = [
        query.strip(),
        "",
        "Source artifact context:",
        "Use the following referenced artifact content as source material for the document when relevant.",
        "Prefer the referenced artifact data for facts and tables and the referenced images for image blocks, instead of inventing them.",
    ]
    for source_artifact in source_artifacts:
        lines.extend(_format_source_artifact_for_prompt(source_artifact))
    image_assets_section = build_docx_image_assets_prompt_section(image_assets or [])
    if image_assets_section:
        lines.extend(["", image_assets_section])
    return "\n".join(lines).strip()


def build_docx_source_artifact_user_message(
    query: str,
    source_artifacts: list[PreparedToolSourceArtifact],
    image_assets: list[SourceImageAsset] | None = None,
) -> dict[str, object]:
    resolved_image_assets = image_assets or []
    return {
        "role": "user",
        "parts": _build_docx_multimodal_parts(
            text=add_source_artifacts_to_docx_prompt(query, source_artifacts, resolved_image_assets),
            image_assets=resolved_image_assets,
        ),
    }


def build_docx_tool_user_message(
    *,
    query: str,
    conversation_context_section: str,
    source_artifacts: list[PreparedToolSourceArtifact],
    image_assets: list[SourceImageAsset],
) -> dict[str, object] | None:
    """Prepend the conversation context to the query, then delegate to the matching source-material builder.

    Returns None when there is nothing to add, so the generation service keeps using the plain prompt.
    """

    if not (conversation_context_section.strip() or source_artifacts or image_assets):
        return None
    text = "\n\n".join(part for part in (query.strip(), conversation_context_section.strip()) if part)
    if source_artifacts:
        return build_docx_source_artifact_user_message(text, source_artifacts, image_assets)
    return build_docx_image_assets_user_message(text, image_assets)


def build_docx_image_assets_user_message(
    query: str,
    image_assets: list[SourceImageAsset],
) -> dict[str, object]:
    text = "\n".join([query.strip(), "", build_docx_image_assets_prompt_section(image_assets)]).strip()
    return {
        "role": "user",
        "parts": _build_docx_multimodal_parts(text=text, image_assets=image_assets),
    }


def build_docx_image_assets_prompt_section(image_assets: list[SourceImageAsset]) -> str:
    if not image_assets:
        return ""

    lines = [
        "Available image assets for this document:",
        "- Use only these exact asset ids in image blocks.",
        "- An image block requires `id`, `type: image`, and `asset_id`.",
        "- It may include a literal `caption`, semantic `width`, and `alignment`.",
    ]
    for index, asset in enumerate(image_assets, start=1):
        label = asset.artifact_name or f"Image {index}"
        lines.append(f"- {asset.asset_id}: {label}")
    return "\n".join(lines)


def _build_docx_multimodal_parts(
    *,
    text: str,
    image_assets: list[SourceImageAsset],
) -> list[dict[str, str]]:
    parts: list[dict[str, str]] = [{"type": "TEXT", "text": text}]
    for index, image_asset in enumerate(image_assets, start=1):
        label = image_asset.artifact_name or f"Image {index}"
        parts.extend(
            [
                {
                    "type": "TEXT",
                    "text": (
                        f"Attached source image {index}: asset_id={image_asset.asset_id}; label={label}. "
                        f"When using this visual in document YAML, reference asset_id `{image_asset.asset_id}`."
                    ),
                },
                {
                    "type": "IMAGE_INLINE",
                    "inlineImage": image_asset.data_base64,
                    "imageMimeType": image_asset.mime_type,
                },
            ]
        )
    return parts


def _format_source_artifact_for_prompt(source_artifact: PreparedToolSourceArtifact) -> list[str]:
    lines = [
        f"- Referenced artifact: {source_artifact.artifact_name or 'Artifact'}",
        "  content:",
    ]
    lines.extend(f"    {line}" for line in (source_artifact.content or "").splitlines())
    return lines
