import sqlite3
from functools import wraps

# Type alias
Cursor = sqlite3.Cursor

# Configs
DB_FILE = "./traceroute.db"


class Database:
    def __init__(self):
        self.db_file = DB_FILE
        self.conn = sqlite3.connect(self.db_file, check_same_thread=False)

        # Return fetch() data as Row objects, instead of tuples.
        self.conn.row_factory = sqlite3.Row

        self.cursor = self.conn.cursor()

    def create_tables(self):
        self.cursor.executescript(
            """
            CREATE TABLE IF NOT EXISTS Traces (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                created TEXT NOT NULL,
                origin TEXT NOT NULL,
                target TEXT NOT NULL,
                unparsed TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS Nodes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                created TEXT NOT NULL,
                name TEXT,
                ip TEXT,
                latency_ms REAL NOT NULL,
            );

            CREATE TABLE IF NOT EXISTS Links (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                created TEXT NOT NULL,
                source TEXT NOT NULL,
                target TEXT NOT NULL,
                latency_ms REAL NOT NULL,
            );
        """
        )

    def end(self):
        """Always call this after you're done with the connection / request."""

        self.conn.commit()
        self.conn.close()

    # TODO
    def list_traces(self):
        # TODO: time filter
        result = []

        self.cursor.execute("SELECT * FROM Traces")
        traces = self.cursor.fetchall()

        for t in traces:
            trace = dict(t)

            self.cursor.execute(
                """
                SELECT number, name, ip, latency, link_latency
                FROM Hops
                WHERE trace_id = ?
                ORDER BY number ASC
            """,
                (trace["id"],),
            )
            hops = self.cursor.fetchall()
            trace["hops"] = hops

            result.append(trace)

        return result

    # TODO
    def create_trace(self, trace):
        self.cursor.execute(
            "INSERT OR IGNORE INTO Traces (created, origin, target) VALUES (?, ?, ?)",
            (trace["created"], trace["origin"], trace["target"]),
        )
        trace_id = self.cursor.lastrowid

        for hop in trace["hops"]:
            self.cursor.execute(
                "INSERT OR IGNORE INTO Hops (trace_id, created, number, name, ip, latency, link_latency) VALUES (?, ?, ?, ?, ?, ?, ?)",
                (
                    trace_id,
                    hop["created"],  # TODO: trace.created
                    hop["number"],
                    hop["name"],
                    hop["ip"],
                    hop["latency"],
                    hop["link_latency"],
                ),
            )

        return self.cursor.fetchone()

    def create_hop(self, name, ip, latency):
        self.cursor.execute(
            "INSERT INTO Hops (name, ip, latency) VALUES (?, ?, ?)",
            (name, ip, latency),
        )


def ensure_table_setup():
    db = Database()
    db.create_tables()
    db.end()