from typing import Optional, Dict, List, Any, Union
import datetime as dt
from sqlmodel import Field, Session, SQLModel, create_engine, select
import threading as th
import queue

# ~~~ Database ~~~~~~~~~~~~~~~

class Database:
    def __init__(self, uri: str):
        self.engine = create_engine(uri)
        SQLModel.metadata.create_all(self.engine)

    def create_all(self, items: List[SQLModel]):
        with Session(self.engine) as session:
            for item in items:
                session.add(item)
            session.commit()

    def get_by_id(self, id: Union[str, int], model: SQLModel):
        with Session(self.engine) as session:
            stmt = select(model).where(model.id == id)
            return session.exec(stmt).first()
    
    def get_by_field(self, key: str, value: Any, model: SQLModel):
        stmt = select(model).where(getattr(model, key) == value)
        print(stmt)
        return self.exec(stmt)
    
    def exec(self, stmt: str, params = {}):
        with Session(self.engine) as session:
            return session.exec(stmt, params=params).all()

class DatabaseWorker(th.Thread):
    def __init__(self, 
            uri: str,
            queue: queue.Queue, 
            batch: int = None,
            timeout: int = 10
        ):
        super().__init__()
        self.q = queue
        self.db = None
        self.uri = uri
        self.timeout = timeout
        self.batch = batch
    
    def run(self):
        self.db = Database(self.uri)
        while True:
            cache = []
            try:
                cache.append(self.q.get(timeout=self.timeout))
                if self.batch:
                    if len(cache) % self.batch == 0:
                        self.db.create_all(cache)
                        cache = []
                else:
                    
                    cache = []
            except queue.Empty:
                self.db.create_all(cache)
                break

# ~~~ Models ~~~~~~~~~~~~~~~~~

class Document(SQLModel, table=True):
    id: str = Field(primary_key=True)
    name: str
    href: str
    date: dt.datetime
    text: Optional[str] = None
    date_collected: dt.datetime
    collected_by: str

class Paragraph(SQLModel, table=True):
    id: str = Field(primary_key=True)
    text: str
    document_id: str = 