import json
from collections import defaultdict
from dataclasses import dataclass
from types import ModuleType
from typing import Any, Callable, Generator, Iterable
from uuid import UUID

from json_multi_merge import merge
from sqlalchemy import Connection, Table
from sqlalchemy.sql import select

Ancestors = Iterable[tuple[UUID, str, Callable | None]]
ParentGetter = Callable[[UUID], Ancestors]


def merge_tree(
    db: Connection,
    table: Table,
    id: UUID,
    type: str,
    json_field: str,
    parents: ModuleType,
    slugs: Iterable[str] | None = None,
    slugs_only: bool = False,
    filters: tuple | None = None,
    buff_query: Callable | None = None,
    debug: str | None = None,
    callback: Callable | None = None,
    parents_callback: Callable | None = None,
) -> dict:
    return TreeMerger(
        db=db,
        table=table,
        json_field=json_field,
        parents=parents,
        filters=filters,
        buff_query=buff_query,
        callback=callback,
        parents_callback=parents_callback,
        debug=debug,
    ).merge(
        resource_id=id,
        resource_type=type,
        slugs=slugs,
        slugs_only=slugs_only,
    )


@dataclass
class TreeMerger:
    db: Connection
    table: Table
    json_field: str
    parents: ModuleType
    filters: tuple | None = None
    buff_query: Callable | None = None
    callback: Callable | None = None
    parents_callback: Callable | None = None
    debug: str | None = None

    def merge(
        self,
        resource_id: UUID,
        resource_type: str,
        slugs: Iterable[str] | None = None,
        slugs_only: bool = False,
    ) -> dict:
        """Take a resource ID and return the merged json object for that page.

        The merged json object is any json saved for that resource, merged into any resources saved
        for its ancestor resources, all the way up the hierarchy.
        """
        # Get a generator that will yield the IDs of a resource's immediate ancestors
        self.resource_id = resource_id
        self.resource_type = resource_type
        self.slugs = slugs
        self.slugs_only = slugs_only
        parent_getter = getattr(self.parents, resource_type, None)

        # Get the IDs of all ancestor resources in the hierarchy, ignoring "inherits"
        # Recursively nested list of IDs of each resource and its ancestors
        resources = self.get_resources(resource_id, resource_type, parent_getter)

        # Get the records for those resources and any requested slugs, in order
        # A dict of resource_id: tuple(record, ...)
        records = self.get_records(resources)

        # Get just the JSON objects that should be merged by checking "inherits" in the records
        # as we recurse through the resources
        json_objects = self.get_json_objects(resources, records)

        # Merge those json objects and return the result
        return merge(*json_objects)

    def get_resources(
        self,
        resource_id: UUID,
        resource_type: str,
        parent_getter: ParentGetter | None,
    ) -> tuple[UUID, str, tuple]:
        """Recursively collect all resource IDs and types in the hierarchy.

        Returns a nested tuple structure like:
        (resource_id, resource_type, (parent1_tuple, parent2_tuple, ...))
        where each parent tuple has the same structure.
        """
        if parent_getter is None:
            return resource_id, resource_type, ()
        return resource_id, resource_type, tuple(self.get_parents(resource_id, parent_getter))

    def get_parents(self, resource_id: UUID, parent_getter: ParentGetter) -> Generator:
        for parent_id, parent_type, grandparent_getter in parent_getter(resource_id):
            yield self.get_resources(parent_id, parent_type, grandparent_getter)

    def get_records(
        self,
        resources: tuple[UUID, str, tuple],
    ) -> dict[str, list]:
        """Fetch all records for resources and requested slugs in one DB call.

        Returns a dict mapping resource_id (as string) to list of records.
        """
        # Flatten to get all unique resource IDs
        resource_ids = list(self.flatten_resource_ids(resources))

        # Build query
        C = self.table.columns
        query = select(self.table)

        # Filter by resource IDs and slugs
        filters: tuple = (C.resource_id.in_(resource_ids),)
        if self.slugs:
            filters = filters + ((C.slug.is_(None) | C.slug.in_(self.slugs)),)
        else:
            filters = filters + (C.slug.is_(None),)

        # Add user filters
        if self.filters is not None:
            filters += self.filters

        # Hook for altering query and filters before executing
        if callable(self.buff_query):
            query, filters = self.buff_query(C, query, filters)

        # Execute query
        query = query.where(*filters).order_by(C.resource_id, C.created_at)
        results = self.db.execute(query)

        # Organize by resource_id (normalize to string for consistent lookups)
        records_by_id = defaultdict(list)
        for record in results:
            resource_id = str(record.resource_id)
            records_by_id[resource_id].append(record)

        return records_by_id

    def flatten_resource_ids(
        self,
        resources: tuple[UUID, str, tuple],
    ) -> Iterable[UUID]:
        """Recursively flatten the nested resource structure to get all IDs."""
        resource_id, resource_type, parents = resources
        yield resource_id
        for parent in parents:
            yield from self.flatten_resource_ids(parent)

    def get_json_objects(
        self,
        resources: tuple[UUID, str, tuple],
        records: dict[str, list],
    ) -> Iterable[dict]:
        """Traverse the resource tree and emit JSON objects in merge order.

        Checks the 'inherits' field to skip parent branches.
        """
        resource_id, resource_type, parents = resources

        # Get this resource's records (convert to string for consistent lookup)
        resource_records = records.get(str(resource_id), [])

        # Find the resource record (no slug) and slug records
        # Take the first record for each (oldest by created_at) to match main branch behavior
        resource_record = None
        slug_records = dict.fromkeys(self.slugs) if self.slugs else {}
        for record in reversed(resource_records):
            if record.slug:
                slug_records[record.slug] = record
            else:
                resource_record = record

        # Check if we should inherit from parents
        inherits = resource_record.inherits if resource_record else True

        # Recurse to parents if inherits is True
        if inherits and parents:
            for parent in parents:
                yield from self.get_json_objects(parent, records)

        # Call parents_callback for every resource in the hierarchy
        if callable(self.parents_callback):
            self.parents_callback(resource_id, resource_type)

        # Yield this resource's JSON (skip when slugs_only)
        if resource_record and not self.slugs_only:
            yield self.record_json(resource_record)

        # Yield slug JSONs in the order specified by self.slugs
        if self.slugs:
            for record in slug_records.values():
                if record:
                    yield self.record_json(record)

    def record_json(self, record: Any) -> dict:
        if callable(self.callback):
            self.callback(record)

        json_data = getattr(record, self.json_field)

        if isinstance(json_data, str):
            json_data = json.loads(json_data)

        if self.debug is not None:
            self.add_debug_info(json_data, record)

        return json_data

    def add_debug_info(self, json_data, record):
        new_items = []
        for k, v in json_data.items():
            if isinstance(v, dict):
                self.add_debug_info(v, record)
            else:
                new_items.append(self.debug_info(k, v, record))
        json_data.update(new_items)

    def debug_info(self, k: str, v: Any, record: Any) -> tuple[str, dict | list]:
        if self.debug == 'history':
            return k + '-history', [self.from_dict(record) | {'value': v}]
        if self.debug == 'annotate':
            return k + '-from', self.from_dict(record)
        return k, self.from_dict(record)

    def from_dict(self, record: Any) -> dict:
        return dict(self.from_pairs(record))

    def from_pairs(self, record: Any) -> Iterable[tuple[str, str]]:
        yield 'id', str(record.resource_id)
        yield 'type', record.resource_type,
        yield self.json_field + '_id', str(record.id)
        if record.slug:
            yield 'slug', record.slug
        if record.draft_set:
            yield 'draft_set', record.draft_set
