from ai.models import *
from ai.utils.show_prompt import *
import datetime
from django_redis import get_redis_connection
import re

# ai排片主流程
def show_main_process():
    # 查看状态
    redis_conn = get_redis_connection()
    redis_key = f'ai_show{datetime.date.today().strftime("%Y%m%d")}'
    if redis_conn.exists(redis_key):
        return False
    # 获取影院列表
    redis_conn.set(redis_key, 1, ex=60*60*20)
    test_cinema_list = TestCinema.objects.filter(is_active=True).all()
    # 生成目标日期
    show_date = datetime.date.strftime(datetime.date.today() + datetime.timedelta(days=3), '%Y-%m-%d')
    for cinema in test_cinema_list:
        print(cinema.name)
        show_ai = ShowAI(cinema, show_date)
        prompt = show_ai.general_prompt()
        start = datetime.datetime.now()
        result, message, tokens = show_ai.get_show_result_ai()
        end = datetime.datetime.now()
        print('prompt:', prompt)
        print('result:', result)
        print('message:', message)
        print('tokens:', tokens)
        # 获取排片数据
        _show = next((s for s in result.split('------') if s.startswith('\n影厅别名,影厅id')), '')
        _show = _show.strip()
        print('_show:', _show)
        # 获取销售额
        _sales = next((s for s in result.split('\n') if re.search('预估销售数据', s)), '')
        _sales = str(re.findall(r'\d+\.?\d*', _sales)[0])
        print('_sales:', _sales)
        # 处理返回结果
        try:
            AiShow.objects.create(
                cinema=cinema.name,
                zz_code=cinema.zz_code,
                show_date=show_date,
                is_ai_show=True,
                show=_show,
                sales=_sales,
                prompt=prompt,
                result=result,
                message=message,
                take_times=int((end - start).seconds),
                take_tokens=tokens
            )
        except Exception as e:
            print(e)
    return True