import logging
from typing import AsyncIterator
from zoneinfo import ZoneInfo

from sqlalchemy import Text, cast, delete, func, insert, literal, outerjoin, select
from sqlalchemy.dialects.postgresql import JSON as PG_JSON
from sqlalchemy.orm import defer, selectinload
from xplan_tools.model import model_factory
from xplan_tools.model.base import BaseCollection, BaseFeature
from xplan_tools.model.orm import Feature, Refs

from xmas_app.models.orm import Document
from xmas_app.settings import get_mappings, get_settings

LOCAL_TZ = ZoneInfo("Europe/Berlin")

logger = logging.getLogger("xmas_app")


async def get_db_plans(chunk_size: int = 500) -> AsyncIterator[list[dict]]:
    """Get plan data from the database.

    Query the database for plan objects, yielded in chunks.

    Args:
        chunk_size: Number of plans per chunk.

    Yields:
        list[dict]: Chunk of dicts containing plan data.
    """
    try:
        async with get_settings().async_repo.AsyncSession() as session:
            join_stmt = outerjoin(Feature, Document, Feature.id == Document.id)

            # Correlated subquery: aggregate bereiche as JSON array
            bereich = Feature.__table__.alias("bereich")
            bereiche_subq = (
                select(
                    func.coalesce(
                        func.json_agg(
                            func.json_build_object(
                                "id",
                                cast(Refs.related_id, Text),
                                "nummer",
                                bereich.c.properties["nummer"].as_string(),
                                "geometry",
                                bereich.c.geometry.isnot(None),
                            )
                        ),
                        cast(literal("[]"), PG_JSON),
                    )
                )
                .select_from(
                    outerjoin(Refs.__table__, bereich, Refs.related_id == bereich.c.id)
                )
                .where(Refs.base_id == Feature.id, Refs.rel == "bereich")
                .correlate(Feature)
                .scalar_subquery()
                .label("bereiche")
            )

            stmt = (
                select(
                    Feature,
                    func.count(Document.id).label("document_count"),
                    func.array_agg(Document.filename)
                    .filter(Document.filetype.startswith("image"))
                    .label("document_filenames"),
                    bereiche_subq,
                )
                .select_from(join_stmt)
                .where(Feature.featuretype.regexp_match("^.*Plan$"))
                .group_by(Feature.pk)
                .execution_options(yield_per=chunk_size)
            )
            result = await session.stream(stmt)
            chunk = []
            async for feature, doc_count, images, bereiche_data in result:
                plan_name = feature.properties["name"]
                updated_local = feature.updated_at.astimezone(LOCAL_TZ)
                created_local = feature.created_at.astimezone(LOCAL_TZ)

                chunk.append(
                    {
                        "id": str(feature.id),
                        "name": plan_name,
                        "featuretype": feature.featuretype,
                        "appschema": f"{feature.appschema.upper()} {feature.version}",
                        "updated": f"{updated_local.strftime('%d.%m.%Y %H:%M')}",
                        "created": f"{created_local.strftime('%d.%m.%Y %H:%M')}",
                        "bereiche": bereiche_data
                        if isinstance(bereiche_data, list)
                        else [],
                        "docs": doc_count,
                        "images": images or [],
                    }
                )
                if len(chunk) >= chunk_size:
                    yield chunk
                    chunk = []
            if chunk:
                yield chunk

    except Exception as e:
        logger.exception("Failed to get plan data: %r", e)
        raise


def get_feature_number(plan_id: str) -> int:
    """Return the number of features of a plan."""
    with (repo := get_settings().repo).Session() as session:
        feature_number = session.scalar(
            func.count(repo.get_feature_graph(plan_id, depth=2).c.id)
        )
    return feature_number


def get_nodes(id: str) -> tuple[list[dict], str, str, str] | None:
    """Build tree nodes for ui.tree element."""

    def build_node(feature: Feature, path: str) -> dict:
        features_set.add(feature.id)
        node = {
            "id": f"{path}.{feature.featuretype}.{feature.id}",
            "label": f"{feature.featuretype} {str(feature.id)[:8]}",
            "icon": get_icon(feature),
        }
        node["children"] = build_nodes(feature)
        return node

    def build_nodes(feature: Feature) -> list[dict]:
        model = model_factory(feature.featuretype, feature.version, feature.appschema)
        node_dict = {}
        for assoc in model.get_associations():
            prop_info = model.get_property_info(assoc)
            if prop_info["assoc_info"]["source_or_target"] == (
                "target"
                if (feature.featuretype == "FP_Plan" and assoc == "bereich")
                or (feature.featuretype == "FP_Bereich" and assoc == "gehoertZuPlan")
                else "source"
            ):
                continue
            node = {
                "id": f"{feature.featuretype}.{feature.id}.{assoc}",
                "label": assoc,
                "selectable": True,
                "icon": "link",
            }
            node_dict[assoc] = node
        for ref in feature.refs:
            path = ".".join(
                [
                    feature.featuretype,
                    str(feature.id),
                    ref.rel,
                ]
            )
            node = build_node(ref.feature_inv, path)
            node_dict[ref.rel].setdefault("children", []).append(node)

        nodes = []
        for node in node_dict.values():
            if children := node.get("children"):
                node["label"] += f" [{len(children)}]"
            nodes.append(node)
        return nodes

    def get_icon(feature: Feature):
        match feature.geometry_type:
            case "polygon":
                icon = "o_space_dashboard"
            case "line":
                icon = "show_chart"
            case "point":
                icon = "o_location_on"
            case _:
                icon = "o_text_snippet"
        return icon

    with get_settings().repo.Session() as session:
        stmt = (
            select(Feature)
            .options(
                selectinload(Feature.refs)
                .selectinload(Refs.feature_inv)
                .selectinload(Feature.refs)
                .selectinload(Refs.feature_inv)
                .selectinload(Feature.refs)
            )
            .where(Feature.id == id)
        )

        feature = session.execute(stmt).scalar_one_or_none()
        # feature = session.get(Feature, id)
        if not feature or "Plan" not in feature.featuretype:
            return

        features_set = {feature.id}

        tree_nodes = [
            {
                "id": feature.id,
                "label": feature.properties["name"],
                "selectable": False,
                "children": build_nodes(feature),
            }
        ]

        return (
            tree_nodes,
            feature.featuretype,
            feature.id,
            feature.appschema,
        )


def _feature_to_table_row(feature: Feature) -> dict:
    extra_property = getattr(
        get_mappings().association_table.extra_properties, feature.appschema
    ).get(feature.featuretype)
    extra_property_label = feature.properties.get(extra_property, "N/A")
    return {
        "id": str(feature.id),
        "featuretype": feature.featuretype,
        "geometry_type": feature.geometry_type,
        "updated": f"{feature.updated_at.strftime('%d.%m.%Y')} {feature.updated_at.strftime('%H:%M')}",
        "extra_property": f"{extra_property}={extra_property_label}"
        if extra_property
        else "N/A",
    }


def _rel_inv_is_list(feature: Feature, rel_inv: str | None) -> bool:
    # if there is no inverse relation, return True
    if not rel_inv:
        return True
    model = model_factory(feature.featuretype, feature.version, feature.appschema)
    return model.get_property_info(rel_inv)["list"]


def get_ref_objects(refs: list[str]) -> list[dict]:
    """Returns a list of data objects for existing feature references to use in table rows."""
    with get_settings().repo.Session() as session:
        stmt = select(Feature).where(Feature.id.in_(refs))
        features = session.execute(stmt).scalars().all()

    return [_feature_to_table_row(feature) for feature in features]


def get_ref_candidates(
    refs: list[str], plan_id: str, model: BaseFeature, rel: str
) -> list[dict]:
    """Returns a list of data objects for potential feature references to use in table rows."""
    prop_info = model.get_property_info(rel)
    typename = prop_info["typename"]
    featuretypes = typename if isinstance(typename, list) else [typename]
    rel_inv = model.get_property_info(rel)["assoc_info"]["reverse"]
    repo = get_settings().repo

    with repo.Session() as session:
        candidates = session.scalars(
            repo.get_feature_graph(plan_id).where(
                Feature.featuretype.in_(featuretypes) & Feature.id.not_in(refs)
            )
        ).all()
        return [
            _feature_to_table_row(feature)
            for feature in candidates
            if _rel_inv_is_list(feature, rel_inv)
        ]


def get_plan_documents(
    plan_id: str, details: bool = False, with_data: bool = False
) -> list[str | dict]:
    """Retrieves binary documents stored in the db for a given plan."""
    repo = get_settings().repo
    with repo.Session() as session:
        stmt = (
            select(Document)
            .where(Document.id == plan_id)
            .options(defer(Document.filedata))
        )
        docs = []
        for doc in session.scalars(stmt).all():
            if not details:
                docs.append(f"file:///{doc.filename}")
            else:
                doc_data = {
                    "pk": doc.pk,
                    "filename": doc.filename,
                    "filetype": doc.filetype,
                }
                if with_data:
                    doc_data["filedata"] = doc.filedata
                docs.append(doc_data)
        return docs


def get_doc_by_name(plan_id: str, filename: str) -> Document:
    """Retrieve a document by filename (case-insensitive)."""
    repo = get_settings().repo
    with repo.Session() as session:
        return session.scalar(
            select(Document)
            .where(Document.id == plan_id)
            .where(func.lower(Document.filename) == func.lower(filename))
        )


def save_documents(data: list[dict]) -> None:
    """Retrieves binary documents stored in the db for a given plan."""
    repo = get_settings().repo
    with repo.Session() as session:
        session.execute(insert(Document), data)
        session.commit()


def save_plan_and_documents(plan: BaseCollection, docs: list[dict]) -> None:
    """Save a plan collection with its documents in one transaction."""
    repo = get_settings().repo
    with repo.Session() as session:
        feature_list, refs_list, refs_inv_list = repo._serialize_bulk_features(plan)
        session.execute(insert(Feature), feature_list)
        refs_list.extend(
            [ref_inv for ref_inv in refs_inv_list if ref_inv not in refs_list]
        )
        if refs_list:
            session.execute(
                insert(Refs),
                refs_list,
            )
        if docs:
            session.execute(insert(Document), docs)
        session.commit()


def delete_documents(pks: list[int]) -> None:
    """Deletes documents from the db."""
    repo = get_settings().repo
    with repo.Session() as session:
        session.execute(delete(Document).where(Document.pk.in_(pks)))
        session.commit()
