import datetime
import json
import time
from collections import defaultdict
from typing import Union

from celery import Task
from prometheus_client import Histogram

from config import config
from consts.consts import (EDA_PROMETHEUS_BUCKETS, MODE_EXPLORE, MODE_REFINE, ContextTypes,
                           EdaTaskJobStatus)
from errors.errors import JobException, parse_exception
from models.ailog import HistoryLog
from models.base import rds, with_session
from models.context import Context
from models.edatask import EDATaskJob
from models.oplog import OpLog
from models.user_preference import UserPreference, is_internal_api_enabled
from svc.aipower import aipower_chat2query_v1
from svc.data_explore import (entities_extract, explore_column_description,
                              explore_db_description,
                              explore_table_description,
                              explore_tables_relationship, setup_data_context,
                              update_column_description,
                              update_table_description)
from svc.dataapi import get_sql_executor
from svc.eda_task import EDATask
from svc.measuring import submit_measure_metrics
from tasks.celery_app import app
from utils.context import fmt_task
from utils.database import get_database_from_uri
from utils.local import (clear_data_context_id, get_local_metrics, get_tl_app_id_pubkey,
                         get_tl_user_token, set_data_context_id,
                         set_local_metrics, set_session, set_tl_app_id_pubkey, set_tl_override_ai_model,
                         set_tl_user_token)
from utils.logging import logging, trace_id_vars

BIZ_LATENCY = Histogram(
    'biz_latency_seconds', 'business latency in seconds(support float number)',
    labelnames=("name", "result", "is_internal", "status_code", "llm_model"), buckets=EDA_PROMETHEUS_BUCKETS,
)
TASK_LATENCY = Histogram(
    'task_latency_seconds', 'task latency in seconds(support float number)',
    labelnames=("name", "result", "einfo", "is_internal"), buckets=EDA_PROMETHEUS_BUCKETS,
)
QUEUE_LATENCY = Histogram(
    'queue_latency_seconds', 'queue latency in seconds(support float number)',
    labelnames=("name", "is_internal"), buckets=EDA_PROMETHEUS_BUCKETS,
)
TASK_TIME_LIMIT = 600  # 10 minutes


def inject_chat2data_result_into_metrics(task_tree: dict):
    if not task_tree:
        logging.info("no result to inject into metrics, skip")
        return

    metrics = get_local_metrics()
    metrics["output"] = fmt_task(task_tree, "0", True)

    set_local_metrics(metrics)


def is_internal_metrics(metrics):
    # we don't observe metrics from internal users
    if metrics and metrics.get("user_email"):
        user_email = str(metrics["user_email"])
        org_id = metrics.get("org_id")
        if is_internal_api_enabled(user_email, org_id):
            return True

    return False


def trim_task_error_msg(einfo):
    msg = str(einfo).splitlines()[-1]
    if ":" in msg:
        msg = msg.split(":")[-1]

    return msg.strip()


def persistant_ai_log(success: bool, error_msg: str = ""):
    metrics = get_local_metrics()
    status = "success" if success else "failed"
    logging.info("persistant_ai_log, success, metrics: %s, %s", success, metrics)
    #  {'sql_generate_mode': 'direct', 'sqleditor_mode': False, 'org_id': '1', 'app_id': 1, 'endpoint': 'POST - /v3/chat2data', 'app_key': '', 'user_id': 0, 'job_id': '714fdb42fce3486690109736e9eed89a', 'debug_action': '', 'user_token': '', 'user_email': '', 'trace_id': '', 'session_context_id': '46163', 'data_context_id': '1', 'start_at': 1720667750, 'preference': {'llm_api_key': None, 'language_code': 'English', 'llm_model': 'gpt-4o-mini'}, 'total_tokens': 1363, 'output': {'type': 'data_retrieval', 'task_id': '0', 'data': {'columns': [], 'rows_count': 0}, 'chart_options': {}, 'sql': 'SELECT `shop`.`shop_name`, SUM(`stock`.`quantity`) AS `total_quantity` FROM `shop` JOIN `stock` ON `shop`.`shop_id` = `stock`.`shop_id` GROUP BY `shop`.`shop_name` ORDER BY `total_quantity` DESC LIMIT 3;', 'sql_error': 'RetryError[<Future at 0x7fa6ddd20a90 state=finished raised ConnectionError>]', 'status': 'done'}}

    with with_session() as s:
        data_summary_id = metrics.get("data_context_id", 0)
        if not data_summary_id:
            logging.info("no data_summary_id in metrics, skip")
            return

        data_summary = Context.get_by_id(s, data_summary_id)

        cluster_id = data_summary.cluster_id
        database = data_summary.db_name or get_database_from_uri(data_summary.database_uri)
        org_id = metrics.get("org_id", "")
        _type = metrics.get("type", "")
        if not _type:
            logging.info("no type in metrics, skip persistant_ai_log")
            return

        trace_id = metrics.get("trace_id", "")
        start_at = datetime.datetime.fromtimestamp(metrics["start_at"])
        session_id = int(metrics.get("session_context_id", 0))
        recalled_knowledge_data = metrics.get("recalled_knowledge_data", {})
        token_count = metrics.get("total_tokens", 0)
        input = {
            "question": metrics.get("question", ""),
            "sql_generate_mode": metrics.get("sql_generate_mode", ""),
            "feedback": metrics.get("feedback", ""),
        }
        settings = metrics.get("preference", {})
        output = metrics.get("output", {})
        if not success:
            output["error"] = error_msg

        history_log = HistoryLog(
            status=status,
            trace_id=trace_id,
            _type=_type,
            org_id=org_id,
            cluster_id=cluster_id,
            database=database,
            question=input["question"],
            start_at=start_at,
            session_id=session_id,
            data_summary_id=data_summary_id,
            input=input,
            output=output,
            recalled_knowledge_data=recalled_knowledge_data,
            settings=settings,
            token_count=token_count,
        )
        s.add(history_log)


class EDACerlryTaskBase(Task):
    default_retry_delay = 1

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        with with_session() as s:
            edatask_job = EDATaskJob.get_by_job_id(s, task_id)
            assert edatask_job is not None

            logging.info("eda task %s marked as failed", task_id)
            edatask_job.update_status(s, EdaTaskJobStatus.FAILED)
            if einfo:
                edatask_job.set_result(s, {"reason": trim_task_error_msg(einfo)})

            OpLog.update_job_status(s, task_id, EdaTaskJobStatus.FAILED)

            persistant_ai_log(False, trim_task_error_msg(einfo))

            metrics = kwargs.get("metrics")
            duration = (datetime.datetime.now() - edatask_job.created_at).total_seconds()
            TASK_LATENCY.labels(edatask_job.fn_name, "failed", str(einfo), is_internal_metrics(metrics)).observe(duration)

    def on_success(self, retval, task_id, args, kwargs):
        with with_session() as s:
            edatask_job = EDATaskJob.get_by_job_id(s, task_id)
            assert edatask_job is not None

            logging.info("eda task %s marked as done", task_id)
            edatask_job.set_result(s, retval)
            edatask_job.update_status(s, EdaTaskJobStatus.DONE)

            OpLog.update_job_status(s, task_id, EdaTaskJobStatus.DONE)

            persistant_ai_log(True)

            metrics = kwargs.get("metrics")
            duration = (datetime.datetime.now() - edatask_job.created_at).total_seconds()
            TASK_LATENCY.labels(edatask_job.fn_name, "success", "", is_internal_metrics(metrics)).observe(duration)

    def before_start(self, task_id, args, kwargs):
        metrics = kwargs.get("metrics")
        if metrics:
            user_token = metrics.get("user_token", "")
            set_tl_user_token(user_token)

            app_id = int(metrics.get("app_id", 0))
            pubkey = metrics.get("app_key", "")
            if app_id and pubkey:
                set_tl_app_id_pubkey(app_id, pubkey)

            trace_id = metrics.get("trace_id", "")
            if trace_id:
                trace_id_vars.set(trace_id)

        with with_session() as s:
            edatask_job = EDATaskJob.get_by_job_id(s, task_id)
            assert edatask_job is not None

            logging.info("eda task %s marked as running", task_id)
            edatask_job.update_status(s, EdaTaskJobStatus.RUNNING)

            duration = (datetime.datetime.now() - edatask_job.created_at).total_seconds()
            QUEUE_LATENCY.labels(edatask_job.fn_name, is_internal_metrics(metrics)).observe(duration)

    def after_return(self, status, retval, task_id, args, kwargs, einfo):
        clear_data_context_id()


def task_gen_column_description(session, data_context_id: int, table_name: str, mode: str, language: str, ai_model: str) -> Union[dict, None]:
    set_data_context_id(data_context_id)
    set_session(session)

    data_context = Context.get_by_id(session, data_context_id)
    if not data_context:
        logging.info("data_context_id = %s, explore column description, failed, table_name = %s, reason = load data context empty, mode = %s", str(data_context_id), table_name, mode)
        return None

    ctx = data_context.load_data()

    logging.info("data_context_id = %s, explore column description, begin, table_name = %s, mode = %s", str(data_context_id), table_name, mode)

    result = explore_column_description(ctx, table_name, mode, language, ai_model)

    logging.info("data_context_id = %s, explore column description, after, table_name = %s, result = %s, mode = %s", str(data_context_id), table_name, str(result), mode)
    return result


def task_gen_table_description(session, data_context_id: int, table_name: str, mode: str, language: str, ai_model: str) -> Union[dict, None]:
    set_data_context_id(data_context_id)
    set_session(session)

    data_context = Context.get_by_id(session, data_context_id)
    if not data_context:
        logging.info("data_context_id = %s, explore table description, failed, reason = load data context empty, mode = %s", str(data_context_id), mode)
        return None

    logging.info("data_context_id = %s, explore table description, begin, table_name = %s, mode = %s", str(data_context_id), table_name, mode)
    result = explore_table_description(table_name, mode, language, ai_model)

    logging.info("data_context_id = %s, explore table description, after, table_name = %s, result = %s, mode = %s", str(data_context_id), table_name, str(result), mode)
    return result


def task_gen_table_relations(session, data_context_id: int, mode: str, language: str, ai_model: str) -> Union[dict, None]:

    set_data_context_id(data_context_id)
    set_session(session)

    data_context = Context.get_by_id(session, data_context_id)
    if not data_context:
        logging.info("data_context_id = %s, explore table relationships, failed, reason = load data context empty, mode = %s", str(data_context_id), mode)
        return None

    logging.info("data_context_id = %s, explore table relationships, begin, mode = %s", str(data_context_id), mode)
    result = None
    for i in (20, 15, 10, 5):
        try:
            result = explore_tables_relationship(mode, i, ai_model)
            break
        except Exception as e:
            logging.info("data_context_id = %s, explore table relationships, failed, reason = %s, relevant_table_param = %s, mode = %s", str(data_context_id), str(e), str(i), mode)

    logging.info("data_context_id = %s, explore table relationships, after, result = %s, mode = %s", str(data_context_id), str(result), mode)
    return result


def task_gen_entity(session, data_context_id: int, language: str, ai_model: str) -> Union[dict, None]:

    set_data_context_id(data_context_id)
    set_session(session)

    data_context = Context.get_by_id(session, data_context_id)
    if not data_context:
        logging.info("data_context_id = %s, extract entities, failed, reason = load data context empty", str(data_context_id))
        return None

    logging.info("data_context_id = %s, extract entities, begin", str(data_context_id))
    result = entities_extract(language, ai_model)

    logging.info("data_context_id = %s, extract entities, after, result = %s", str(data_context_id), str(result))
    return result


def task_gen_db_description(session, data_context_id: int, language: str, ai_model: str) -> Union[dict, None]:

    set_data_context_id(data_context_id)
    set_session(session)

    data_context = Context.get_by_id(session, data_context_id)
    if not data_context:
        logging.info("data_context_id = %s, explore db summary, failed, reason = load data context empty", str(data_context_id))
        return None

    logging.info("data_context_id = %s, explore db summary, begin", str(data_context_id))

    result = explore_db_description(language, ai_model)

    logging.info("data_context_id = %s, explore db summary, after, result = %s", str(data_context_id), str(result))
    return result


def set_measuring_metrics(metrics: Union[dict, None], job_id: str, data_context_id: int, session_context_id: int):
    metrics = metrics or {}

    sqleditor_mode = metrics.get("sqleditor_mode", False)
    preference = UserPreference.load(metrics.get("app_id", 0)) if not sqleditor_mode else None

    # set measure metrics
    metrics.update({
        "job_id": str(job_id),
        "session_context_id": str(session_context_id),
        "data_context_id": str(data_context_id),
        "start_at": int(time.time()),
        "preference": preference.to_history_log() if preference else {},
        "sqleditor_mode": sqleditor_mode,
    })

    set_local_metrics(metrics)


def upsert_local_metrics(metrics: dict):
    m = get_local_metrics()

    m.update(metrics)

    set_local_metrics(m)


@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_gen_db_schema(
    self, data_context_id, mode=MODE_EXPLORE, need_refine=False, set_default=False, metrics=None,
):
    job_id = self.request.id
    start_at = time.time()
    set_measuring_metrics(metrics, job_id, data_context_id, 0)

    if metrics:
        set_tl_override_ai_model(metrics.get("override_ai_model", ""))

    app_id, public_key = get_tl_app_id_pubkey()
    preference = UserPreference.load(app_id)

    try:
        with with_session() as session:
            data_context = Context.get_by_id(session, data_context_id)
            assert data_context is not None
            data = data_context.load_data()

            logging.info("setup_data_context with data_context_id = %s, data %s", data_context_id, str(data))
            user_token = get_tl_user_token()
            data = setup_data_context(data, app_id, public_key, user_token=user_token)
            data_context.update_data(data)

        with with_session() as session:
            edatask_job = EDATaskJob.get_by_job_id(session, job_id)
            assert edatask_job is not None

            data_context = Context.get_by_id(session, data_context_id)
            assert data_context is not None

            data = data_context.load_data(escape=True)
            language = preference.language_code
            ai_model = preference.llm_model

            if mode == MODE_EXPLORE and data.get("status") == "done":
                logging.info("data_context_id = %s, data context is already generated, ignore", str(data_context_id))

                if set_default:
                    data_context.update_description_and_default(session, data_context.description, set_default)

                return data_context.fmt_data_summary_v2()

            table_names = data["db_schema"].keys()
            logging.info("table_names %s", table_names)

            # 1. generate column schema
            for table_name in table_names:
                task_gen_column_description(session, data_context.id, table_name, mode, language, ai_model)
                data_context = Context.get_by_id(session, data_context_id)
                edatask_job.set_result(session, data_context.fmt_data_summary_v2())
                session.commit()

            # 2. generate table schema
            for table_name in table_names:
                task_gen_table_description(session, data_context.id, table_name, MODE_EXPLORE, language, ai_model)
                data_context = Context.get_by_id(session, data_context_id)
                edatask_job.set_result(session, data_context.fmt_data_summary_v2())
                session.commit()

            # 3. generate table relationships
            task_gen_table_relations(session, data_context.id, MODE_EXPLORE, language, ai_model)
            data_context = Context.get_by_id(session, data_context_id)
            edatask_job.set_result(session, data_context.fmt_data_summary_v2())
            session.commit()

            # 4. generate entity
            task_gen_entity(session, data_context.id, language, ai_model)
            data_context = Context.get_by_id(session, data_context_id)
            edatask_job.set_result(session, data_context.fmt_data_summary_v2())
            session.commit()

            # 5. generate db summary
            task_gen_db_description(session, data_context.id, language, ai_model)
            data_context = Context.get_by_id(session, data_context_id)
            edatask_job.set_result(session, data_context.fmt_data_summary_v2())
            session.commit()

            if need_refine:
                task_gen_db_schema(session, data_context_id, job_id, mode=MODE_REFINE, need_refine=False)

            # 6. return the whole data context
            data_context = Context.get_by_id(session, data_context.id)

            BIZ_LATENCY.labels(
                "task_understand_db", True, is_internal_metrics(metrics), 200, preference.llm_model,
            ).observe(time.time() - start_at)
            logging.info("task_gen_db_schema submit metrics: %s", True)

            if set_default:
                data_context.update_description_and_default(session, data_context.description, set_default)

            return data_context.fmt_data_summary_v2()
    except Exception as e:
        logging.error("failed to generate db summary: %s", e, exc_info=True)
        _, status_code = parse_exception(e)
        BIZ_LATENCY.labels(
            "task_understand_db", False, is_internal_metrics(metrics), status_code, preference.llm_model,
        ).observe(time.time() - start_at)
        logging.info("task_gen_db_schema submit metrics: %s", False)
        raise JobException(str(e))


@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_suggest_questions(
    self, session_context_id: int, data_context_id: int, app_id, public_key, user_token,
    question_num=5, question_only=True, metrics=None,
):
    logging.info("task_suggest_questions, data_context_id = %s, app_id = %s, public_key = %s, user_token = %s, question_num = %s, question_only = %s", data_context_id, app_id, public_key, user_token, question_num, question_only)

    job_id = self.request.id
    start_at = time.time()

    if metrics:
        set_tl_override_ai_model(metrics.get("override_ai_model", ""))

    preference = UserPreference.load(app_id)

    try:
        with with_session() as session:
            # understand db schema
            data_summary = Context.get_by_id(session, data_context_id)
            data = data_summary.load_data()
            data = setup_data_context(data, app_id, public_key, user_token=user_token)
            data_summary.update_data(data)

            edatask_job = EDATaskJob.get_by_job_id(session, job_id)
            assert edatask_job is not None

            logging.info("run task_suggest_questions with %s, %s", session_context_id, job_id)
            set_data_context_id(data_context_id)
            set_session(session)

            set_measuring_metrics(metrics, job_id, data_context_id, 0)
            eda = EDATask(session_context_id)

            resp = eda.task_suggest_questions(question_num=question_num)
            resp = resp or []
            result = [i["question"] for i in resp] if question_only else resp

            edatask_job.set_result(session, result)

            question_suggestion_generated = isinstance(result, list) and len(result) > 0

            BIZ_LATENCY.labels(
                "task_suggest_questions", question_suggestion_generated, is_internal_metrics(metrics), 200, preference.llm_model,
            ).observe(time.time() - start_at)
            logging.info("task_suggest_questions submit metrics question_suggestion_generated: %s", question_suggestion_generated)

            session.commit()

            return result
    except Exception as e:
        logging.error("failed to suggest questions: %s", e, exc_info=True)
        _, status_code = parse_exception(e)
        BIZ_LATENCY.labels(
            "task_suggest_questions", False, is_internal_metrics(metrics), status_code, preference.llm_model,
        ).observe(time.time() - start_at)
        raise e


@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_chat2data(
    self, data_context_id: int, session_context_id: int, raw_question: str, direct_resolve, auto_run, sql_to_data,
    data_to_chart, evidence, app_id, public_key, user_token=None, query_id="", task_id="", v3_format=False,
    metrics=None,
) -> Union[dict, None]:  # noqa
    job_id = self.request.id
    start_at = time.time()
    preference = UserPreference.load(app_id)

    if metrics:
        set_tl_override_ai_model(metrics.get("override_ai_model", ""))

    try:
        with with_session() as session:
            # understand db schema
            data_summary = Context.get_by_id(session, data_context_id)
            data = data_summary.load_data()
            data = setup_data_context(data, app_id, public_key, user_token=user_token)
            data_summary.update_data(data)

            edatask_job = EDATaskJob.get_by_job_id(session, job_id)
            assert edatask_job is not None

            set_data_context_id(data_context_id)
            set_session(session)

            data_context = Context.get_by_id(session, data_context_id)
            if not data_context:
                logging.info("session_id = %s, run task analyse failed, data_context_id does not exist, query = %s, data_context_id = %s", str(session_context_id), raw_question, str(data_context_id))
                return None

            logging.info("session_id = %s, run task analyse begin, query = %s, data_context_id = %s", str(session_context_id), raw_question, str(data_context_id))

            session_context = Context.get_by_id(session, session_context_id)
            assert session_context is not None
            thread_id = session_context.thread_id if session_context._type == ContextTypes.SESSION_CHAT else ""
            
            set_measuring_metrics(metrics, job_id, data_context_id, session_context_id)
            eda = EDATask(session_context_id, thread_id)

            result = None
            task_tree = {}
            for task_tree, question_id in eda.task_proxy(app_id, raw_question, 
                                                           explaination=evidence, 
                                                           query_id=query_id, 
                                                           task_id=task_id,
                                                           direct_resolve=direct_resolve, 
                                                           auto_run=auto_run, 
                                                           sql_to_data=sql_to_data, 
                                                           data_to_chart=data_to_chart):
                # eda.task_proxy + yield
                # question_id
                logging.info("session_id = %s, run task analyse successfully, yield result = %s", str(session_context_id), str(task_tree))

                if v3_format:
                    result = fmt_task(task_tree, "0")
                else:
                    result = {
                        "question_id": question_id,
                        "raw_question": raw_question if not query_id else "",
                        "feedback": raw_question if query_id else "",
                        "task_tree": task_tree,
                        "time_elapsed": time.time() - start_at,
                        "session_context_id": session_context_id,
                    }

                edatask_job.set_result(session, result)
                session.commit()

            sql_generated = eda.is_sql_generated()
            duration = time.time() - start_at

            label = "task_chat2data" if sql_to_data else "task_chat2sql"
            BIZ_LATENCY.labels(
                label, sql_generated, is_internal_metrics(metrics), 200, preference.llm_model,
            ).observe(duration)
            logging.info("task_chat2data submit metrics: %s", sql_generated)

            inject_chat2data_result_into_metrics(task_tree)

            return result
    except Exception as e:
        logging.error("session_id = %s, run task analyse failed, error = %s", str(session_context_id), str(e), exc_info=True)
        _, status_code = parse_exception(e)
        BIZ_LATENCY.labels(
            "task_chat2data", False, is_internal_metrics(metrics), status_code, preference.llm_model,
        ).observe(time.time() - start_at)
        raise e


@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_text2sql(
    self, data_context_id: int, session_context_id: int, task_id: str, question_id: str, metrics=None,
) -> Union[dict, None]:
    job_id = self.request.id
    start_at = time.time()
    app_id, public_key = get_tl_app_id_pubkey()
    preference = UserPreference.load(app_id)

    if metrics:
        set_tl_override_ai_model(metrics.get("override_ai_model", ""))

    try:
        with with_session() as session:
            set_data_context_id(data_context_id)
            set_session(session)

            data_context = Context.get_by_id(session, data_context_id)
            if not data_context:
                return None

            db_context = data_context.load_data()
            preference = UserPreference.load(db_context.get("app_id"))
            language = preference.language_code
            ai_model = preference.llm_model

            logging.info("session_id = %s, query_id = %s, task_id = %s, run text2sql - begin", str(session_context_id), str(question_id), str(task_id))

            set_measuring_metrics(metrics, job_id, data_context_id, session_context_id)
            eda = EDATask(session_context_id)
            task = eda.task_router(question_id, task_id, language=language, ai_model=ai_model)

            logging.info("session_id = %s, query_id = %s, task_id = %s, run text2sql - end, resp = %s", str(session_context_id), str(question_id), str(task_id), str(task))
            edatask_job = EDATaskJob.get_by_job_id(session, job_id)
            assert edatask_job is not None
            edatask_job.refresh_last_updated(session)

            sql_generated = eda.is_sql_generated()
            duration = time.time() - start_at
            BIZ_LATENCY.labels(
                "task_text2sql", sql_generated, is_internal_metrics(metrics), 200, preference.llm_model,
            ).observe(duration)
            logging.info("task_text2sql submit metrics: %s", sql_generated)

            return task
    except Exception as e:
        logging.error("failed to text2sql: %s", e, exc_info=True)
        _, status_code = parse_exception(e)
        BIZ_LATENCY.labels(
            "task_text2sql", False, is_internal_metrics(metrics), status_code, preference.llm_model,
        ).observe(time.time() - start_at)
        raise e


@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_update_data_summary(self, data_summary_id: int, tables: dict, metrics=None):
    set_data_context_id(data_summary_id)

    with with_session() as s:
        set_session(s)

        for table_name, v in tables.items():
            description = v["description"]

            # update table description
            update_table_description(table_name, description)

            for column_name, column_description in v["columns"].items():
                # update column description
                update_column_description(table_name, column_name, column_description)


def mark_expired_jobs():
    logging.info("start to mark expired jobs")
    threshold = datetime.datetime.now() - datetime.timedelta(minutes=config.MAX_EXPIRED_THRESHOLD_IN_MINUTES)

    with with_session() as s:
        for i in EDATaskJob.get_all_running_jobs_before(s, threshold, True):
            logging.info("mark job %s as failed(because of expired)", i.job_id)
            i.append_output(s, "mark job to expired as it's timeout")
            i.update_status(s, EdaTaskJobStatus.FAILED)
            s.flush()


def send_measure_metrics():
    logging.info("start to send measure metrics")
    d = defaultdict(list)

    for i in rds.scan_iter(match="measure:*", _type="LIST"):
        key = i.decode()
        logging.info("iter metrics of %s", key)
        items = rds.rpop(key, rds.llen(key))
        if not items:
            logging.info("no metrics of %s", key)
            continue

        for item in items:
            value = json.loads(item)
            org_id = value["tenant_id"]
            unit = value["unit"]
            key = f"{org_id}:{unit}"
            d[key].append(value)

    for key, metrics in d.items():
        logging.info("iter metrics of %s, length: %s", key, len(metrics))
        if len(metrics) == 0:
            continue

        cost = 0
        for metric in metrics:
            cost += metric["cost"]

        metric = metrics[-1]  # use the latest as base, replace the cost
        metric["cost"] = cost

        try:
            submit_measure_metrics(metric)
        except Exception as e:
            logging.error("failed to send measure metrics %s: %s", metric, e)


@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_sql_refine(
    self, data_context_id, session_context_id, sql, raw_question, feedback, sql_error, user_token=None, metrics=None,
):
    job_id = self.request.id
    set_measuring_metrics(metrics, job_id, data_context_id, session_context_id)

    if metrics:
        set_tl_override_ai_model(metrics.get("override_ai_model", ""))

    with with_session() as session:
        edatask_job = EDATaskJob.get_by_job_id(session, job_id)
        assert edatask_job is not None

        set_data_context_id(data_context_id)
        set_session(session)

        data_context = Context.get_by_id(session, data_context_id)
        assert data_context is not None
        data = data_context.load_data()

        logging.info("session_id = %s, setup_data_context with data_context_id = %s, data = %s", session_context_id, data_context_id, str(data))
        app_id, public_key = get_tl_app_id_pubkey()
        user_token = get_tl_user_token()
        data = setup_data_context(data, app_id, public_key, user_token=user_token)
        data_context.update_data(data)

        ctx = data_context.load_data()

        app_id, public_key = get_tl_app_id_pubkey()
        user_token = user_token
        executor = get_sql_executor(ctx, app_id, public_key, user_token=user_token)
        eda = EDATask(session_context_id)
        try:
            executor.explain_sql(sql)
            # rewrite sql(no sql error)
            result = eda.task_sql_rewrite(raw_question, sql, feedback)
        except Exception as e:
            logging.info("session_id = %s, failed to explain sql, explain_error = %s", session_context_id, e)
            # refine sql (has error), use input `sql_error` firstly
            result = eda.task_sql_refine(raw_question, sql, sql_error or str(e), feedback)
    
        return result
        

@app.task(bind=True, base=EDACerlryTaskBase, autoretry_for=(Exception,), retry_kwargs={'max_retries': 3}, time_limit=TASK_TIME_LIMIT)
def task_aipower_chat2query_v1(self, org_id, project_id, cluster_id, database, question, caller, metrics=None):
    job_id = self.request.id

    with with_session() as session:
        edatask_job = EDATaskJob.get_by_job_id(session, job_id)
        assert edatask_job is not None

        set_session(session)

        aipower_resp = aipower_chat2query_v1(org_id, project_id, cluster_id, database, question, caller)

        result = {
            "aipower_resp": aipower_resp,
            "task_tree": {
                "0": {
                    "sql": aipower_resp["response"],
                }
            }
        }

        return result