from flask import Blueprint, g
from consts.consts import IdentifyTypes
from handlers.base import check_org_id, fail, get_limit_offset, get_user_id, success

from handlers.edatask import _handle_chat2data, _handle_sql_refine_fix, _handle_create_session_context, _handle_update_session_context, _handle_reset_session_context, _handle_get_session_context, _handle_multi_round_chat2sql, _handle_multi_round_feedback
from models.context import Context
from models.edatask import EDATaskJob


sqleditor_v1_bp = Blueprint("sqleditor_v1_bp", __name__)


@sqleditor_v1_bp.post("/chat2sql")
def chat2data(org_id: int, project_id: int):
    return _handle_chat2data(False, True, project_id, data_to_chart=False)


@sqleditor_v1_bp.post("/sqlAutoFix")
def sql_auto_fix(org_id: int, project_id: int):
    return _handle_sql_refine_fix(sqleditor_mode=True)


@sqleditor_v1_bp.post("/sqlRefine")
def sql_refine(org_id: int, project_id: int):
    return _handle_sql_refine_fix(sqleditor_mode=True)


@sqleditor_v1_bp.get("/jobs/<job_id>")
def query_job_detail(org_id: int, project_id: int, job_id: str):
    edatask_job = EDATaskJob.get_by_job_id(g.session, job_id)
    if not edatask_job:
        return fail(404)

    check_org_id(edatask_job.org_id)

    job_detail = edatask_job.to_json()
    job_detail.pop("variables", None)
    job_detail.pop("output", None)

    return success(result=job_detail)


@sqleditor_v1_bp.get("/sessions")
def query_sessions(org_id: int, project_id: int):
    user_id = get_user_id()
    limit, offset = get_limit_offset()

    sessions = []
    items, count = Context.get_all_session_contexts_by_identifier_and_type(g.session, IdentifyTypes.USER_ID, user_id, limit, offset)
    for i in items:
        sessions.append(i.to_session_context(simple=True))

    return success(result={"sessions": sessions, "count": count})


@sqleditor_v1_bp.post("/sessions")
def create_session(org_id: int, project_id: int):
    return _handle_create_session_context(True)


@sqleditor_v1_bp.get("/sessions/<int:session_context_id>")
def query_session(org_id: int, project_id: int, session_context_id: int):
    return _handle_get_session_context(session_context_id, True)


@sqleditor_v1_bp.put("/sessions/<int:session_context_id>/reset")
def reset_session(org_id: int, project_id: int, session_context_id: int):
    return _handle_reset_session_context(session_context_id, True)


@sqleditor_v1_bp.put("/sessions/<int:session_context_id>")
def update_session(org_id: int, project_id: int, session_context_id: int):
    return _handle_update_session_context(session_context_id, True)


@sqleditor_v1_bp.post("/sessions/<int:session_id>/multiRoundChat2sql")
def multi_round_chat2sql(org_id: int, project_id: int, session_id: int):
    return _handle_multi_round_chat2sql(session_id, True, data_to_chart=False)


@sqleditor_v1_bp.post("/multiRoundFeedback")
def multi_round_feedback(org_id: int, project_id: int):
    return _handle_multi_round_feedback(True)