from typing import Any, List

from backend.models.sources import AggregatedToolSources, RAGImage, Source
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]} (p. {page_start}"
            title += ")" if page_end == page_start else f"- {page_end})"

    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 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 and len(aid.split(":")) >= 3:
        parts = aid.split(":")
        agent_name = get_llm_friendly_name(llm_id=f"{parts[1]}:{parts[2]}", project_key=parts[0])
    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 == "GENERATED_SQL_QUERY":
                items.append(Source(textSnippet=source.get("performedQuery"), type=source_type))
            elif source_type == "RECORDS" and "records" in source:
                items.append(
                    Source(
                        records=convert_to_rows_and_columns(source["records"]["columns"], source["records"]["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 {}),
                    )
                )
        aggregated_sources.append(
            AggregatedToolSources(
                toolCallDescription=aggregated_source.get("toolCallDescription", agent_name), items=items
            )
        )

    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":
                records[index].append({"columns": item["records"]["columns"], "data": item["records"]["data"]})
    return records
