from marshmallow import Schema, fields, validate

from consts.consts import GPT_4, GPT_4O,GPT_DEFAULT, LANGUAGE_CN, LANGUAGE_EN, MODE_EXPLORE, MODE_TYPE_ASK_BILLING, MODE_TYPE_LENGJING, MODE_TYPE_SPIDER, MODE_TYPE_TIINSIGHT, Claude3, ContextTypes


class CheckConnectionSchema(Schema):
    database_uri = fields.Str(required=True)


class CreateContextSchema(Schema):
    type = fields.Str(required=True, validate=validate.OneOf([
        ContextTypes.DATA, ContextTypes.DOMAIN, ContextTypes.SESSION,
    ]))
    name = fields.Str(required=True)
    project_id = fields.Int(required=True)
    cluster_id = fields.Int(load_default=0)
    db_name = fields.String(load_default="")
    session_id = fields.String(load_default="")
    creator = fields.String(load_default="")
    data = fields.Dict(required=True)
    data_context_id = fields.Int(load_default=None)
    domain_context_id = fields.Int(load_default=None)


class UpsertContextSchema(Schema):
    type = fields.Str(required=True, validate=validate.OneOf([
        ContextTypes.DATA, ContextTypes.DOMAIN, ContextTypes.SESSION,
    ]))
    name = fields.Str(required=True)
    context_id = fields.Int(load_default=0)
    cluster_id = fields.Int(load_default=0)
    db_name = fields.String(load_default="")
    session_id = fields.String(load_default="")
    creator = fields.String(load_default="")
    data = fields.Dict(required=True)
    data_context_id = fields.Int(load_default=None)
    domain_context_id = fields.Int(load_default=None)


class UpdateContextSchema(Schema):
    data = fields.Dict(required=True)


class EvalTemplateSchema(Schema):
    session_id = fields.String(required=True)
    template = fields.String(required=True)
    vars = fields.Dict(keys=fields.Str(), values=fields.Str())


class ClearContextSchema(Schema):
    key = fields.String(required=True)
    all = fields.Boolean(load_default=False)


class FormatPromptSchema(Schema):
    name = fields.String(required=True)
    vars = fields.Dict(keys=fields.Str(), values=fields.Str())


class CreateEDATaskSchema(Schema):
    name = fields.String(required=True)
    url_path = fields.String(required=True)
    method = fields.String(required=True, validate=lambda i: i in ("GET", "POST", "PUT", "DELETE"))
    description = fields.String(load_default="")
    timeout = fields.Int(required=True)
    prompt_template = fields.String(required=True)
    exprs = fields.Dict(keys=fields.Str(), values=fields.Str(), required=True)


class UpdateEDATaskSchema(CreateEDATaskSchema):
    pass


class RunEDATaskSchema(Schema):
    variables = fields.Dict(keys=fields.Str(), values=fields.Str(), required=True)


class GenerateColumnSchema(Schema):
    table_name = fields.String(required=True)
    mode = fields.String(load_default=MODE_EXPLORE)


class RefineColumnSchema(Schema):
    data_context_id = fields.Integer(required=True)
    summary = fields.String(required=True)


class GenerateTableSchema(Schema):
    table_name = fields.String(required=True)
    mode = fields.String(load_default=MODE_EXPLORE)


class RefineTableSchema(Schema):
    data_context_id = fields.Integer(required=True)
    table_name = fields.String(required=True)
    summary = fields.String(required=True)


class GenerateTableRelationSchema(Schema):
    mode = fields.String(load_default=MODE_EXPLORE)


class GenerateEntitySchema(Schema):
    pass


class GenerateDBSummarySchema(Schema):
    pass


class GenerateDBSchema(Schema):
    pass


class Chat2DataSchema(Schema):
    data_summary_id = fields.Int(load_default=0)
    raw_question = fields.String(load_default="")
    question = fields.String(load_default="")
    flags = fields.Dict(required=False, keys=fields.String(validate=validate.OneOf(["sql_generate_mode", "auto_run", "data_to_chart"])))
    cluster_id = fields.String(load_default="")
    database = fields.String(load_default="")
    session_context_id = fields.Int(load_default=0)
    evidence = fields.String(load_default="")
    query_id = fields.String(load_default="")
    sql_generate_mode = fields.String(load_default="direct", validate=validate.OneOf(["direct", "auto_breakdown"]))



class QuestionBreakdownSchema(Schema):
    raw_question = fields.String(required=True)
    auto_run = fields.Boolean(load_default=False)
    direct_resolve = fields.Boolean(load_default=False)
    data_to_chart = fields.Boolean(load_default=True)
    query_id = fields.String(load_default="")


class Text2SQLSchema(Schema):
    task_id = fields.String(required=True)
    question_id = fields.String(required=True)


class CreateDataContextSchema(Schema):
    org_id = fields.Str(required=True)
    database_uri = fields.Str(required=True)
    force = fields.Boolean(load_default=False)
    creator = fields.String(load_default="")
    data = fields.Dict(required=True)


class CreateDataSummarySchema(Schema):
    database_uri = fields.Str(load_default="")
    creator = fields.String(load_default="")
    cluster_id = fields.String(load_default="")
    database = fields.String(load_default="")
    description = fields.String(load_default="")
    reuse = fields.Boolean(load_default=False)
    default = fields.Boolean(load_default=False)


class CreateSessionContextSchema(Schema):
    org_id = fields.Str(required=True)
    creator = fields.String(load_default="")
    data = fields.Dict(required=True)
    data_context_id = fields.Int(required=True)


class LikeOrNotSchema(Schema):
    task_id = fields.String(required=True)
    question_id = fields.String(required=True)
    like_or_not = fields.Boolean(load_default=None)


class SQLEditSchema(Schema):
    task_id = fields.String(required=True)
    question_id = fields.String(required=True)
    sql = fields.String(required=True)


class UpdateTableDescriptionSchema(Schema):
    description = fields.String(required=True)


class UpdateColumnsDescriptionSchema(Schema):
    columns = fields.List(fields.Dict(keys=fields.String(), values=fields.String()), required=True)


class TaskExecutionSchema(Schema):
    task_id = fields.String(required=True)
    question_id = fields.String(required=True)


class SQLExecutionSchema(Schema):
    session_context_id = fields.Int(required=True)
    task_id = fields.String(required=True)
    question_id = fields.String(required=True)


class QueryExecutionSchema(Schema):
    cluster_id = fields.String(required=True)
    database = fields.String(required=True, validate=validate.Length(min=1))
    sql = fields.String(required=True)
    limit = fields.Int(load_default=0)


class UserFeedbackSchema(Schema):
    session_context_id = fields.Int(required=True)
    question_id = fields.String(required=True)
    feedback = fields.String(required=True, validate=validate.OneOf(["like", "dislike"]))
    comment = fields.String(load_default="")


class UpdateDataTableDescriptionSchema(Schema):
    description = fields.String(required=True)
    columns = fields.Dict(required=True, keys=fields.String(), values=fields.String())


class UpdateDataDescriptionSchema(Schema):
    description = fields.String(required=True)
    default = fields.Boolean(required=True)


class UpsertSettingsSchema(Schema):
    openai_api_key = fields.Str(missing=None)
    llm_api_key = fields.Str(missing=None)
    language_code = fields.Str(load_default=LANGUAGE_EN, validate=validate.OneOf([LANGUAGE_EN, LANGUAGE_CN]))
    ai_model = fields.Str(load_default=Claude3.DEFAULT, validate=validate.OneOf([GPT_DEFAULT, GPT_4, GPT_4O, Claude3.HAIKU, Claude3.SONNET, Claude3.OPUS]))
    llm_model = fields.Str(load_default=Claude3.DEFAULT, validate=validate.OneOf([GPT_DEFAULT, GPT_4, GPT_4O, Claude3.HAIKU, Claude3.SONNET, Claude3.OPUS]))
    mode_type = fields.Str(load_default="", validate=validate.OneOf(["", MODE_TYPE_ASK_BILLING, MODE_TYPE_LENGJING, MODE_TYPE_TIINSIGHT, MODE_TYPE_SPIDER]))


class DeleteSettingsSchema(Schema):
    settings = fields.List(fields.Str(validate=validate.OneOf(["openai_api_key", "language", "ai_model", "mode_type"])), required=True)


class CreateSessionSchema(Schema):
    cluster_id = fields.String(load_default="")
    database = fields.String(load_default="")
    name = fields.String(load_default="")
    database_uri = fields.String(load_default="")
    session_context_id = fields.Int(load_default=0)


class UpdateSessionSchema(Schema):
    name = fields.String(load_default="")


class RefineSQLSchema(Schema):
    cluster_id = fields.String(required=True)
    database = fields.String(required=True, validate=validate.Length(min=1))
    sql = fields.String(required=True)
    raw_question = fields.String(load_default="")
    question = fields.String(load_default="")
    feedback = fields.String(load_default="")
    sql_error = fields.String(load_default="")


class CreateDomainItemSchema(Schema):
    type = fields.Str(required=True)
    meta_data = fields.Dict(required=False)
    raw_data = fields.Dict(required=True)


class UpdateDomainItemSchema(Schema):
    meta_data = fields.Dict(required=True)
    raw_data = fields.Dict(required=True)


class MultiRoundChat2SQLSchema(Schema):
    question = fields.String(required=True, validate=validate.Length(min=1))
    feedback_answer_id = fields.String(load_default="")
    feedback_task_id = fields.String(load_default="")
    evidence = fields.String(load_default="")
    sql_generate_mode = fields.String(load_default="direct", validate=validate.OneOf(["direct", "auto_breakdown"]))


class MultiRoundFeedbackSchema(Schema):
    question = fields.String(required=True, validate=validate.Length(min=1))
    session_id = fields.String(required=True)
    feedback_answer_id = fields.String(load_default="")
    feedback_task_id = fields.String(load_default="")


class CreateKnowledgeBaseSchema(Schema):
    cluster_id = fields.String(required=True)
    database = fields.String(required=True)
    description = fields.String(load_default="")


class UpdateKnowledgeBaseSchema(Schema):
    description = fields.String(required=True)
    default = fields.Boolean(required=True)


class UpdateColumnDescriptionSchema(Schema):
    name = fields.String(required=True)
    description = fields.String(required=True)


class SuggestQuestionSchema(Schema):
    cluster_id = fields.String(load_default="")
    database = fields.String(load_default="")
    database_uri = fields.String(load_default="")