import sqlite3
from contextlib import contextmanager
from typing import ContextManager


class Storage:
    def __init__(self):
        self.con = sqlite3.connect("database.db")

    @contextmanager
    def get_cursor(self) -> ContextManager[sqlite3.Cursor]:
        cur = self.con.cursor()
        try:
            yield cur
        finally:
            cur.close()

    def create_tables(self):
        with self.get_cursor() as cur:
            cur.execute("""
            CREATE TABLE IF NOT EXISTS ltl (
                id INTEGER PRIMARY KEY,
                doc_filepath TEXT,
                link TEXT,
                total_cost INTEGER
            );""")

    def add_link(self, doc_filepath: str, link: str, total_cost: int):
        with self.get_cursor() as cur:
            cur.execute("INSERT INTO ltl (doc_filepath, link, total_cost) VALUES (?, ?, ?)",
                        (doc_filepath, link, total_cost))
            self.con.commit()

    def is_link_exists(self, link: str) -> bool:
        with self.get_cursor() as cur:
            res = cur.execute("SELECT * FROM ltl WHERE link = ?", (link,))
            return res.fetchone()

    def is_doc_exists(self, doc_filepath: str) -> bool:
        with self.get_cursor() as cur:
            res = cur.execute("SELECT * FROM ltl WHERE doc_filepath = ?", (doc_filepath,))
            return res.fetchone()