# Copyright 2020-2026 Ternaris
# SPDX-License-Identifier: Apache-2.0
"""Sqlite3 storage."""

from __future__ import annotations

import sqlite3
import sys
from pathlib import Path
from typing import TYPE_CHECKING, cast

import apsw

from rosbags.interfaces import (
    Connection,
    ConnectionExtRosbag2,
    MessageDefinition,
    MessageDefinitionFormat,
)
from rosbags.typesys.msg import get_types_from_msg
from rosbags.typesys.store import Typestore

from .enums import CompressionMode
from .errors import ReaderError, WriterError
from .metadata import ReaderMetadata, parse_qos

if sys.version_info >= (3, 12):
    from typing import override
else:  # pragma: no cover
    from typing_extensions import override

if TYPE_CHECKING:
    from collections.abc import Generator, Iterable

    from rosbags.interfaces.typing import RPath


def make_vfs(rpath: RPath) -> apsw.VFS:  # pragma: no cover
    """Create sqlite vfs driver for path."""

    class VFS(apsw.VFS):
        class VFSFile(apsw.VFSFile):
            def __init__(
                self,
                _vfs: VFS,
                _path: str | apsw.URIFilename | None,
                _flags: list[int],
            ) -> None:
                self._path = rpath.open('rb')

            @override
            def xDeviceCharacteristics(self) -> int:
                return 0

            @override
            def xSectorSize(self) -> int:
                return 0

            @override
            def xFileControl(self, op: int, ptr: int) -> bool:
                if op in (14, 15, 30):
                    return False
                return super().xFileControl(op, ptr)

            @override
            def xFileSize(self) -> int:
                return rpath.stat().st_size

            @override
            def xRead(self, amount: int, offset: int) -> bytes:
                self._path.seek(offset)
                return self._path.read(amount)

            @override
            def xClose(self) -> None:
                self._path.close()

        @override
        def __init__(self, name: str = 'rpathvfs') -> None:
            super().__init__(name, '')

        @override
        def xOpen(
            self,
            name: str | apsw.URIFilename | None,
            flags: list[int],
        ) -> apsw.VFSFile:
            if name is None:
                return super().xOpen(name, flags)

            write_flags = {
                apsw.SQLITE_OPEN_READWRITE,
                apsw.SQLITE_OPEN_CREATE,
                apsw.SQLITE_OPEN_DELETEONCLOSE,
                apsw.SQLITE_OPEN_EXCLUSIVE,
                apsw.SQLITE_OPEN_AUTOPROXY,
                apsw.SQLITE_OPEN_WAL,
                apsw.SQLITE_OPEN_SUPER_JOURNAL,
            }
            if any(x in write_flags for x in flags):
                msg = 'Only read-only access supported'
                raise RuntimeError(msg)

            return self.VFSFile(self, name, flags)

    return VFS()


class Sqlite3Reader:
    """Sqlite3 storage reader."""

    def __init__(self, path: RPath) -> None:
        """Set up storage reader.

        Args:
            path: Paths of storage files.

        """
        self.path = path
        self.dbconn: apsw.Connection | None = None
        self.schema = 0
        self.msgtypes: list[dict[str, str]] = []
        self.connections: list[Connection] = []
        self.metadata = ReaderMetadata(0, 2**63 - 1, 0, 0, None, None, None, None)

    def open(self) -> None:
        """Open sqlite3 storage file."""
        if isinstance(self.path, Path):
            vfs = None
        else:  # pragma: no cover
            vfs = 'rpathvfs'
            _vfs = make_vfs(self.path)

        conn = apsw.Connection(
            f'file:{self.path}?immutable=1',
            flags=apsw.SQLITE_OPEN_READONLY | apsw.SQLITE_OPEN_URI,
            vfs=vfs,
        )

        cur = conn.cursor()
        _ = cur.execute(
            (
                'SELECT count(*) FROM sqlite_master '
                'WHERE type="table" AND name IN ("messages", "topics")'
            ),
        )
        if (x := cur.fetchone()) is None or x[0] != 2:
            msg = f'Cannot open database {self.path} or database missing tables.'
            conn.close()
            raise ReaderError(msg)

        self.dbconn = conn

        cur = conn.cursor()
        if cur.execute('PRAGMA table_info(schema)').fetchall():
            schema: int
            (schema,) = cur.execute('SELECT schema_version FROM schema').fetchone() or (-1,)
        elif any(
            x[1] == 'offered_qos_profiles'
            for x in cast('Iterable[tuple[str, str]]', cur.execute('PRAGMA table_info(topics)'))
        ):
            schema = 2
        else:
            schema = 1

        if schema >= 4:
            msgtypes: list[dict[str, str]] = [
                {
                    'name': x[0],
                    'encoding': x[1],
                    'msgdef': x[2],
                    'digest': x[3],
                }
                for x in cast(
                    'Iterable[tuple[str, str, str, str]]',
                    cur.execute(
                        (
                            'SELECT topic_type, encoding, encoded_message_definition,'
                            ' type_description_hash FROM message_definitions ORDER BY id'
                        ),
                    ),
                )
            ]
            for typ in msgtypes:
                assert typ['encoding'] == 'ros2msg'
                types = get_types_from_msg(typ['msgdef'], typ['name'])

                store = Typestore()
                store.register(types)

                assert typ['digest'] == store.hash_rihs01(
                    typ['name'],
                ), f'Failed to parse {typ["name"]}'
        else:
            msgtypes = []

        def get_msgdef(name: str) -> MessageDefinition:
            """Get message definition for name."""
            fmtmap = {
                'ros2msg': MessageDefinitionFormat.MSG,
                'ros2idl': MessageDefinitionFormat.IDL,
            }
            if msgtype := next((x for x in msgtypes if x['name'] == name), None):
                return MessageDefinition(fmtmap[msgtype['encoding']], msgtype['msgdef'])
            return MessageDefinition(MessageDefinitionFormat.NONE, '')

        if schema >= 4:
            connections = [
                Connection(
                    cid,
                    topic,
                    msgtype,
                    get_msgdef(msgtype),
                    digest,
                    msgcount,
                    ConnectionExtRosbag2(serialization_format, parse_qos(offered_qos_profiles)),
                    self,
                )
                for (
                    cid,
                    topic,
                    msgtype,
                    msgcount,
                    serialization_format,
                    offered_qos_profiles,
                    digest,
                ) in cast(
                    'Iterable[tuple[int, str, str, int, str, str, str]]',
                    cur.execute(
                        (
                            'SELECT topics.id, name, type, count(messages.id), '
                            'serialization_format, offered_qos_profiles, type_description_hash '
                            'FROM topics LEFT JOIN messages ON topics.id == messages.topic_id '
                            'GROUP BY topics.id ORDER BY topics.id'
                        ),
                    ),
                )
            ]
        elif schema >= 2:
            connections = [
                Connection(
                    cid,
                    topic,
                    msgtype,
                    get_msgdef(msgtype),
                    '',
                    msgcount,
                    ConnectionExtRosbag2(serialization_format, parse_qos(offered_qos_profiles)),
                    self,
                )
                for (
                    cid,
                    topic,
                    msgtype,
                    msgcount,
                    serialization_format,
                    offered_qos_profiles,
                ) in cast(
                    'Iterable[tuple[int, str, str, int, str, str]]',
                    cur.execute(
                        (
                            'SELECT topics.id, name, type, count(*), '
                            'serialization_format, offered_qos_profiles '
                            'FROM topics JOIN messages ON topics.id == messages.topic_id '
                            'GROUP BY topics.id ORDER BY topics.id'
                        ),
                    ),
                )
            ]
        else:
            connections = [
                Connection(
                    cid,
                    topic,
                    msgtype,
                    get_msgdef(msgtype),
                    '',
                    msgcount,
                    ConnectionExtRosbag2(serialization_format, []),
                    self,
                )
                for (
                    cid,
                    topic,
                    msgtype,
                    msgcount,
                    serialization_format,
                ) in cast(
                    'Iterable[tuple[int, str, str, int, str]]',
                    cur.execute(
                        (
                            'SELECT topics.id, name, type, count(*), '
                            'serialization_format '
                            'FROM topics JOIN messages ON topics.id == messages.topic_id '
                            'GROUP BY topics.id ORDER BY topics.id'
                        ),
                    ),
                )
            ]

        self.schema = schema
        self.msgtypes = msgtypes
        self.connections = connections

        ((start_time, end_time, msgcount),) = cast(
            'Iterable[tuple[int, int, int]]',
            cur.execute(
                'SELECT MIN(timestamp), MAX(timestamp) + 1, COUNT(*) FROM messages',
            ),
        )
        self.metadata = ReaderMetadata(
            end_time - start_time if start_time is not None else 0,
            start_time if start_time is not None else 2**63 - 1,
            end_time if end_time is not None else 0,
            msgcount,
            None,
            None,
            None,
            None,
        )

    def close(self) -> None:
        """Close rosbag2."""
        assert self.dbconn
        self.dbconn.close()
        self.dbconn = None

    def messages(
        self,
        connections: Iterable[Connection],
        start: int | None = None,
        stop: int | None = None,
    ) -> Generator[tuple[Connection, int, bytes], None, None]:
        """Read messages from bag.

        Args:
            connections: Iterable with connections to filter for. An empty
                iterable disables filtering on connections.
            start: Yield only messages at or after this timestamp (ns).
            stop: Yield only messages before this timestamp (ns).

        Yields:
            tuples of connection, timestamp (ns), and rawdata.

        Raises:
            ReaderError: Bag not open.

        """
        assert self.dbconn

        query = [
            'SELECT topics.id,messages.timestamp,messages.data',
            'FROM messages JOIN topics ON messages.topic_id=topics.id',
        ]
        args: list[apsw.Binding] = []
        clause = 'WHERE'

        topics = tuple({x.topic for x in connections})
        query.append(f'{clause} topics.name IN ({",".join("?" for _ in topics)})')
        args += topics
        clause = 'AND'

        if start is not None:
            query.append(f'{clause} messages.timestamp >= ?')
            args.append(start)
            clause = 'AND'

        if stop is not None:
            query.append(f'{clause} messages.timestamp < ?')
            args.append(stop)
            clause = 'AND'

        query.append('ORDER BY timestamp')
        querystr = ' '.join(query)

        connmap = {x.id: x for x in self.connections}

        cur = cast('Iterable[tuple[int, int, bytes]]', self.dbconn.execute(querystr, args))

        for cid, timestamp, data in cur:
            yield connmap[cid], timestamp, data


class Sqlite3Writer:
    """Sqlite3 Storage Writer."""

    SQLITE_SCHEMA = """
    CREATE TABLE schema(
      schema_version INTEGER PRIMARY KEY,
      ros_distro TEXT NOT NULL
    );
    CREATE TABLE metadata(
      id INTEGER PRIMARY KEY,
      metadata_version INTEGER NOT NULL,
      metadata TEXT NOT NULL
    );
    CREATE TABLE topics(
      id INTEGER PRIMARY KEY,
      name TEXT NOT NULL,
      type TEXT NOT NULL,
      serialization_format TEXT NOT NULL,
      offered_qos_profiles TEXT NOT NULL,
      type_description_hash TEXT NOT NULL
    );
    CREATE TABLE message_definitions(
      id INTEGER PRIMARY KEY,
      topic_type TEXT NOT NULL,
      encoding TEXT NOT NULL,
      encoded_message_definition TEXT NOT NULL,
      type_description_hash TEXT NOT NULL
    );
    CREATE TABLE messages(
      id INTEGER PRIMARY KEY,
      topic_id INTEGER NOT NULL,
      timestamp INTEGER NOT NULL,
      data BLOB NOT NULL
    );
    CREATE INDEX timestamp_idx ON messages (timestamp ASC);
    INSERT INTO schema(schema_version, ros_distro) VALUES (4, 'rosbags');
    """

    def __init__(self, path: Path, compression: CompressionMode) -> None:
        """Initialize sqlite3 storage."""
        if compression == CompressionMode.STORAGE:
            msg = 'SQLITE3 writer does not support storage-side compression.'
            raise WriterError(msg)

        self.path = path / f'{path.name}.db3'
        self.conn = sqlite3.connect(f'file:{self.path}', uri=True)
        _ = self.conn.executescript(self.SQLITE_SCHEMA)
        self.cursor = self.conn.cursor()

    def add_msgtype(self, connection: Connection) -> None:
        """Add a msgtype.

        Args:
            connection: Connection.

        """
        _ = self.cursor.execute(
            (
                'INSERT INTO message_definitions (topic_type, encoding,'
                ' encoded_message_definition, type_description_hash) VALUES(?, ?, ?, ?)'
            ),
            (
                connection.msgtype,
                'ros2msg' if connection.msgdef.format == MessageDefinitionFormat.MSG else 'ros2idl',
                connection.msgdef.data,
                connection.digest,
            ),
        )

    def add_connection(self, connection: Connection, offered_qos_profiles: str) -> None:
        """Add a connection.

        Args:
            connection: Connection.
            offered_qos_profiles: Serialized QoS profiles.

        """
        _ = self.cursor.execute(
            'INSERT INTO topics VALUES(?, ?, ?, ?, ?, ?)',
            (
                connection.id,
                connection.topic,
                connection.msgtype,
                cast('ConnectionExtRosbag2', connection.ext).serialization_format,
                offered_qos_profiles,
                connection.digest,
            ),
        )

    def write(self, connection: Connection, timestamp: int, data: bytes | memoryview) -> None:
        """Write message to rosbag2.

        Args:
            connection: Connection to write message to.
            timestamp: Message timestamp (ns).
            data: Serialized message data.

        """
        _ = self.cursor.execute(
            'INSERT INTO messages (topic_id, timestamp, data) VALUES(?, ?, ?)',
            (connection.id, timestamp, data),
        )

    def close(self, version: int, metadata: str) -> None:
        """Close rosbag2 after writing.

        Closes open database transactions and writes metadata.yaml.

        """
        self.cursor.execute(
            'INSERT INTO metadata(metadata_version, metadata) VALUES(?, ?)',
            (version, metadata),
        )

        self.conn.commit()
        _ = self.conn.execute('PRAGMA optimize')
        self.conn.close()
