# -------------------------------------------------------------------------
#  This file is part of the MindStudio project.
# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
#
# MindStudio is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# -------------------------------------------------------------------------

import os
from datetime import datetime
from celery.app.base import Celery

from atk.db.db_tools import DatabasePool
from atk.db.models import CaseResult
from atk.common.output_manager import OutputManager
from atk.common.utils import chmod_output_path

DB_DIR = OutputManager.get_dir_by_env("db")
DB_NAME = os.path.join(DB_DIR, f"celerydb.sqlite3")
RESULT_NAME = os.path.join(DB_DIR, f"result.sqlite3")
ATK_NAME = os.path.join(DB_DIR, f"atk.sqlite3")
ALL_DB = [DB_NAME, RESULT_NAME, ATK_NAME]

# 创建Celery应用并配置Broker
app = Celery("atk", broker=f"sqla+sqlite:///{DB_NAME}?timeout=30")

# 配置其他设置，如任务序列化方式等
app.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_backend=f"db+sqlite:///{RESULT_NAME}?timeout=30",  # 存储任务结果
    result_expires=3600,  # 结果过期时间
    task_reject_on_worker_lost=True,
    task_acks_late=True,
    worker_log_color=False,
    broker_connection_retry_on_startup=True,
    worker_hijack_root_logger=False,
    worker_redirect_stdouts=False,
    task_track_started=True,
    include=["atk.tasks.celery_tasks"],  # 确保任务模块已加载
)

db_pool = DatabasePool(db_url=f"sqlite:///{ATK_NAME}")
with db_pool as db:
    r = CaseResult()
    db.add(r)
    db.delete(r)
    db.execute("PRAGMA journal_mode=WAL;")
with DatabasePool(db_url=f"sqlite:///{DB_NAME}") as db:
    db.execute("PRAGMA journal_mode=WAL;")
with DatabasePool(db_url=f"sqlite:///{RESULT_NAME}") as db:
    db.execute("PRAGMA journal_mode=WAL;")
chmod_output_path(DB_DIR)


def clear_db_all_table(db_name):
    from sqlalchemy import create_engine, MetaData

    # 创建 SQLAlchemy Engine
    engine = create_engine(f"sqlite:///{db_name}")

    # 获取数据库中的所有表
    metadata = MetaData()
    metadata.reflect(bind=engine)

    # 删除所有表
    metadata.drop_all(bind=engine)

    # 重新创建所有表
    metadata.create_all(bind=engine)


def clear_celery_db():
    clear_db_all_table(DB_NAME)
    clear_db_all_table(RESULT_NAME)


def get_base_time():
    if not hasattr(app.conf, "GLOBAL_BASE_TIME"):
        data_utc = datetime.utcnow()
        timestamp = data_utc.timestamp()
        base_time = datetime.utcfromtimestamp(timestamp)
        app.conf.update(GLOBAL_BASE_TIME=base_time)
    return app.conf.GLOBAL_BASE_TIME