import functools
from typing import Dict, Generator

import dataiku
from backend.config import get_project_key
from backend.models.events import EventKind
from backend.utils.artifact_summaries import render_artifact_summary_messages


def add_history_to_completion(
    completion,
    messages,
    include_artifact_summaries: bool = False,
):
    # Artifact summaries are not generic history context. At the moment only orchestrator
    # flows should opt in and expose them when serializing history for an LLM.
    if include_artifact_summaries:
        messages = render_artifact_summary_messages(messages)
    # TODO improve to make sure we always have messages of type dict or MessageRead
    for msg in messages:
        if isinstance(msg, dict):
            completion.with_message(message=_message_text_for_completion(msg), role=msg["role"])
        elif hasattr(msg, "role") and hasattr(msg, "content"):
            completion.with_message(message=msg.content, role=msg.role)
    return completion


def build_conversation_context_section(*, conversation_context: str | None) -> str:
    """Compose the conversation context block embedded in a generation tool's user message."""

    if not (conversation_context and conversation_context.strip()):
        return ""
    return f"Conversation context for this request:\n{conversation_context.strip()}"


def normalize_llm_id(llm_id: str) -> str:
    """Ensure enterprise/augmented IDs carry the project prefix when missing."""
    if not llm_id:
        return llm_id
    if llm_id.count(":") >= 2:
        return llm_id
    if llm_id.startswith("agent:") or llm_id.startswith("retrieval-augmented-llm:"):
        return f"{get_project_key()}:{llm_id}"
    return llm_id


@functools.lru_cache(maxsize=128)
def get_llm_display_name(llm_id: str) -> str:
    """Support short and long llm_id"""
    if not llm_id:
        return ""
    project_key = get_project_key()
    short_id = llm_id
    if llm_id.startswith(f"{project_key}:"):
        short_id = llm_id[len(project_key) + 1 :]
    name = get_llm_friendly_name(short_id, project_key)
    if name:
        return name
    name = get_llm_friendly_name(llm_id, project_key)
    if name:
        return name
    return short_id or llm_id


@functools.lru_cache(maxsize=128)
def get_llm_friendly_name(llm_id: str, project_key: str) -> str:
    project = dataiku.api_client().get_project(project_key)
    llms = project.list_llms()

    for llm in llms:
        if llm.get("id") == llm_id:
            return llm.get("friendlyName", llm_id)
    return ""


def get_tool_validation_requests_msg(tool_validation_requests):
    return {
        "role": "toolValidationRequests",
        "toolValidationRequests": tool_validation_requests,
    }


def get_memory_fragment_msg(memory_fragment):
    return {
        "role": "memoryFragment",
        "memoryFragment": memory_fragment,
    }


def get_tool_validation_responses_msg(tool_validation_responses):
    return {
        "role": "toolValidationResponses",
        "toolValidationResponses": tool_validation_responses,
    }


def add_completion_msgs(comp, messages):
    from backend.utils.utils import append_documents

    for m in messages:
        if m["role"] in ["system", "user", "assistant"]:
            if m.get("parts"):
                msg = comp.new_multipart_message(role=m["role"])
                msg.parts = m["parts"]
                msg.add()
            else:
                comp = comp.with_message(m["content"], role=m["role"])
            docs = m.get("__documents__") or []
            if docs:
                append_documents(docs, comp)
        if m["role"] == "memoryFragment":
            comp = comp.with_memory_fragment(m["memoryFragment"])
        if m["role"] == "toolValidationResponses":
            for response in m["toolValidationResponses"]:
                comp = comp.with_tool_validation_response(
                    validation_request_id=response["validationRequestId"],
                    validated=response["validated"],
                    arguments=response.get("arguments"),
                )
        if m["role"] == "toolValidationRequests":
            comp = comp.with_tool_validation_requests(m["toolValidationRequests"])
    return comp


def _message_text_for_completion(message: dict) -> str:
    content = message.get("content")
    if isinstance(content, str):
        return content

    parts = message.get("parts")
    if not isinstance(parts, list):
        return ""

    text_parts = []
    for part in parts:
        if not isinstance(part, dict):
            continue
        if part.get("type") != "TEXT":
            continue
        text = part.get("text")
        if isinstance(text, str) and text:
            text_parts.append(text)

    return "\n\n".join(text_parts)


def is_agent_like(llm_id: str) -> bool:
    """Detect user/enterprise agents or augmented llms from id string."""
    if not llm_id:
        return False
    parts = llm_id.split(":")
    if parts[0] in ("agent", "retrieval-augmented-llm"):
        return True
    if len(parts) >= 2 and parts[1] in ("agent", "retrieval-augmented-llm"):
        return True
    return False
