import mimetypes
from collections import defaultdict
from urllib.parse import quote

import pandas as pd
import numpy as np
from fastapi import APIRouter, Depends, Request
from motor.motor_asyncio import AsyncIOMotorDatabase
from starlette.responses import StreamingResponse

import crud, schemas
from common import *

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

from models.behavior_analysis import BehaviorAnalysis
from models.user_analysis import UserAnalysis
from models.x_analysis import XAnalysis
from utils import DfToStream

router = APIRouter()


@router.post("/ltv_model_sql")
async def ltv_model_sql(
        request: Request,
        game: str,
        analysis: XAnalysis = Depends(XAnalysis),
        current_user: schemas.UserDB = Depends(deps.get_current_user)
) -> schemas.Msg:
    """ ltv模型sql """
    await analysis.init()
    data = analysis.ltv_model_sql()
    return schemas.Msg(code=0, msg='ok', data=[data])


@router.post("/ltv_model")
async def ltv_model_sql(
        request: Request,
        game: str,
        analysis: XAnalysis = Depends(XAnalysis),
        ckdb: CKDrive = Depends(get_ck_db),
        current_user: schemas.UserDB = Depends(deps.get_current_user)
) -> schemas.Msg:
    """ ltv模型sql """
    await analysis.init()

    res = analysis.ltv_model_sql()
    sql = res['sql']
    quota = res['quota']
    df = await ckdb.query_dataframe(sql)
    df.fillna(0, inplace=True)
    df.rename(columns={'date': '注册日期'}, inplace=True)
    if quota == '#account_id':
        df.rename(columns={'cnt1': '角色数'}, inplace=True)
    elif quota in ('#distinct_id',):
        df.rename(columns={'cnt1': '设备数'}, inplace=True)
    data = {
        'title': df.columns.tolist(),
        'rows': df.values.tolist()
    }

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


@router.post("/ltv_model_export")
async def ltv_model_export(request: Request,
                           game: str,
                           ckdb: CKDrive = Depends(get_ck_db),
                           analysis: XAnalysis = Depends(XAnalysis),
                           current_user: schemas.UserDB = Depends(deps.get_current_user)
                           ):
    """ ltv分析 数据导出"""
    await analysis.init()
    data = analysis.ltv_model_sql()
    file_name = quote(f'lvt.xlsx')
    mime = mimetypes.guess_type(file_name)[0]

    sql = data['sql']
    df = await ckdb.query_dataframe(sql)
    df_to_stream = DfToStream((df, 'ltv'))
    with df_to_stream as d:
        export = d.to_stream()
    return StreamingResponse(export, media_type=mime, headers={'Content-Disposition': f'filename="{file_name}"'})