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 = Field(foreign_key="document.id")
    sentiment: str
    sent_score: float

class Entity(SQLModel, table=True):
    id: str = Field(primary_key=True)
    name: str
    description: Optional[str]

class EntityMention(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    text: str
    score: Optional[float]
    label: str
    start: int
    end: int
    paragraph_id: str = Field(foreign_key="paragraph.id")
    kb_id: Optional[str] = Field(foreign_key="entity.id")

class EntityFeature(SQLModel, table=True):
    id: int = Field(primary_key=True)
    kb_id: str = Field(foreign_key="entity.id")
    key: str
    value: str





sqlmodel.create_engine sqlmodel.select sqlmodel.Session sqlmodel.SQLModel.metadata.create_all sqlmodel.Field
[((2043, 2066), 'sqlmodel.Field', 'Field', ([], {'primary_key': '(True)'}), '(primary_key=True)\n', (2048, 2066), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2256, 2279), 'sqlmodel.Field', 'Field', ([], {'primary_key': '(True)'}), '(primary_key=True)\n', (2261, 2279), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2317, 2349), 'sqlmodel.Field', 'Field', ([], {'foreign_key': '"""document.id"""'}), "(foreign_key='document.id')\n", (2322, 2349), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2442, 2465), 'sqlmodel.Field', 'Field', ([], {'primary_key': '(True)'}), '(primary_key=True)\n', (2447, 2465), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2579, 2616), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (2584, 2616), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2725, 2758), 'sqlmodel.Field', 'Field', ([], {'foreign_key': '"""paragraph.id"""'}), "(foreign_key='paragraph.id')\n", (2730, 2758), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2786, 2816), 'sqlmodel.Field', 'Field', ([], {'foreign_key': '"""entity.id"""'}), "(foreign_key='entity.id')\n", (2791, 2816), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2875, 2898), 'sqlmodel.Field', 'Field', ([], {'primary_key': '(True)'}), '(primary_key=True)\n', (2880, 2898), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2916, 2946), 'sqlmodel.Field', 'Field', ([], {'foreign_key': '"""entity.id"""'}), "(foreign_key='entity.id')\n", (2921, 2946), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((284, 302), 'sqlmodel.create_engine', 'create_engine', (['uri'], {}), '(uri)\n', (297, 302), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((311, 352), 'sqlmodel.SQLModel.metadata.create_all', 'SQLModel.metadata.create_all', (['self.engine'], {}), '(self.engine)\n', (339, 352), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((416, 436), 'sqlmodel.Session', 'Session', (['self.engine'], {}), '(self.engine)\n', (423, 436), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((620, 640), 'sqlmodel.Session', 'Session', (['self.engine'], {}), '(self.engine)\n', (627, 640), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1004, 1024), 'sqlmodel.Session', 'Session', (['self.engine'], {}), '(self.engine)\n', (1011, 1024), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((841, 854), 'sqlmodel.select', 'select', (['model'], {}), '(model)\n', (847, 854), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((672, 685), 'sqlmodel.select', 'select', (['model'], {}), '(model)\n', (678, 685), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n')]