from typing import Optional
import pandas as pd
import streamlit as st
import os
from sqlmodel import Field, Session, SQLModel, create_engine, select


class Hero(SQLModel, table=True):
    __table_args__ = {'extend_existing': True}  # required for streamlit refreshing
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    secret_name: str
    age: Optional[int] = None


sqlite_file_name = 'database.db'
sqlite_url = f"sqlite:///{sqlite_file_name}"

engine = create_engine(sqlite_url, echo=True)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def create_heroes():
    hero_1 = Hero(name="Rusty-Man", secret_name="<NAME>", age=36)
    hero_2 = Hero(name="Tarantula", secret_name="<NAME>", age=26)
    hero_3 = Hero(name="<NAME>", secret_name="<NAME>", age=33)
    with Session(engine) as session:
        session.add(hero_1)
        session.add(hero_2)
        session.add(hero_3)
        session.commit()


def get_db_size():
    with Session(engine) as session:
        heroes = session.exec(select(Hero)).all()
    return len(heroes)


def select_heros():
    with Session(engine) as session:
        statement = select(Hero).where(Hero.age <= 35)
        results = session.exec(statement)
        for hero in results:
            st.text(hero)
        #st.text(len(results))


def show_table():
    with Session(engine) as session:
        heroes = session.exec(select(Hero)).all()
        st.table(pd.DataFrame([s.dict() for s in heroes[-5:]]))


def delete_db():
    with Session(engine) as session:
        heroes = session.exec(select(Hero)).all()
        for hero in heroes:
            session.delete(hero)  
        session.commit()  
    st.text("Deleted all rows")  


def write_new_row():

    with st.form('new_row'):
        name_input = st.text_input('Name', value="<NAME>")
        secret_name = st.text_input('Secret alias', value="Dr. Jon")

        submitted = st.form_submit_button("Submit")
        if submitted:
            hero = Hero(name=name_input, secret_name=secret_name, age=23)
            st.write('submitted')

            hero_1 = Hero(name="Rusty-Man", secret_name="<NAME>", age=36)
            with Session(engine) as session:
                session.add(hero_1)
                session.commit()


# ====================================== main ====================================== #

def main():
    create_db_and_tables()
    st.title('🦄 SQLModel Demo')
    b1, b2, b3, b4= st.columns(4)
    #if b1.button('Add Filter'):
    #    pass
        #select_heros()  # todo
    if b4.button("♻️ Empty db"):
        delete_db()
    #if b2.button('Create db'): 
    #    create_db_and_tables()
    if b3.button('+ Add 3 rows'): 
        create_heroes()

    if st.button("➡️ Insert Row"):
        write_new_row()

    show_table()
    col0, col1, col2 = st.columns(3)
    file_size = os.path.getsize('database.db')
    col1.metric("💾 database.db", f"{get_db_size()}", "total rows")
    col2.metric("filesize", f"{file_size/1000:0.1f}", 'kb')


if __name__ == '__main__':
    st.set_page_config(
        page_title="SQLModel Demo",
        page_icon="🦄",
        layout="centered",
        initial_sidebar_state="auto")
    main()

streamlit.form streamlit.set_page_config streamlit.columns sqlmodel.create_engine sqlmodel.Session streamlit.text_input sqlmodel.SQLModel.metadata.create_all os.path.getsize sqlmodel.Field sqlmodel.select streamlit.title streamlit.form_submit_button streamlit.write streamlit.text streamlit.button
[((487, 523), 'sqlmodel.create_engine', 'create_engine', (['sqlite_url'], {'echo': '(True)'}), '(sqlite_url, echo=True)\n', (500, 523), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((294, 331), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (299, 331), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((558, 594), 'sqlmodel.SQLModel.metadata.create_all', 'SQLModel.metadata.create_all', (['engine'], {}), '(engine)\n', (586, 594), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1703, 1730), 'streamlit.text', 'st.text', (['"""Deleted all rows"""'], {}), "('Deleted all rows')\n", (1710, 1730), True, 'import streamlit as st\n'), ((2419, 2446), 'streamlit.title', 'st.title', (['"""🦄 SQLModel Demo"""'], {}), "('🦄 SQLModel Demo')\n", (2427, 2446), True, 'import streamlit as st\n'), ((2467, 2480), 'streamlit.columns', 'st.columns', (['(4)'], {}), '(4)\n', (2477, 2480), True, 'import streamlit as st\n'), ((2745, 2771), 'streamlit.button', 'st.button', (['"""➡️ Insert Row"""'], {}), "('➡️ Insert Row')\n", (2754, 2771), True, 'import streamlit as st\n'), ((2838, 2851), 'streamlit.columns', 'st.columns', (['(3)'], {}), '(3)\n', (2848, 2851), True, 'import streamlit as st\n'), ((2868, 2898), 'os.path.getsize', 'os.path.getsize', (['"""database.db"""'], {}), "('database.db')\n", (2883, 2898), False, 'import os\n'), ((3059, 3174), 'streamlit.set_page_config', 'st.set_page_config', ([], {'page_title': '"""SQLModel Demo"""', 'page_icon': '"""🦄"""', 'layout': '"""centered"""', 'initial_sidebar_state': '"""auto"""'}), "(page_title='SQLModel Demo', page_icon='🦄', layout=\n    'centered', initial_sidebar_state='auto')\n", (3077, 3174), True, 'import streamlit as st\n'), ((822, 837), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (829, 837), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((989, 1004), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (996, 1004), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1121, 1136), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1128, 1136), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1361, 1376), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1368, 1376), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1531, 1546), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1538, 1546), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1766, 1784), 'streamlit.form', 'st.form', (['"""new_row"""'], {}), "('new_row')\n", (1773, 1784), True, 'import streamlit as st\n'), ((1807, 1844), 'streamlit.text_input', 'st.text_input', (['"""Name"""'], {'value': '"""<NAME>"""'}), "('Name', value='<NAME>')\n", (1820, 1844), True, 'import streamlit as st\n'), ((1867, 1913), 'streamlit.text_input', 'st.text_input', (['"""Secret alias"""'], {'value': '"""Dr. Jon"""'}), "('Secret alias', value='Dr. Jon')\n", (1880, 1913), True, 'import streamlit as st\n'), ((1935, 1966), 'streamlit.form_submit_button', 'st.form_submit_button', (['"""Submit"""'], {}), "('Submit')\n", (1956, 1966), True, 'import streamlit as st\n'), ((1287, 1300), 'streamlit.text', 'st.text', (['hero'], {}), '(hero)\n', (1294, 1300), True, 'import streamlit as st\n'), ((2075, 2096), 'streamlit.write', 'st.write', (['"""submitted"""'], {}), "('submitted')\n", (2083, 2096), True, 'import streamlit as st\n'), ((1169, 1181), 'sqlmodel.select', 'select', (['Hero'], {}), '(Hero)\n', (1175, 1181), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2189, 2204), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (2196, 2204), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1047, 1059), 'sqlmodel.select', 'select', (['Hero'], {}), '(Hero)\n', (1053, 1059), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1419, 1431), 'sqlmodel.select', 'select', (['Hero'], {}), '(Hero)\n', (1425, 1431), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1589, 1601), 'sqlmodel.select', 'select', (['Hero'], {}), '(Hero)\n', (1595, 1601), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n')]