"""Flat component-per-row table storage model.

Every component across all OHRs lives in a single "long" DataFrame.
OHR structure is reconstructed via graph-like views on the table.

Flat columns (tone_function, etc.) enable fast SQL queries.  The full
serialized properties JSON is preserved alongside for lossless round-trip.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import Any

import pandas as pd

from flexohr.core.component import AbsentComponent, Component, UndefinedComponent
from flexohr.core.ohr import (
    OHR,
    _component_from_dict,
    _component_to_dict,
    _properties_from_json,
    _properties_to_json,
)
from flexohr.harmony.enums import (
    EMPTY_PROPS,
    get_value_type_name,
)

logger = logging.getLogger(__name__)


@dataclass
class _OHRMeta:
    """Lightweight OHR-level metadata stored alongside the component table."""

    ohr_id: int
    parent_ohr_id: int | None
    body_position: int | None  # position in parent body (-1 = not a body child)
    ordered: bool
    horizontal: bool
    properties_json: str  # serialized OHR-level properties
    reference_ohr_id: int | None


class FlatComponentTable:
    """Flat component-per-row storage for OHR trees.

    Each row is one component.  OHR structure is metadata that allows
    reconstruction of the original tree.  Properties are stored both as
    flat query columns (tone_function) and as a full JSON blob for lossless
    round-trip.
    """

    def __init__(self) -> None:
        self._rows: list[dict[str, Any]] = []
        self._ohr_meta: dict[int, _OHRMeta] = {}
        self._next_component_id = 0
        self._next_ohr_id = 0

    # ── Ingest ──

    def add_ohr(
        self,
        ohr: OHR,
        parent_ohr_id: int | None = None,
        body_position: int | None = None,
    ) -> int:
        """Flatten *ohr* into the table, returning its ohr_id."""
        ohr_id = self._next_ohr_id
        self._next_ohr_id += 1

        # Handle reference_ohr recursively (not a body child, so body_position=None)
        ref_ohr_id: int | None = None
        if ohr.reference_ohr is not None:
            ref_ohr_id = self.add_ohr(ohr.reference_ohr, parent_ohr_id=ohr_id, body_position=None)

        # Store OHR-level metadata
        ohr_props_json = json.dumps(_properties_to_json(ohr.properties) if ohr.properties else [])
        self._ohr_meta[ohr_id] = _OHRMeta(
            ohr_id=ohr_id,
            parent_ohr_id=parent_ohr_id,
            body_position=body_position,
            ordered=ohr.ordered,
            horizontal=ohr.horizontal,
            properties_json=ohr_props_json,
            reference_ohr_id=ref_ohr_id,
        )

        # Reference component (position=-1 sentinel)
        if ohr.reference_component is not None:
            self._add_component_row(ohr.reference_component, ohr_id, position=-1, is_reference=True)

        # Body items — both Components and nested OHRs get a position
        for pos, item in enumerate(ohr.body):
            if isinstance(item, OHR):
                self.add_ohr(item, parent_ohr_id=ohr_id, body_position=pos)
            else:
                self._add_component_row(item, ohr_id, position=pos, is_reference=False)

        return ohr_id

    def add_ohrs(self, ohrs: list[OHR]) -> list[int]:
        """Add multiple top-level OHRs, returning their ohr_ids."""
        return [self.add_ohr(ohr) for ohr in ohrs]

    def _add_component_row(
        self,
        comp: Component,
        ohr_id: int,
        position: int,
        is_reference: bool,
    ) -> None:
        cid = self._next_component_id
        self._next_component_id += 1

        # Serialize the component to dict for lossless storage
        comp_dict = _component_to_dict(comp)
        comp_json = json.dumps(comp_dict)

        # Extract flat query columns
        value_raw: Any = None
        value_type: str | None = None
        if comp.value is not None:
            vt_name = get_value_type_name(comp.value)
            if vt_name is not None:
                value_type = vt_name
                if isinstance(comp.value, tuple):
                    value_raw = [int(comp.value[0]), int(comp.value[1])]
                else:
                    value_raw = int(comp.value)
            else:
                value_raw = comp.value

        tone_function = _get_prop_str(comp, "tone_function")

        comp_type = "component"
        if isinstance(comp, AbsentComponent):
            comp_type = "absent_component"
        elif isinstance(comp, UndefinedComponent):
            comp_type = "undefined_component"

        self._rows.append(
            {
                "component_id": cid,
                "ohr_id": ohr_id,
                "position": position,
                "is_reference": is_reference,
                "value": value_raw,
                "value_type": value_type,
                "component_type": comp_type,
                "status": comp.status.value,
                "tone_function": tone_function,
                "component_json": comp_json,  # lossless backup
            }
        )

    # ── DataFrame access ──

    def to_dataframe(self) -> pd.DataFrame:
        """Return the flat component table as a DataFrame."""
        if not self._rows:
            return pd.DataFrame(
                columns=[
                    "component_id",
                    "ohr_id",
                    "position",
                    "is_reference",
                    "value",
                    "value_type",
                    "component_type",
                    "status",
                    "tone_function",
                    "component_json",
                ]
            )
        return pd.DataFrame(self._rows)

    def ohr_meta_dataframe(self) -> pd.DataFrame:
        """Return OHR-level metadata as a DataFrame."""
        rows = []
        for meta in self._ohr_meta.values():
            rows.append(
                {
                    "ohr_id": meta.ohr_id,
                    "parent_ohr_id": meta.parent_ohr_id,
                    "body_position": meta.body_position,
                    "ordered": meta.ordered,
                    "horizontal": meta.horizontal,
                    "reference_ohr_id": meta.reference_ohr_id,
                    "properties_json": meta.properties_json,
                }
            )
        return (
            pd.DataFrame(rows)
            if rows
            else pd.DataFrame(
                columns=[
                    "ohr_id",
                    "parent_ohr_id",
                    "body_position",
                    "ordered",
                    "horizontal",
                    "reference_ohr_id",
                    "properties_json",
                ]
            )
        )

    # ── Reconstruction ──

    def reconstruct_ohr(self, ohr_id: int, _df: pd.DataFrame | None = None) -> OHR:
        """Reconstruct a full OHR tree from the flat table."""
        meta = self._ohr_meta[ohr_id]
        if _df is None:
            _df = self.to_dataframe()

        # Reference component
        ref_rows = _df[(_df["ohr_id"] == ohr_id) & (_df["is_reference"])]
        ref_comp: Component | None = None
        if len(ref_rows) > 0:
            ref_comp = _component_from_dict(json.loads(ref_rows.iloc[0]["component_json"]))

        # Reference OHR
        ref_ohr: OHR | None = None
        if meta.reference_ohr_id is not None:
            ref_ohr = self.reconstruct_ohr(meta.reference_ohr_id, _df=_df)

        # Body: interleave components and child OHRs by position
        body_comps = _df[(_df["ohr_id"] == ohr_id) & (~_df["is_reference"])]

        # Child OHRs that are body items (have body_position != None)
        child_body_ohrs: list[tuple[int, int]] = []  # (position, child_ohr_id)
        for m in self._ohr_meta.values():
            if m.parent_ohr_id == ohr_id and m.ohr_id != (meta.reference_ohr_id or -1):
                if m.body_position is not None:
                    child_body_ohrs.append((m.body_position, m.ohr_id))

        # Build position -> item mapping
        body_items: dict[int, Any] = {}
        for _, row in body_comps.iterrows():
            pos = row["position"]
            body_items[pos] = _component_from_dict(json.loads(row["component_json"]))

        for pos, child_id in child_body_ohrs:
            body_items[pos] = self.reconstruct_ohr(child_id, _df=_df)

        # Assemble body in position order
        body = [body_items[k] for k in sorted(body_items.keys())]

        # Reconstruct OHR-level properties
        ohr_props_data = json.loads(meta.properties_json)
        ohr_props = _properties_from_json(ohr_props_data) if ohr_props_data else EMPTY_PROPS

        return OHR(
            reference_component=ref_comp,
            reference_ohr=ref_ohr,
            body=tuple(body),
            ordered=meta.ordered,
            horizontal=meta.horizontal,
            properties=ohr_props,
        )

    def reconstruct_all_top_level(self) -> list[OHR]:
        """Reconstruct all top-level OHRs (those with no parent)."""
        # Pre-build indices for fast reconstruction
        comp_index, child_index = self._build_indices()
        top_ids = [m.ohr_id for m in self._ohr_meta.values() if m.parent_ohr_id is None]
        return [self._reconstruct_fast(oid, comp_index, child_index) for oid in sorted(top_ids)]

    def _build_indices(self) -> tuple[dict, dict]:
        """Build lookup indices for fast reconstruction.

        Returns (comp_index, child_index) where:
        - comp_index[ohr_id] = list of (position, is_reference, component_json)
        - child_index[parent_ohr_id] = list of (body_position, child_ohr_id)
        """
        comp_index: dict[int, list[tuple[int, bool, str]]] = {}
        for row in self._rows:
            oid = row["ohr_id"]
            comp_index.setdefault(oid, []).append((row["position"], row["is_reference"], row["component_json"]))

        child_index: dict[int, list[tuple[int, int]]] = {}
        for meta in self._ohr_meta.values():
            if meta.parent_ohr_id is not None and meta.body_position is not None:
                parent = meta.parent_ohr_id
                # Exclude reference_ohr children
                parent_meta = self._ohr_meta.get(parent)
                if parent_meta and parent_meta.reference_ohr_id == meta.ohr_id:
                    continue
                child_index.setdefault(parent, []).append((meta.body_position, meta.ohr_id))

        return comp_index, child_index

    def _reconstruct_fast(
        self,
        ohr_id: int,
        comp_index: dict,
        child_index: dict,
    ) -> OHR:
        """Fast OHR reconstruction using pre-built indices."""
        meta = self._ohr_meta[ohr_id]

        # Reference component
        ref_comp: Component | None = None
        body_items: dict[int, Any] = {}

        for pos, is_ref, cjson in comp_index.get(ohr_id, []):
            comp = _component_from_dict(json.loads(cjson))
            if is_ref:
                ref_comp = comp
            else:
                body_items[pos] = comp

        # Reference OHR
        ref_ohr: OHR | None = None
        if meta.reference_ohr_id is not None:
            ref_ohr = self._reconstruct_fast(meta.reference_ohr_id, comp_index, child_index)

        # Child OHRs in body
        for pos, child_id in child_index.get(ohr_id, []):
            body_items[pos] = self._reconstruct_fast(child_id, comp_index, child_index)

        body = [body_items[k] for k in sorted(body_items.keys())]

        ohr_props_data = json.loads(meta.properties_json)
        ohr_props = _properties_from_json(ohr_props_data) if ohr_props_data else EMPTY_PROPS

        return OHR(
            reference_component=ref_comp,
            reference_ohr=ref_ohr,
            body=tuple(body),
            ordered=meta.ordered,
            horizontal=meta.horizontal,
            properties=ohr_props,
        )

    # ── Queries ──

    def query_components_by_property(self, key: str, value: Any) -> pd.DataFrame:
        """Find all components matching a flat property column value."""
        df = self.to_dataframe()
        if key in df.columns:
            return df[df[key] == value]
        return df.iloc[0:0]

    def query_ohrs_by_property(self, key: str, value: Any) -> list[int]:
        """Find all OHR IDs whose metadata properties match."""
        results = []
        for meta in self._ohr_meta.values():
            props_data = json.loads(meta.properties_json)
            for prop in props_data:
                if prop.get("type") == "enum" and prop.get("value") == value:
                    results.append(meta.ohr_id)
                    break
                if prop.get("type") == "key_value" and prop.get("key") == key and prop.get("value") == value:
                    results.append(meta.ohr_id)
                    break
        return results

    def count_components(self) -> int:
        """Total number of component rows."""
        return len(self._rows)

    def count_ohrs(self) -> int:
        """Total number of OHR nodes."""
        return len(self._ohr_meta)

    @property
    def memory_bytes(self) -> int:
        """Approximate memory usage."""
        import sys

        df = self.to_dataframe()
        return df.memory_usage(deep=True).sum() + sys.getsizeof(self._ohr_meta)


def _get_prop_str(comp: Component, key: str) -> str | None:
    """Extract a string property value from a component."""
    val = comp.get_property(key)
    if val is None:
        return None
    if hasattr(val, "value"):
        return val.value
    return str(val)
