import json

import pandas as pd
from fastapi import APIRouter, Depends, Request
import crud, schemas

from api import deps
from db.ckdb import get_ck_db, CKDrive
from db.redisdb import get_redis_pool, RedisDrive
from models import ToSql

router = APIRouter()


@router.post("/sql")
async def query_sql(
        request: Request,
        data_in: schemas.Sql,
        ckdb: CKDrive = Depends(get_ck_db),
        current_user: schemas.UserDB = Depends(deps.get_current_user)
) -> schemas.Msg:
    """原 sql 查询 """
    data = await ckdb.execute(data_in.sql)
    return schemas.Msg(code=0, msg='ok', data=data)


@router.post("/event_model_sql")
async def event_model_sql(
        request: Request,
        game: str,
        data_in: schemas.CkQuery,
        ckdb: CKDrive = Depends(get_ck_db),
        rdb: RedisDrive = Depends(get_redis_pool),
        current_user: schemas.UserDB = Depends(deps.get_current_user)
) -> schemas.Msg:
    """ 事件分析模型 sql"""

    columns_json = await rdb.get(f'{game}_event')
    columns = json.loads(columns_json)
    to_sql = ToSql(data_in.dict(), game, 'event', columns.keys())
    res = to_sql.get_sql_query_event_model()
    return schemas.Msg(code=0, msg='ok', data=res)


@router.post("/event_model")
async def event_model(
        request: Request,
        game: str,
        data_in: schemas.CkQuery,
        ckdb: CKDrive = Depends(get_ck_db),
        rdb: RedisDrive = Depends(get_redis_pool),
        current_user: schemas.UserDB = Depends(deps.get_current_user)
) -> schemas.Msg:
    """ 事件分析"""
    columns_json = await rdb.get(f'{game}_event')
    columns = json.loads(columns_json)
    to_sql = ToSql(data_in.dict(), game, 'event', columns.keys())
    sqls = to_sql.get_sql_query_event_model()
    res = []
    for item in sqls:
        q = {
            'groups': [],
            'values': [],
            'event_name': item['event_name']
        }
        sql = item['sql']
        groupby = item['groupby'][1:]
        date_range = item['date_range']
        q['date_range'] = date_range
        df = await ckdb.query_dataframe(sql)
        if df.shape[0]==0:
            return schemas.Msg(code=0, msg='ok', data=q)
        df['date'] = df['date'].apply(lambda x: str(x))
        # todo 时间粒度 暂时按天
        df['date'] = df['date'].apply(
            lambda x: pd.Timestamp(year=int(x[:4]), month=int(x[4:6]), day=int(x[6:])).strftime('%Y-%m-%d'))

        if groupby:
            # 有分组
            for group, df_group in df.groupby(groupby):
                df_group.reset_index(drop=True, inplace=True)
                q['groups'].append(group)

                concat_data = []
                for i in set(date_range) - set(df_group['date']):
                    if len(groupby) > 1:
                        concat_data.append((i, *group, 0))
                    else:
                        concat_data.append((i, group, 0))
                df_group = pd.concat([df_group, pd.DataFrame(concat_data, columns=df_group.columns)])
                df_group.sort_values('date', inplace=True)
                q['values'].append(df_group['values'].to_list())

        else:
            #     无分组
            concat_data = []
            for i in set(date_range) - set(df['date']):
                concat_data.append((i, 0))
            df = pd.concat([df, pd.DataFrame(concat_data, columns=df.columns)])
            q['values'].append(df['values'].to_list())

        res.append(q)
    return schemas.Msg(code=0, msg='ok', data=res)