import uuid

import dataiku
from backend.config import get_charts_generation_llm_id, get_config, get_enterprise_agents
from backend.constants import (
    GENERATE_DOCX_TOOL_REF,
    GENERATE_PPTX_TOOL_REF,
    PUBLISHED_ZONE,
    REVISE_DOCX_TOOL_REF,
    REVISE_PPTX_TOOL_REF,
)
from backend.models.events import EventKind
from backend.schemas import schemas
from backend.utils.events_utils import extract_artifacts_preview, get_artifacts_size_mb, has_inline, has_records
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import get_visual_agent_in_zone
from backend.utils.sources_utils import (
    extract_nested_vector_search_tool_calls,
    format_sources_for_ui,
    get_records_from_artifacts,
)

logger = get_logger(__name__)

# Track text generation state per streaming session
_stream_text_generating = set()

# Track reasoning chunks per streaming session
_stream_reasoning_chunks = {}


# Checks if the artifact was generated by an AgentConnect built-in document tool.
# This is used to determine if the artifacts should be persisted or not.
AGENT_HUB_DOCUMENT_TOOL_REFS = {
    GENERATE_DOCX_TOOL_REF,
    GENERATE_PPTX_TOOL_REF,
    REVISE_DOCX_TOOL_REF,
    REVISE_PPTX_TOOL_REF,
}


def _is_agent_hub_document_artifact(artifacts: list[dict]) -> bool:
    for artifact in artifacts or []:
        for hierarchy_entry in artifact.get("hierarchy") or []:
            if hierarchy_entry.get("toolRef") in AGENT_HUB_DOCUMENT_TOOL_REFS:
                return True
    return False


def _is_agent_hub_reusable_image_artifact(artifacts: list[dict]) -> bool:
    """Return True for Agent Hub artifacts that expose a reusable inline image payload.

    We keep Agent Hub artifact persistence narrow on purpose: most non-document artifacts are transient tool outputs
    that should not become part of the persisted conversation artifact registry. PPTX generation/revision is the
    exception because a later turn may need to reread the original image artifact referenced by the deck metadata.
    """
    for artifact in artifacts or []:
        for part in artifact.get("parts") or []:
            if not isinstance(part, dict):
                continue
            if part.get("type") not in {"DATA_INLINE", "SIMPLE_DOCUMENT"}:
                continue
            mime_type = part.get("mimeType")
            if isinstance(mime_type, str) and mime_type.lower().startswith("image/"):
                return True
    return False


def get_vector_search_tool_calls(additional_information: dict) -> list[dict]:
    trajectory = additional_information.get("trajectory") or {}
    if not isinstance(trajectory, dict):
        return []

    agent_loop = trajectory.get("agentLoop") or []
    if not isinstance(agent_loop, list):
        return []

    tool_calls = []
    for entry in agent_loop:
        if (
            not isinstance(entry, dict)
            or entry.get("type") != "TOOLCALL"
            or entry.get("toolType") != "VectorStoreSearch"
        ):
            continue
        input_content = entry.get("inputContent") or {}
        # Only forward the fields the UI actually renders (search query + filter for
        # FilterTagDisplay). The raw trajectory entry also carries outputContent (raw KB
        # search results) and other internal metadata that must not reach the browser
        # or the persisted event log.
        tool_calls.append(
            {
                "toolCallId": entry.get("toolCallId"),
                "toolName": entry.get("toolName"),
                "toolFunctionName": entry.get("toolFunctionName"),
                "inputContent": {
                    "searchQuery": input_content.get("searchQuery"),
                    "filter": input_content.get("filter"),
                },
            }
        )
    return tool_calls


def normalise_stream_event(
    ev,
    tcb,
    pcb,
    msgs,
    aid=None,
    trace=None,
    aname=None,
    query=None,
    artifacts_meta=None,
    tool_validation_requests=None,
    memory_fragment=None,
    stream_id=None,
):
    if artifacts_meta is None:
        artifacts_meta = {}

    if "chunk" in ev:
        chunk = ev["chunk"]
        if chunk.get("text"):
            # emit answer_stream_start when text generation starts (first chunk or after interruption)
            if stream_id and stream_id not in _stream_text_generating:
                if aname:
                    pcb({"eventKind": EventKind.ANSWER_STREAM_START, "eventData": {"agentName": aname}})
                else:
                    pcb({"eventKind": EventKind.ANSWER_STREAM_START, "eventData": {"agentName": "Agent Hub"}})
                _stream_text_generating.add(stream_id)
            tcb(chunk["text"])
        elif chunk.get("type") == "event":
            data = {}
            if chunk.get("eventData"):
                data = chunk["eventData"]
            data = chunk.get("eventData") or {}
            if aname and "agentName" not in data:
                data["agentName"] = aname
            pcb({"eventKind": chunk["eventKind"], "eventData": data})
        elif "artifacts" in chunk:
            artifacts = chunk["artifacts"]

            # Detect reasoning artifacts
            has_reasoning = isinstance(artifacts, list) and any(
                artifact.get("type") == "REASONING" for artifact in artifacts
            )
            if has_reasoning:
                for art in artifacts:
                    if art.get("type") == "REASONING":
                        pcb(
                            {
                                "eventKind": EventKind.AGENT_REASONING,
                                "eventData": {
                                    "artf_id": art.get("id"),
                                    "parts": art.get("parts"),
                                    "agentName": aname,
                                },
                            }
                        )
            if (
                aname in {"Agent Hub", "Agent Chat"}
                and not has_reasoning
                and not _is_agent_hub_document_artifact(artifacts)
                and not _is_agent_hub_reusable_image_artifact(artifacts)
            ):
                # Most Agent Hub artifacts are transient tool outputs. Keep the old guardrail and persist only the
                # small reusable subset that later turns may target again: generated documents and reusable images.
                return

            # Persist artifacts
            # Preserve the runtime artifacts_id so Artifact Summary references and persisted metadata stay aligned.
            artifacts_id = chunk.get("artifactsId") or str(uuid.uuid4())
            size = get_artifacts_size_mb(artifacts)
            with_records = has_records(artifacts)
            with_inline = has_inline(artifacts)
            max_size = get_config().get("max_artifacts_size_mb", 2) if with_records else 0
            if with_inline:
                max_size = 20
            meta = schemas.ArtifactsMetadata(
                size_mb=size,
                artifacts_id=artifacts_id,
                agentName=aname or "Used tool",
                agentId=aid or "",
                query=query or "",
                has_records=(with_records or with_inline),
                artifacts=artifacts,
                preview=True if size > max_size else False,
            )
            artifacts_meta[artifacts_id] = meta
            if aname in {"Agent Hub", "Agent Chat"}:
                # Agent Hub artifacts that pass the filter above must survive for cross-turn rereads, but we still do
                # not emit download-card events from this aggregation layer.
                return
            # Don't send the Downloads event for REASONING artifacts
            if not has_reasoning:
                pcb(
                    {
                        "eventKind": "artifacts",
                        "eventData": {
                            **meta,
                            "artifacts": extract_artifacts_preview(artifacts) if size > max_size else artifacts,
                        },
                    },
                    store_event=False,
                )
            if (with_records or with_inline) and size > max_size:
                logger.info(
                    f"Artifacts size {size} MB exceeds the maximum allowed size, skipping visualization generation."
                )
            elif (with_records or with_inline) and size <= max_size:
                logger.info(f"Artifacts size: {size} MB")
                handle_artifacts(
                    artifacts=artifacts,
                    artifacts_id=artifacts_id,
                    msgs=msgs,
                    pcb=pcb,
                    trace=trace,
                    aname=aname,
                    query=query,
                )
        elif "toolValidationRequests" in chunk:
            tool_validation_requests_chunk = chunk["toolValidationRequests"]
            # Store as a list directly
            if tool_validation_requests is not None:
                # tool_validation_requests.clear()
                tool_validation_requests.extend(tool_validation_requests_chunk)

            event_data = chunk.get("eventData", {})
            event_data["requests"] = tool_validation_requests_chunk
            event_data["agentName"] = aname or "Agent Hub"
            event_data["agentId"] = aid
            pcb(
                {
                    "eventKind": EventKind.TOOL_VALIDATION_REQUESTS,
                    "eventData": event_data,
                },
                store_event=True,
            )
        elif "memoryFragment" in chunk:
            memory_fragment_chunk = chunk["memoryFragment"]
            # Store the object directly
            if memory_fragment is not None:
                # memory_fragment.clear()
                memory_fragment.update(memory_fragment_chunk)
            # else:
            # memory_fragment["value"] = memory_fragment_chunk
            event_data = chunk.get("eventData", {})
    elif "footer" in ev:
        additional_information = ev["footer"].get("additionalInformation", {})
        raw_sources = additional_information.get("sources", [])
        sources = format_sources_for_ui(raw_sources, aid) if raw_sources else []
        # Single-agent calls expose the trajectory directly in the footer. Multi-agent
        # orchestration (agents called as tools) has no top-level trajectory; instead,
        # each sub-agent's tool calls are embedded on its sources by AHDSSAgentTool.run
        # and recovered here (see extract_nested_vector_search_tool_calls).
        kb_tool_calls = get_vector_search_tool_calls(additional_information) or extract_nested_vector_search_tool_calls(
            raw_sources
        )
        if sources or kb_tool_calls:
            event_data = {"sources": sources}
            if kb_tool_calls:
                event_data["toolCalls"] = kb_tool_calls
            pcb({"eventKind": "references", "eventData": event_data})
        # Clean up reasoning chunks when stream ends (footer indicates end of stream)
        if stream_id and stream_id in _stream_reasoning_chunks:
            accumulated_reasoning = "".join(_stream_reasoning_chunks[stream_id])
            if accumulated_reasoning:
                logger.debug(
                    f"Stream completed (stream_id: {stream_id[:8]}...). "
                    f"Total reasoning text accumulated: {len(accumulated_reasoning)} characters. "
                    f"Full accumulated reasoning: {accumulated_reasoning}"
                )
            del _stream_reasoning_chunks[stream_id]
    # elif "trace_ready" in ev:
    #     trace_data = ev["trace_ready"].get("trace")
    #     if trace_data:
    #         pcb({"eventKind": "TRACE", "eventData": {"trace": trace_data}})


def handle_artifacts(artifacts, artifacts_id, msgs, pcb, trace=None, aname=None, query=None):
    from backend.models.chart import MAX_CHART_SAMPLE_ROWS
    from backend.services.visualization_service import VisualizationService

    vis_mode = get_config().get("visualization_generation_mode")
    artifacts_records = get_records_from_artifacts(artifacts)
    graph_llm_id = get_charts_generation_llm_id()
    logger.info(f"Visualization mode: {vis_mode}, graph_llm_id: {graph_llm_id}")
    # We generate a chart per records for the chart generation
    if vis_mode == "AUTO" and artifacts_records and graph_llm_id:
        for art_index in range(len(artifacts)):
            for idx, r in enumerate(artifacts_records[art_index] or []):
                pcb({"eventKind": EventKind.GENERATING_CHART, "eventData": {}})
                plan = None
                excep = None
                try:
                    plan = VisualizationService.prepare_chart_plan(
                        llm_id=graph_llm_id,
                        messages=msgs,
                        data=r["data"][:MAX_CHART_SAMPLE_ROWS],
                        columns=r["columns"],
                        trace=trace,
                        row_count=len(r["data"]),
                    )
                except Exception as e:
                    logger.exception(f"Failed to generate graph {e}")
                    excep = e
                if plan:
                    pcb(
                        {
                            "eventKind": EventKind.CHART_PLAN,
                            "eventData": {
                                "chart_plan": plan,
                                "source_name": aname or "Used tool",
                                "artifacts_id": artifacts_id,
                                "artifact_index": art_index,
                                "records_index": idx,
                                "query": query,
                            },
                        }
                    )
                else:
                    pcb(
                        {
                            "eventKind": EventKind.CHART_PLAN_ERROR,
                            "eventData": f"Failed to get chart {excep if excep else ''}",
                        }
                    )


def get_selected_agents_as_objs(
    store,
    sel_ids,
    *,
    enterprise_agents_by_id: dict[str, dict] | None = None,
    user_agents: dict[str, object] | None = None,
) -> list[dict]:
    """Convert selected ids into the DSS agent objects expected by orchestration code.

    When `user_agents` is omitted, this falls back to `store.get_all_agents()` and
    rebuilds the full user-agent catalog before filtering to `sel_ids`. The
    conversation flow is expected to pass the targeted subset explicitly so this
    helper can stay scoped to the selected ids only.
    """
    from backend.schemas import schemas
    from flask import g

    agents_obj = []
    current_user = getattr(g, "authIdentifier", None)
    enterprise = (
        list((enterprise_agents_by_id or {}).values())
        if enterprise_agents_by_id is not None
        else get_enterprise_agents(current_user)
    )
    client = dataiku.api_client()
    user_agents_list: list[schemas.AgentRead]
    if user_agents is None:
        user_agents_list = store.get_all_agents()
    else:
        user_agents_list = [ua for ua in user_agents.values() if isinstance(ua, schemas.AgentRead)]
    for ea in enterprise:
        if ea.get("id") in sel_ids:
            agents_obj.append(ea)
    for ua in user_agents_list:
        if ua.id in sel_ids:
            aid = ua.id
            is_owner = ua.owner == current_user
            # MUST have published version
            if not ua.published_version:
                logger.warning(f"Agent {aid} has no published version, skipping")
                continue  # Skip this agent
            logger.debug(f"Agent {aid}: is_owner={is_owner}, use_published={True}")
            project_key = ua.id
            project = client.get_project(project_key)
            vis_agent = get_visual_agent_in_zone(project, PUBLISHED_ZONE)
            if not vis_agent:
                logger.warning(f"Agent {aid} has no visual agent in published zone, skipping")
                continue
            agents_obj.append(
                {
                    **ua.model_dump(mode="json"),
                    "id": f"{project_key}:agent:{vis_agent.id}",
                    "uaid": aid,
                    "name": ua.published_version.name or "Unnamed Agent",
                }
            )  # include uaid to keep original id needed later in orchestrator service
    return agents_obj
