from typing import Any, List, Optional

from backend.models.sources import AggregatedToolSources, RAGImage, Source
from backend.utils.agent_id_utils import parse_agent_id
from backend.utils.llm_utils import get_llm_friendly_name
from backend.utils.logging_utils import get_logger
from backend.utils.picture_utils import b64encode_image_from_full_folder_id

logger = get_logger(__name__)


def _enrich_metadata_with_tags(metadata: dict) -> dict:
    """
    We don't display metadata directly but we display tags which are
    built into metadata.
    """
    tags = []
    for key, value in metadata.items():
        try:
            value_as_str = str(value)

            tags.append({"name": key, "value": value_as_str, "type": "string"})

        except Exception:
            continue
    if tags:
        metadata["tags"] = tags
    return metadata


def _build_file_based_document_source(source: Source) -> Source:
    """
    Build a source object for frontend based on DSS source
    """
    rag_images: List[RAGImage] = []
    metadata = source.get("metadata")
    textSnippet, htmlSnippet, markdownSnippet, jsonSnippet = "", "", "", ""

    # if there's some imageRefs we display them instead of snippets
    # displaying both would be too much
    if imageRefs := source.get("imageRefs"):
        for image_ref in imageRefs:
            b64_image = b64encode_image_from_full_folder_id(image_ref.get("path"), image_ref.get("folderId"))
            rag_images.append(
                RAGImage(
                    full_folder_id=image_ref.get("folderId") or "",
                    file_path=image_ref.get("path") or "",
                    index=0,
                    file_data=b64_image,
                )
            )
    else:
        textSnippet = source.get("textSnippet") or ""
        htmlSnippet = source.get("htmlSnippet") or ""
        markdownSnippet = source.get("markdownSnippet") or ""
        jsonSnippet = source.get("jsonSnippet") or ""

    title = source.get("title")
    fileRef = source.get("fileRef")
    if isinstance(fileRef, dict):
        page_range = fileRef.get("pageRange") or {}
        page_start = page_range.get("start") or ""
        page_end = page_range.get("end") or ""
        filepath = fileRef.get("path") or ""
        if not title:  # title can be set in the settings (from a metadata) but now always
            title = f"{filepath.split('/')[-1]}"
            # page start / end can not be defined
            # (e.g.: docx with text extraction in KB)
            if page_start:
                page_info = f" (p. {page_start})"
                if page_end and page_end != page_start:
                    page_info = f" (p. {page_start}-{page_end})"
                title += page_info

    if metadata is not None and isinstance(metadata, dict):
        metadata = _enrich_metadata_with_tags(metadata)

    source_for_asw = Source(
        type="SIMPLE_DOCUMENT",
        title=title,
        url=source.get("url"),
        textSnippet=textSnippet,
        htmlSnippet=htmlSnippet,
        markdownSnippet=markdownSnippet,
        images=rag_images,
        jsonSnippet=jsonSnippet,
        metadata=metadata,
        fileRef=fileRef,
    )

    return source_for_asw


def _get_tool_call_id(aggregated_source: AggregatedToolSources) -> Optional[str]:
    for hierarchy_entry in aggregated_source.get("hierarchy") or []:
        if hierarchy_entry.get("type") == "TOOL":
            return hierarchy_entry.get("toolCallId")
    return None


def extract_nested_vector_search_tool_calls(raw_sources: List[AggregatedToolSources]) -> List[dict]:
    """Recover KB search tool calls performed by an agent called as a tool.

    AHDSSAgentTool.run() embeds these on each source it returns (as
    nestedVectorSearchToolCalls), because tools_using_v2.py's generic tool-call
    plumbing only forwards the `sources` field across the tool-call boundary. By the
    time raw_sources reaches here, `_call_tools` has already prepended the outer
    "call the sub-agent" TOOL hierarchy entry, so we re-tag each embedded tool call
    (which still carries its inner, sub-agent-local toolCallId) with that outer id,
    matching what _get_tool_call_id/format_sources_for_ui put on the formatted source.
    """
    tool_calls: List[dict] = []
    seen_outer_ids: set = set()
    for aggregated_source in raw_sources:
        embedded = aggregated_source.get("nestedVectorSearchToolCalls")
        if not embedded:
            continue
        outer_id = _get_tool_call_id(aggregated_source)
        if not outer_id or outer_id in seen_outer_ids:
            continue
        seen_outer_ids.add(outer_id)
        tool_calls.extend({**tool_call, "toolCallId": outer_id} for tool_call in embedded)
    return tool_calls


def convert_to_rows_and_columns(columns: List[str], data: List[Any]):
    rows = [dict(zip(columns, row)) for row in data]

    columns_formatted = [{"name": col, "label": col, "field": col, "align": "left"} for col in columns]

    return {"rows": rows, "columns": columns_formatted}


def format_sources_for_ui(raw_sources: List[AggregatedToolSources], aid: str = None) -> List[AggregatedToolSources]:
    aggregated_sources: List[AggregatedToolSources] = []
    agent_name = aid or "Used tool"
    if aid:
        try:
            project_key, llm_id = parse_agent_id(aid)
            resolved_name = get_llm_friendly_name(llm_id=llm_id, project_key=project_key)
            if resolved_name:
                agent_name = resolved_name
        except ValueError:
            # Invalid id, keep the fallback display name.
            agent_name = "Used tool"
    for aggregated_source in raw_sources:
        items: List[Source] = []

        for source in aggregated_source.get("items") or []:
            source_type = source.get("type", "")
            if source_type == "FILE_BASED_DOCUMENT":
                source_for_asw = _build_file_based_document_source(source)
                items.append(source_for_asw)
            elif source_type in {"GENERATED_SQL_QUERY", "GENERATED_CYPHER_QUERY"}:
                performed_queries = source.get("performedQueries")
                query_text = "\n\n".join(performed_queries) if performed_queries else source.get("performedQuery")
                items.append(Source(textSnippet=query_text, type=source_type))
            elif source_type == "RECORDS" and "records" in source:
                columns = source.get("records", {}).get("columns", None)
                data = source.get(f"records", {}).get("data", None)
                if (columns is not None) and (data is not None):
                    items.append(
                        Source(
                            records=convert_to_rows_and_columns(columns, data),
                            type=source_type,
                        )
                    )
            else:
                title = source.get("title")
                items.append(
                    Source(
                        title=source.get("title"),
                        type=source.get("type", ""),
                        textSnippet=source.get("textSnippet"),
                        htmlSnippet=source.get("htmlSnippet"),
                        markdownSnippet=source.get("markdownSnippet"),
                        jsonSnippet=source.get("jsonSnippet"),
                        url=source.get("url", ""),
                        metadata=_enrich_metadata_with_tags(source.get("metadata") or {}),
                    )
                )
        formatted_source = AggregatedToolSources(
            toolCallDescription=aggregated_source.get("toolCallDescription", agent_name), items=items
        )
        if tool_call_id := _get_tool_call_id(aggregated_source):
            formatted_source["toolCallId"] = tool_call_id
        aggregated_sources.append(formatted_source)

    return aggregated_sources


def convert_back_to_columns_data(formatted: dict) -> tuple[List[str], List[List[Any]]]:
    rows = formatted["rows"]
    columns = [col["name"] for col in formatted["columns"]]

    # Extract data rows back into list of lists
    data = [[row[col] for col in columns] for row in rows]

    return columns, data


def get_records_from_artifacts(artifacts: list[dict]) -> list:
    records = []
    for index, a in enumerate(artifacts or []):
        records.append([])
        for item in a.get("parts") or []:
            if item.get("type") == "RECORDS":
                columns = item.get("records", {}).get("columns", None)
                data = item.get(f"records", {}).get("data", None)
                if (columns is not None) and (data is not None):
                    records[index].append({"columns": columns, "data": data})
    return records
