from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Optional

import dataiku

from backend.config import agents_as_tools
from backend.models.events import EventKind
from backend.services.helpers.attachment_history_utils import history_has_snapshots
from backend.services.helpers.multimodal_attachment_handling import (
    AttachmentForwardingPolicy,
    is_llm_multimodal,
    resolve_attachment_forwarding_policy,
)
from backend.utils.agent_id_utils import parse_agent_id
from backend.utils.llm_utils import is_agent_like, normalize_llm_id
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)


@dataclass
class AttachmentPreparationResult:
    """Prepared attachment routing for one target (agent or base LLM)."""

    target_id: str
    target_name: str
    selected_target: dict | None
    history: list[dict] | None
    excluded: bool
    missing_text_fallback_files: list[str]


@dataclass
class AttachmentPreparationResults:
    """Global routing output across all targets for one conversation turn."""

    by_target: dict[str, AttachmentPreparationResult]
    events: list[dict]
    blocking_error: Optional[str] = None

    @property
    def filtered_targets(self) -> list[dict]:
        return [r.selected_target for r in self.by_target.values() if r.selected_target and not r.excluded]

    @property
    def history_by_target(self) -> dict[str, list[dict]]:
        return {
            target_id: r.history for target_id, r in self.by_target.items() if r.history is not None and not r.excluded
        }


@dataclass
class StrippedHistoryResult:
    """Text-only history plus counters used for UX events and blocking checks."""

    history: list[dict]
    replaced_snapshots_count: int
    replaced_documents_count: int
    missing_text_fallback_files: set[str]


def resolve_multimodal_capability(project, llm_id: str) -> bool:
    """
    Return whether a target supports IMAGE_INPUT using compatibility fallbacks.

    This wraps the low-level `is_llm_multimodal(...)` check by retrying with normalized
    and short-id variants, because agent/llm ids do not always come in the exact format
    returned by `project.list_llms(...)`.
    """
    if not llm_id:
        return False
    normalized = normalize_llm_id(llm_id)
    if is_llm_multimodal(project, llm_id):
        return True
    if normalized != llm_id and is_llm_multimodal(project, normalized):
        return True
    if ":" in normalized:
        short = normalized.split(":", 1)[1]
        return is_llm_multimodal(project, short)
    return False


def prepare_attachment_routing(
    *,
    sel_agents: list[dict],
    history: list[dict],
    mode_agents: bool,
    project,
    user_agents_by_id: dict,
    get_default_llm_id: Callable[[], str],
) -> AttachmentPreparationResults:
    """Decide per-target history (`original` vs `text-only`) and any blocking/error events."""
    if not history_has_snapshots(history):
        by_target = {
            a.get("agentId"): AttachmentPreparationResult(
                target_id=a.get("agentId"),
                target_name=a.get("agentName") or a.get("agentId") or "unknown",
                selected_target=a,
                history=history,
                excluded=False,
                missing_text_fallback_files=[],
            )
            for a in sel_agents
            if a.get("agentId")
        }
        return AttachmentPreparationResults(
            by_target=by_target,
            events=[],
        )

    strip_result = remove_history_snapshots(history)

    # No agents to route to (e.g. agents mode is on but the project has none selected/available)
    # falls back to the base LLM, same as explicit LLM mode, instead of reporting a blocking error.
    if not sel_agents:
        return _prepare_base_llm_attachment_routing(
            history=history,
            strip_result=strip_result,
            project=project,
            get_default_llm_id=get_default_llm_id,
        )

    return _prepare_agents_attachment_routing(
        sel_agents=sel_agents,
        history=history,
        strip_result=strip_result,
        mode_agents=mode_agents,
        project=project,
        user_agents_by_id=user_agents_by_id,
    )


def _prepare_base_llm_attachment_routing(
    *,
    history: list[dict],
    strip_result: StrippedHistoryResult,
    project,
    get_default_llm_id: Callable[[], str],
) -> AttachmentPreparationResults:
    base_llm_id = normalize_llm_id(get_default_llm_id())
    policy = resolve_attachment_forwarding_policy(
        multimodal=resolve_multimodal_capability(project, base_llm_id),
        has_text_fallback=not strip_result.missing_text_fallback_files,
    )
    if policy == AttachmentForwardingPolicy.CANT_FORWARD:
        return AttachmentPreparationResults(
            by_target={},
            events=[],
            blocking_error=build_blocking_error(
                is_base_llm=True,
                impacted_agent_names=[base_llm_id],
                missing_files=strip_result.missing_text_fallback_files,
            ),
        )

    routed_history = (
        history if policy == AttachmentForwardingPolicy.FORWARD_TEXT_AND_SNAPSHOTS else strip_result.history
    )
    impacted_agents = []
    if policy != AttachmentForwardingPolicy.FORWARD_TEXT_AND_SNAPSHOTS:
        impacted_agents = [{"agentId": base_llm_id, "agentName": base_llm_id, "excluded": False}]

    return AttachmentPreparationResults(
        by_target={
            "__base_llm__": AttachmentPreparationResult(
                target_id="__base_llm__",
                target_name=base_llm_id,
                selected_target=None,
                history=routed_history,
                excluded=False,
                missing_text_fallback_files=[],
            )
        },
        events=build_attachment_events(
            impacted_agents=impacted_agents,
            replaced_snapshots_count=strip_result.replaced_snapshots_count,
            replaced_documents_count=strip_result.replaced_documents_count,
        ),
    )


def _prepare_agents_attachment_routing(
    *,
    sel_agents: list[dict],
    history: list[dict],
    strip_result: StrippedHistoryResult,
    mode_agents: bool,
    project,
    user_agents_by_id: dict,
) -> AttachmentPreparationResults:
    by_target: dict[str, AttachmentPreparationResult] = {}
    viable_sel_agents: list[dict] = []
    impacted_agents: list[dict] = []
    tools_mode = mode_agents and agents_as_tools() and len(sel_agents) > 1
    force_text_only_in_tools_mode = False

    for sel in sel_agents:
        agent_id = sel.get("agentId")
        if not agent_id:
            continue
        multimodal = target_is_multimodal(target_id=agent_id, project=project, user_agents_by_id=user_agents_by_id)
        policy = resolve_attachment_forwarding_policy(
            multimodal=multimodal,
            has_text_fallback=not strip_result.missing_text_fallback_files,
        )
        if policy == AttachmentForwardingPolicy.FORWARD_TEXT_AND_SNAPSHOTS:
            viable_sel_agents.append(sel)
            by_target[agent_id] = AttachmentPreparationResult(
                target_id=agent_id,
                target_name=sel.get("agentName") or agent_id,
                selected_target=sel,
                history=history,
                excluded=False,
                missing_text_fallback_files=[],
            )
            continue
        excluded = policy == AttachmentForwardingPolicy.CANT_FORWARD
        if tools_mode and not multimodal and not excluded:
            force_text_only_in_tools_mode = True
        impacted_agents.append(
            {
                "agentId": agent_id,
                "agentName": sel.get("agentName") or agent_id,
                "excluded": excluded,
                "missingTextFallbackFiles": sorted(strip_result.missing_text_fallback_files),
            }
        )
        if excluded:
            by_target[agent_id] = AttachmentPreparationResult(
                target_id=agent_id,
                target_name=sel.get("agentName") or agent_id,
                selected_target=sel,
                history=None,
                excluded=True,
                missing_text_fallback_files=sorted(strip_result.missing_text_fallback_files),
            )
            continue
        viable_sel_agents.append(sel)
        by_target[agent_id] = AttachmentPreparationResult(
            target_id=agent_id,
            target_name=sel.get("agentName") or agent_id,
            selected_target=sel,
            history=strip_result.history,
            excluded=False,
            missing_text_fallback_files=[],
        )

    if tools_mode and force_text_only_in_tools_mode:
        for target in by_target.values():
            if not target.excluded:
                target.history = strip_result.history

    if not viable_sel_agents:
        return AttachmentPreparationResults(
            by_target=by_target,
            events=[],
            blocking_error=build_blocking_error(
                is_base_llm=False,
                impacted_agent_names=[a.get("agentName") or a.get("agentId") or "unknown" for a in impacted_agents],
                missing_files=strip_result.missing_text_fallback_files,
            ),
        )

    return AttachmentPreparationResults(
        by_target=by_target,
        events=build_attachment_events(
            impacted_agents=impacted_agents,
            replaced_snapshots_count=strip_result.replaced_snapshots_count,
            replaced_documents_count=strip_result.replaced_documents_count,
        ),
    )


def remove_history_snapshots(history: list[dict]) -> StrippedHistoryResult:
    """Remove snapshots and keep text fallback; only latest user turn can make fallback missing blocking."""
    replaced_snapshots = 0
    replaced_documents = 0
    missing_files: set[str] = set()
    out: list[dict] = []
    latest_user_idx = None

    for idx, msg in enumerate(history):
        if msg.get("role") == "user":
            latest_user_idx = idx

    for idx, msg in enumerate(history):
        copied = dict(msg)
        docs = msg.get("__documents__") or []
        if not docs:
            out.append(copied)
            continue

        filtered_docs: list[dict] = []
        for doc in docs:
            text = (doc.get("text") or "").strip()
            has_text_fallback = bool(doc.get("has_text_fallback"))
            snapshots = doc.get("snapshots") or []

            if text and (has_text_fallback or not snapshots):
                doc_copy = dict(doc)
                doc_copy["snapshots"] = []
                filtered_docs.append(doc_copy)
                if snapshots:
                    replaced_snapshots += len(snapshots)
                    replaced_documents += 1
            elif snapshots:
                if latest_user_idx is not None and idx == latest_user_idx and msg.get("role") == "user":
                    missing_files.add(doc.get("name") or "unknown")
            else:
                doc_copy = dict(doc)
                doc_copy["snapshots"] = []
                filtered_docs.append(doc_copy)

        if filtered_docs:
            copied["__documents__"] = filtered_docs
        else:
            copied.pop("__documents__", None)
        out.append(copied)

    return StrippedHistoryResult(
        history=out,
        replaced_snapshots_count=replaced_snapshots,
        replaced_documents_count=replaced_documents,
        missing_text_fallback_files=missing_files,
    )


def build_attachment_events(
    *,
    impacted_agents: list[dict],
    replaced_snapshots_count: int,
    replaced_documents_count: int,
) -> list[dict]:
    """Create chat events after filtering/replacing attachments for non-multimodal targets."""
    events: list[dict] = []
    if not impacted_agents:
        return events

    fallback_text_agents = [a for a in impacted_agents if not a.get("excluded")]
    excluded_agents = [a for a in impacted_agents if a.get("excluded")]

    if fallback_text_agents and replaced_snapshots_count > 0:
        impacted_agent_names = ", ".join(
            [a.get("agentName") or a.get("agentId") or "unknown" for a in fallback_text_agents]
        )
        events.append(
            {
                "eventKind": EventKind.ATTACHMENTS_FILTERED,
                "eventData": {
                    "reason": "non_multimodal_replaced_with_text",
                    "impactedAgentNames": impacted_agent_names,
                    "replacedSnapshotsCount": replaced_snapshots_count,
                    "replacedDocumentsCount": replaced_documents_count,
                },
            }
        )

    for excluded in excluded_agents:
        name = excluded.get("agentName") or excluded.get("agentId") or "unknown"
        missing = excluded.get("missingTextFallbackFiles") or []
        events.append(
            {
                "eventKind": EventKind.AGENT_ERROR,
                "eventData": {
                    "agentId": excluded.get("agentId") or "",
                    "agentName": name,
                    "reason": "non_multimodal_agent_excluded_for_image_input",
                    "missingTextFallbackFiles": missing,
                },
            }
        )

    return events


def build_blocking_error(*, is_base_llm: bool, impacted_agent_names: list[str], missing_files: set[str]) -> str:
    """Format the chat_error message when no target can accept the current attachments."""
    missing_msg = f" No text fallback for: {', '.join(sorted(missing_files))}." if missing_files else ""
    if is_base_llm:
        return "Unable to process attachments for selected non-multimodal target." + missing_msg
    return (
        "Unable to process attachments for selected non-multimodal agent(s): "
        + ", ".join(impacted_agent_names)
        + "."
        + missing_msg
    )


def target_is_multimodal(*, target_id: str, project, user_agents_by_id: dict) -> bool:
    """Resolve multimodality for user agents and enterprise/base targets."""
    # Resolve My Agents from the preloaded user-agent summaries.
    if target_id in user_agents_by_id:
        ua = user_agents_by_id[target_id]
        llm_id = ua.llmid or ""
        if not llm_id and ua.published_version and ua.published_version.llmid:
            llm_id = ua.published_version.llmid
        return resolve_multimodal_capability(project, llm_id)

    # Resolve enterprise agents and direct LLM targets from their DSS reference.
    target_project, llm_ref = resolve_target_project_and_llm_ref(target_id=target_id, default_project=project)
    if not target_project or not llm_ref:
        return False
    return resolve_multimodal_capability(target_project, llm_ref)


def resolve_target_project_and_llm_ref(*, target_id: str, default_project) -> tuple[object | None, str]:
    """Resolve `(project, llm_ref)` expected by `project.list_llms(...)` id space."""
    # Provider LLM ids also use colons (for example ``openai:connection:model``),
    # so resolve non-agent targets against the default project.
    if not is_agent_like(target_id):
        return default_project, normalize_llm_id(target_id)

    try:
        project_key, parsed = parse_agent_id(target_id)
        project = dataiku.api_client().get_project(project_key)
        return project, parsed
    except Exception:
        try:
            return default_project, normalize_llm_id(target_id)
        except Exception:
            return None, ""
