"""Lazy-loading table registry.

Knows every table/view name at startup but defers materialization
until first access via ensure_loaded().
"""

import logging
import threading
from dataclasses import dataclass, field
from typing import Callable

import pyarrow as pa
from fastapi import FastAPI

from n6k_server.formats import to_arrow

_log = logging.getLogger("n6k_server")


@dataclass
class TableSpec:
    name: str
    loader: Callable | None = None
    parent: str | None = None
    derive_fn: Callable | None = None
    loaded: bool = False
    children: list[str] = field(default_factory=list)


@dataclass
class ViewSpec:
    name: str
    sql: str
    depends_on: list[str] = field(default_factory=list)
    registered: bool = False


@dataclass
class RpcSpec:
    name: str
    fn: Callable
    parent_table: str | None = None


class TableRegistry:
    def __init__(self, app: FastAPI):
        self._app = app
        self._lock = threading.RLock()
        self._tables: dict[str, TableSpec] = {}
        self._views: dict[str, ViewSpec] = {}
        self._rpcs: dict[str, RpcSpec] = {}

    def register_source(self, name: str, fn: Callable):
        self._tables[name] = TableSpec(name=name, loader=fn)

    def register_derivation(self, parent: str, name: str, fn: Callable):
        spec = TableSpec(name=name, parent=parent, derive_fn=fn)
        self._tables[name] = spec
        if parent in self._tables:
            self._tables[parent].children.append(name)

    def register_view(self, name: str, sql: str, depends_on: list[str] | None = None):
        self._views[name] = ViewSpec(name=name, sql=sql, depends_on=depends_on or [])
        self._app.state.views[name] = sql

    def register_rpc(self, name: str, fn: Callable, parent_table: str | None = None):
        self._rpcs[name] = RpcSpec(name=name, fn=fn, parent_table=parent_table)

    def get_rpc(self, name: str) -> RpcSpec | None:
        return self._rpcs.get(name)

    @property
    def table_names(self) -> list[str]:
        return list(self._tables.keys())

    @property
    def view_names(self) -> list[str]:
        return list(self._views.keys())

    @property
    def rpc_names(self) -> list[str]:
        return list(self._rpcs.keys())

    @property
    def all_names(self) -> list[str]:
        return self.table_names + self.view_names

    def is_known(self, name: str) -> bool:
        return name in self._tables or name in self._views

    def ensure_loaded(self, name: str):
        if name in self._tables:
            self._ensure_table_loaded(name)
        elif name in self._views:
            self._ensure_view_loaded(name)

    def ensure_all_loaded(self):
        for name in self._tables:
            self._ensure_table_loaded(name)
        for name in self._views:
            self._ensure_view_loaded(name)

    def _ensure_table_loaded(self, name: str):
        spec = self._tables.get(name)
        if spec is None or spec.loaded:
            return

        with self._lock:
            if spec.loaded:
                return

            if spec.parent and spec.derive_fn is not None:
                self._ensure_table_loaded(spec.parent)
                parent_arrow = self._app.state.cached_tables[spec.parent]
                parent_df = parent_arrow.to_pandas()
                _log.info("Loading derived table '%s' from '%s'", name, spec.parent)
                arrow_table = to_arrow(spec.derive_fn(parent_df))
            elif spec.loader is not None:
                _log.info("Loading data source '%s'", name)
                arrow_table = to_arrow(spec.loader())
            else:
                return

            self._app.state.cached_tables[name] = arrow_table
            self._app.state.table_meta[name] = {"schema": arrow_table.schema}
            _register_table_in_db(self._app.state.shared_db, name, arrow_table)
            for session in self._app.state.sessions.values():
                _register_table_in_db(session.db, name, arrow_table)

            spec.loaded = True

            for child_name in spec.children:
                self._ensure_table_loaded(child_name)

    def _ensure_view_loaded(self, name: str):
        spec = self._views.get(name)
        if spec is None or spec.registered:
            return

        with self._lock:
            if spec.registered:
                return

            for dep in spec.depends_on:
                self.ensure_loaded(dep)

            _register_view_in_db(self._app.state.shared_db, name, spec.sql)
            for session in self._app.state.sessions.values():
                _register_view_in_db(session.db, name, spec.sql)

            spec.registered = True


def _register_table_in_db(db, name: str, arrow_table: pa.Table):
    db.register(name, arrow_table)
    db.execute(f'CREATE VIEW IF NOT EXISTS db."{name}" AS SELECT * FROM "{name}"')


def _register_view_in_db(db, name: str, sql: str):
    db.execute(f'CREATE OR REPLACE VIEW db."{name}" AS {sql}')
