from utils.rdbmsUtil.databaseUtil import DatabaseUtil
from utils.rdbmsUtil.modelBase import Base
from utils.rdbmsUtil.modelUser import User, user_factory  # user_factory 确保工厂注册代码被执行
from utils.rdbmsUtil.modelUserOperations import UserOperations
# company_factory 确保工厂注册代码被执行
from utils.rdbmsUtil.modelCompany import Company, company_factory
from utils.rdbmsUtil.modelCompanyOperations import CompanyOperations

from sqlalchemy.sql import text
import pandas as pd


def access_sqlite():
    db = DatabaseUtil('sqlite:///mydatabase.db')
    Base.metadata.create_all(db.engine)

    companys = [Company(id=1, name="DS-Tech"), Company(id=2, name="ZN-Tech")]
    users = [User(name="Alice", age=31, companyid=1),
            User(name="Bob", age=25, companyid=2)]


    with db.session_scope() as session:
        company_ops = CompanyOperations()

        user_ops = UserOperations()
        
        company_ops.upsert_companys(session, companys)
        user_ops.upsert_users(session, users)
        # session.flush()  # 确保所有的 SQL 操作都被推送到数据库

        user_ops.update(session, name="Alice", age=32)

        result = user_ops.get_user(session, name="Alice", age=32)
        for user in result:
            print(user)


def access_tidb():
    print('access MySQL/TiDB')
    # DatabaseUtil(url='mysql+pymysql://root:123456@192.168.50.234:4000/energy',
    #                     pool_size=5,
    #                     max_overflow=3,
    #                     connect_timeout=1800)
    # db = DatabaseUtil('mysql+pymysql://root:123456@192.168.50.234:4000/energy')

    db = DatabaseUtil('mysql+pymysql://root:admin123456@192.168.50.235:30306/energy_data_prod')

    query_text = f"""SELECT * FROM analysis_result  """
    query_text=f"""SELECT  wind_turbine_number, time_stamp, wind_velocity, active_power, yaw_error1 FROM `WOF055000062_second` WHERE wind_turbine_number IN ('WOG02325') AND time_stamp >= '2026-02-01 00:00:00' AND time_stamp <= '2026-04-26 00:00:00' AND lab in (0,1,2,3,4)"""

    chunks = []
    with db.session_scope() as session:
        # query_result = session.execute(text(query_text)).fetchall()

        # dataFrame = pd.DataFrame(query_result)

        conn=session.connection()
        
        result=conn.execution_options(stream_results=True).execute(text(query_text))

        # 每次取 100000 行，可根据实际情况调整
        while True:
            rows = result.fetchmany(100000)
            if not rows:
                break
            chunks.append(pd.DataFrame(rows))

    if not chunks:
        return pd.DataFrame()
    
    dataFrame=pd.concat(chunks, ignore_index=True)        
    
    print(len(dataFrame))
    print(dataFrame.head(5))

access_tidb()