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

# ai排片主流程
def show_main_process(zz_code=None):
    # 查看状态
    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)
    if not zz_code:
        test_cinema_list = TestCinema.objects.filter(is_active=True).all()
    else:
        test_cinema_list = TestCinema.objects.filter(Q(zz_code=zz_code) & Q(is_active=True)).all()
    # 获取提示词版本
    prompt_ver = PromptTemplate.objects.filter(
        Q(del_flag=False) & Q(prompt_type='version') & Q(prompt_key='ShowPromptVersion')).first().prompt_val
    # 开始处理提示词
    # 生成目标日期
    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()
        # 方式二
        if _show == '':
            _show = next((s for s in result.split('------') if s.startswith('\n```csv\n影厅别名,')), '')
            _show = _show.replace('```', '').replace('csv', '').strip()
        # 方式三
        if _show == '':
            # 正则提取中间内容
            csv_pattern = r"```csv\s*([\s\S]*?)```"
            match = re.search(csv_pattern, result)
            if match:
                _show = match.group(1)
        print('_show:', _show)
        # 获取销售额
        _sales = next((s for s in result.split('\n') if re.search(r'\d{3,6}元', s)), '')
        _sales = str(re.findall(r'\d{3,6}元', _sales)[-1])
        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,
                prompt_version=prompt_ver,
            )
        except Exception as e:
            print(e)
    return True