import abc

import pandas as pd
from pathlib import Path
from dodgeballtracker.typing import TrackedDetectedObject, VideoTimestamp
from sqlite3 import Connection

import pandera as pa


class DataDumper(abc.ABC):

    @abc.abstractmethod
    def write_boxes(self, obj: list[TrackedDetectedObject], ts: VideoTimestamp,
                    frame_index: int) -> None:
        pass

    @abc.abstractmethod
    def start(self, video_path: Path, hexa_hash: str) -> None:
        pass


class SQLiteDataDumper(DataDumper):
    def __init__(self, conn: Connection):
        self._conn = conn
        self._video_id = None

    def _create_tables(self):
        cur = self._conn.cursor()

        cur.executescript('''
                CREATE TABLE IF NOT EXISTS video (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    file_name TEXT NOT NULL,
                    file_sha256_hex TEXT NOT NULL,
                    
                    UNIQUE(file_sha256_hex)
                );
                ''')

        cur.executescript('''
                                CREATE TABLE IF NOT EXISTS video_frame(
                                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                                    timestamp_ DOUBLE PRECISION NOT NULL,
                                    index_ INTEGER NOT NULL,
                                    video_id INTEGER NOT NULL,
                                    
                                    FOREIGN KEY(video_id) REFERENCES video(id),
                                    UNIQUE(video_id, index_)
                                );
                                ''')

        cur.execute('''
                                CREATE TABLE IF NOT EXISTS tracked_detected_object(
                                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                                    xmin DOUBLE PRECISION NOT NULL,
                                    ymin DOUBLE PRECISION NOT NULL,
                                    xmax DOUBLE PRECISION NOT NULL,
                                    ymax DOUBLE PRECISION NOT NULL,
                                    score DOUBLE PRECISION NOT NULL,
                                    category INTEGER NOT NULL,
                                    video_frame_id INTEGER NOT NULL,
                                    
                                    FOREIGN KEY(video_frame_id) REFERENCES video_frame(id)
                                );
                                ''')

        self._conn.commit()

    def start(self, video_path: Path, hexa_hash: str) -> None:
        self._create_tables()
        cur = self._conn.cursor()
        res = cur.execute("""
        INSERT INTO video(file_name, file_sha256_hex) 
        VALUES (?, ?) 
        ON CONFLICT(file_sha256_hex) DO NOTHING
        RETURNING id
        """, (video_path.name, hexa_hash))
        self._video_id = res.fetchone()[0]
        self._conn.commit()

    @pa.check_io(df=TrackedDetectedObject)
    def write_boxes(self,
                    df: pd.DataFrame,
                    ts: VideoTimestamp,
                    frame_index: int) -> None:
        cur = self._conn.cursor()
        res = cur.execute('''
        INSERT INTO video_frame(timestamp_, index_, video_id )
        VALUES (?, ?, ?)
        ON CONFLICT(video_id, index_) DO NOTHING
        RETURNING id
        ''', (ts, frame_index, self._video_id))
        frame_id = res.fetchone()[0]
        cols = ['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category']
        cur.executemany('''
            INSERT INTO tracked_detected_object(xmin, ymin, xmax, ymax, score, category, video_frame_id)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', (list(x[cols]) + [frame_id] for x in df.iloc))
        self._conn.commit()