from flask import Blueprint, request, g
from jsonpath_ng.ext import parse
from consts.consts import ContextTypes

from handlers.base import (
    success,
    fail,
    check_schema,
)
from schema.schema import (
    CreateContextSchema,
    EvalTemplateSchema,
    UpdateContextSchema,
    UpsertContextSchema,
    ClearContextSchema,
)
from models.context import Context
from utils.jsonpath import (
    parse_jsonpath_context_type,
    is_valid_jsonpath,
)
from utils.template import fmt_template
from utils.logging import logging

context_bp = Blueprint("context_bp", __name__)


def eval_context_by_expr(session, session_context, expr: str):
    context_type, expr = parse_jsonpath_context_type(expr)
    logging.info("parse_jsonpath_context_type got context_type, expr: %s, %s", context_type, expr)
    if context_type == ContextTypes.UNKNOWN:
        return None

    jsonpath_expr = parse(expr)
    result = None
    if context_type == ContextTypes.SESSION:
        data = session_context.load_data()
        result = jsonpath_expr.find(data)
        logging.info("find data from session context %s got result %s", str(data), result)
    elif context_type == ContextTypes.DATA:
        data_context = Context.get_by_id(session, session_context.data_context_id)
        if data_context:
            data = data_context.load_data()
            result = jsonpath_expr.find(data)
            logging.info("find data from data context %s got result %s", str(data), result)
    else:  # domain context, eval expr in order: sesion -> data -> domain
        eval_datas = []
        # add override domain context in session context
        eval_datas.append(session_context.load_data().get("domain"))
        # add override context in data context
        if session_context.data_context_id:
            data_context = Context.get_by_id(session, session_context.data_context_id)
            if data_context:
                eval_datas.append(data_context.load_data().get("domain"))
        # add context in domain context
        if session_context.domain_context_id:
            domain_context = Context.get_by_id(session, session_context.domain_context_id)
            if domain_context:
                eval_datas.append(domain_context.load_data())

        # start to eval
        for i in eval_datas:
            if i:
                logging.info("find data from domain context %s got result %s", str(i), result)
                result = jsonpath_expr.find(i)
                if result:
                    break

    if result and len(result) > 0:
        result = result[0].value

    return result


@context_bp.get("/eval")
def eval_context():
    expr = request.args["expr"]
    session_id = request.args["session_id"]
    logging.info("eval_context with expr %s, session_id %s", expr, session_id)

    if not is_valid_jsonpath(expr):
        return fail(msg="invalid jsonpath")

    session_context = Context.get_session_context_by_session_id(g.session, session_id)
    if not session_context:
        return fail(404)

    result = eval_context_by_expr(g.session, session_context, expr)

    return success(result=result)


@context_bp.post("/eval")
def eval_prompt():
    payload = check_schema(EvalTemplateSchema, request.get_json())
    session_id = payload["session_id"]
    template = payload["template"]
    variables = payload["vars"]

    for _, expr in variables.items():
        if not is_valid_jsonpath(expr):
            return fail(msg="invalid jsonpath")

    session_context = Context.get_session_context_by_session_id(g.session, session_id)
    if not session_context:
        return fail(404)

    fmt_vars = {}
    for var, expr in variables.items():
        fmt_vars[var] = eval_context_by_expr(g.session, session_context, expr)

    result = fmt_template(template, fmt_vars)

    return success(result=result)


@context_bp.get("/orgs/<org_id>/contexts")
def get_contexts(org_id: int):
    context_id = request.args.get("context_id", type=int)
    cluster_id = request.args.get("cluster_id", type=int)
    db_name = request.args.get("db_name")
    session_id = request.args.get("session_id")

    context = None
    if context_id:
        context = Context.get_by_id(g.session, context_id)
    elif cluster_id and db_name:
        context = Context.get_data_context_by_cluster_and_db(g.session, cluster_id, db_name)
    elif session_id:
        context = Context.get_session_context_by_session_id(g.session, session_id)
        logging.info("get context by session_id %s: %s", session_id, context)

    if not context:
        return fail(404)

    return success(result=context.to_json())


@context_bp.put("/orgs/<org_id>/contexts")
def upsert_contexts(org_id: int):
    payload = check_schema(UpsertContextSchema, request.get_json())
    logging.info("upsert context with args %s", str(payload))

    _type = payload["type"]
    context_id = payload.get("context_id")
    cluster_id = payload.get("cluster_id")
    db_name = payload.get("db_name")
    session_id = payload.get("session_id")

    context = None
    if context_id:
        context = Context.get_by_id(g.session, context_id)
        logging.info("get context by id %s: %s", context_id, context)
    elif _type == ContextTypes.DATA and cluster_id and db_name:
        context = Context.get_data_context_by_cluster_and_db(g.session, cluster_id, db_name)
        logging.info("get context by cluster_id, db_name %s, %s: %s", cluster_id, db_name, context)
    elif _type == ContextTypes.SESSION and session_id:
        context = Context.get_session_context_by_session_id(g.session, session_id)
        logging.info("get context by session_id %s: %s", session_id, context)

    if context:
        if payload["data"]:
            context.update_data(payload["data"])
        context.data_context_id = payload["data_context_id"]
        context.domain_context_id = payload["domain_context_id"]
    else:
        context = Context.create(
            g.session, payload["type"], payload["name"], org_id, payload["cluster_id"], payload["db_name"],
            payload["session_id"], payload["creator"], payload["data"], payload["data_context_id"],
            payload["domain_context_id"],
        )

    return success(result=context.to_json())


@context_bp.get("/orgs/<org_id>/contexts/<context_type>")
def get_contexts_by_type(org_id: int, context_type: str):
    contexts = Context.get_all_by_org_id_and_type(g.session, org_id, context_type)

    result = {
        "items": [i.to_json() for i in contexts]
    }

    return success(result=result)


@context_bp.put("/orgs/<org_id>/contexts/<context_id>")
def update_context(org_id: int, context_id: int):
    payload = check_schema(UpdateContextSchema, request.get_json())

    context = Context.get_by_id(g.session, context_id)
    if not context:
        return fail(404)

    if payload["data"]:
        context.update_data(payload["data"])
    return success()


@context_bp.post("/orgs/<org_id>/contexts")
def create_context(org_id: int):
    payload = check_schema(CreateContextSchema, request.get_json())

    context = Context.create(
        g.session, payload["type"], payload["name"], org_id, payload["project_id"], payload["cluster_id"],
        payload["db_name"], payload["session_id"], payload["creator"], payload["data"],
        payload["data_context_id"], payload["domain_context_id"],
    )

    return success(result={"context_id": context.id})


@context_bp.delete("/orgs/<org_id>/contexts/<context_id>")
def delete_context(org_id: int, context_id: int):
    context = Context.get_by_id(g.session, context_id)
    if not context:
        return fail(code=404, msg="context not found")

    payload = check_schema(ClearContextSchema, request.get_json())
    data = context.load_data()

    if payload["all"]:
        data = {}
    else:
        data.pop(payload["key"], None)

    context.update_data(data)

    return success()