from __future__ import annotations

from collections.abc import Callable
from contextlib import nullcontext
from types import SimpleNamespace
from typing import Any, Optional

import dataiku
from backend.agents.tools.generated_document_fields import (
    CONVERSATION_CONTEXT_DESCRIPTION,
    REVISE_PPTX_FILENAME_DESCRIPTION,
    REVISION_SOURCE_ARTIFACTS_DESCRIPTION,
)
from backend.agents.tools_using_v2 import LLMTool
from backend.config import get_default_llm_id
from backend.constants import REVISE_PPTX_TOOL_NAME, REVISE_PPTX_TOOL_REF
from backend.services.artifact_models import ArtifactReference, ArtifactReferenceSpec
from backend.services.generated_document.generated_document_metadata import (
    GENERATED_DOCUMENT_FORMAT_PPTX,
    build_revision_document_metadata,
)
from backend.services.generated_document.generated_document_revision_source import (
    resolve_generated_document_revision_source,
)
from backend.services.generated_document.revision_models import GeneratedDocumentRevisionSource
from backend.services.pptx_artifact_generation.image_assets import (
    build_image_assets_from_source_artifacts,
    merge_image_assets,
)
from backend.services.pptx_artifact_generation.service import revise_pptx_artifacts
from backend.services.pptx_artifact_generation.source_artifacts import build_pptx_tool_user_message
from backend.services.tool_source_artifacts.source_artifact_input_service import prepare_source_artifacts_for_tool
from backend.socket import socketio
from backend.utils.llm_utils import build_conversation_context_section
from backend.utils.logging_utils import get_logger
from pydantic import BaseModel, Field

logger = get_logger(__name__)
PPTX_REVISION_PROGRESS_EVENT_DATA = {"toolName": "Revise PPTX"}


class RevisePptxSpec(BaseModel):
    source: ArtifactReferenceSpec = Field(
        ...,
        description=(
            "Artifact Reference for the existing Agent Hub generated PPTX to revise. Use the message_id, "
            "artifacts_id, and artifact_index exactly as shown in the Artifact Summary."
        ),
    )
    revision_instructions: str = Field(
        ...,
        description=(
            "Concise, complete, self-contained natural-language request describing the changes to apply. Preserve "
            "unchanged slides and styling unless the user explicitly asks for broader changes."
        ),
    )
    filename: Optional[str] = Field(
        default=None,
        description=REVISE_PPTX_FILENAME_DESCRIPTION,
    )
    conversation_context: str | None = Field(
        default=None,
        description=CONVERSATION_CONTEXT_DESCRIPTION,
    )
    source_artifacts: list[ArtifactReferenceSpec] | None = Field(
        default=None,
        description=REVISION_SOURCE_ARTIFACTS_DESCRIPTION,
    )


class RevisePptxArtifactTool:
    def __init__(
        self,
        *,
        app=None,
        store=None,
        llm_id: str | None = None,
        progress_event_callback: Callable[[dict[str, Any]], None] | None = None,
        conversation_history: list[dict] | None = None,
    ):
        self._app = app
        self._store = store
        self._llm_id = llm_id
        self._progress_event_callback = progress_event_callback
        self._conversation_history = conversation_history or []

    project_key = "AgentHub"
    tool_id = REVISE_PPTX_TOOL_NAME
    tool_name = REVISE_PPTX_TOOL_NAME

    def get_settings(self):
        return SimpleNamespace(_settings={"type": "builtin"})

    def describe_tool_call(self, input, llm_descriptor, context=None, subtool_name=None):
        revision_instructions = ((input or {}).get("revision_instructions") or "PPTX revision").strip()
        return f"Revise a PPTX with '{revision_instructions}'."

    def run(
        self,
        input,
        context=None,
        subtool_name=None,
        memory_fragment=None,
        tool_validation_responses=None,
        tool_validation_requests=None,
    ):
        spec = RevisePptxSpec.model_validate(input or {})
        project = dataiku.api_client().get_default_project()
        llm_id = self._llm_id or get_default_llm_id()
        emit_progress_event = _build_progress_emitter(
            progress_event_callback=self._progress_event_callback,
            context=context,
        )
        revision_source = _resolve_revision_source(spec, context, self._app, self._store)
        source_artifact_references = [
            ArtifactReference(
                message_id=source_artifact.message_id,
                artifacts_id=source_artifact.artifacts_id,
                artifact_index=source_artifact.artifact_index,
            )
            for source_artifact in (spec.source_artifacts or [])
        ]
        prepared_source_artifacts = prepare_source_artifacts_for_tool(
            source_artifacts=source_artifact_references,
            context=context,
            app=self._app,
            accepted_read_modes={"records_preview", "inline_image"},
        )
        merged_image_assets = merge_image_assets(
            revision_source.image_assets,
            build_image_assets_from_source_artifacts(prepared_source_artifacts),
        )
        revision_user_message = build_pptx_tool_user_message(
            query=spec.revision_instructions,
            conversation_context_section=build_conversation_context_section(
                conversation_context=spec.conversation_context,
            ),
            source_artifacts=prepared_source_artifacts,
            image_assets=merged_image_assets,
        )
        document_metadata = _build_revision_metadata(revision_source)
        revised = revise_pptx_artifacts(
            project=project,
            llm_id=llm_id,
            source_deck_yaml=revision_source.source_yaml,
            source_image_assets=merged_image_assets,
            revision_instructions=spec.revision_instructions,
            user_message=revision_user_message,
            requested_filename=spec.filename,
            document_metadata=document_metadata,
            emit_progress_event=emit_progress_event,
            conversation_history=self._conversation_history,
        )

        logger.info(
            "revise_pptx tool produced artifact: llm_id=%s filename=%s artifacts=%d revision_instructions=%s",
            llm_id,
            revised.filename,
            len(revised.artifacts),
            spec.revision_instructions,
        )

        return {
            "output": {
                "result": f"Revised PPTX: {revised.filename}",
            },
            "sources": [],
            "artifacts": revised.artifacts,
            "trace": None,
            "toolValidationRequests": None,
            "memoryFragment": None,
        }


def _build_progress_emitter(
    *,
    progress_event_callback: Callable[[dict[str, Any]], None] | None,
    context: dict | None,
):
    if progress_event_callback is not None:

        def emit_progress(event_kind) -> None:
            progress_event_callback(
                {
                    "eventKind": event_kind,
                    "eventData": PPTX_REVISION_PROGRESS_EVENT_DATA,
                }
            )

        return emit_progress

    return _build_live_progress_emitter(context)


def _build_live_progress_emitter(context: dict | None):
    if not isinstance(context, dict):
        return None

    conversation_id = context.get("conversationId")
    user_login = context.get("dku_user_login")
    if not conversation_id or not user_login:
        return None

    room = f"user:{user_login}"

    def emit_progress(event_kind) -> None:
        socketio.emit(
            "chat_event",
            {
                "conversationId": conversation_id,
                "eventKind": event_kind,
                "eventData": PPTX_REVISION_PROGRESS_EVENT_DATA,
            },
            room=room,
        )

    return emit_progress


def build_revise_pptx_llm_tool(
    app=None,
    store=None,
    pcb: Callable[[dict[str, Any]], None] | None = None,
    llm_id: str | None = None,
    conversation_history: list[dict] | None = None,
) -> LLMTool:
    return LLMTool(
        llm_tool_name=REVISE_PPTX_TOOL_NAME,
        llm_descriptor={
            "type": "function",
            "function": {
                "name": REVISE_PPTX_TOOL_NAME,
                "description": (
                    "Create a new revised PPTX generated document artifact from an existing Agent Hub generated PPTX. "
                    "Pass the source Artifact Reference from the Artifact Summary; Agent Hub resolves the editable "
                    "deck source and lineage metadata server-side. Use this only for Agent Hub generated PPTX "
                    "artifacts with editable source. Preserve unchanged slides, styling, and structure unless the "
                    "user asks for broader changes. Do not use it for uploaded files or format conversion."
                ),
                "parameters": RevisePptxSpec.model_json_schema(),
            },
        },
        dku_tool=RevisePptxArtifactTool(
            app=app,
            store=store,
            progress_event_callback=pcb,
            llm_id=llm_id,
            conversation_history=conversation_history,
        ),
        dku_tool_ref=REVISE_PPTX_TOOL_REF,
        dku_tool_type="builtin",
        dku_tool_name="Revise PPTX",
        dku_subtool_name=None,
        require_validation=False,
        allow_editing_inputs=False,
    )


def _resolve_revision_source(
    spec: RevisePptxSpec,
    context: dict | None,
    app,
    store=None,
) -> GeneratedDocumentRevisionSource:
    if not isinstance(context, dict):
        raise ValueError("Conversation context is required to resolve the revision source")

    conversation_id = context.get("conversationId")
    if not conversation_id:
        raise ValueError("Conversation id is required to resolve the revision source")

    app_context = app.app_context() if app is not None else nullcontext()
    with app_context:
        resolved_store = store
        if resolved_store is None:
            from backend.utils.utils import get_store

            resolved_store = get_store()

        return resolve_generated_document_revision_source(
            store=resolved_store,
            conversation_id=conversation_id,
            reference=spec.source,
            expected_format=GENERATED_DOCUMENT_FORMAT_PPTX,
            runtime_context=context,
        )


def _build_revision_metadata(revision_source: GeneratedDocumentRevisionSource) -> dict[str, Any]:
    return build_revision_document_metadata(
        lineage_id=revision_source.lineage_id,
        source=revision_source.source,
        source_lineage_role=revision_source.lineage_role,
    )
