#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import json
import logging
from numbers import Real

from quart import Response, request
from api.db.services.dialog_service import async_ask
from api.apps import current_user, login_required

from api.constants import DATASET_NAME_LIMIT
from api.db.db_models import DB
from api.db.services import duplicate_name
from api.db.services.search_service import SearchService
from api.db.services.user_service import TenantService, UserTenantService
from common.misc_utils import get_uuid
from common.constants import RetCode, StatusEnum
from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request
from api.utils.pagination_utils import validate_rest_api_page_size


def _full_text_weight(vector_similarity_weight):
    if isinstance(vector_similarity_weight, Real):
        return 1 - vector_similarity_weight
    return None


@manager.route("/searches", methods=["POST"])  # noqa: F821
@login_required
@validate_request("name")
async def create():
    req = await get_request_json()
    search_name = req["name"]
    description = req.get("description", "")
    if not isinstance(search_name, str):
        return get_data_error_result(message="Search name must be string.")
    if search_name.strip() == "":
        return get_data_error_result(message="Search name can't be empty.")
    if len(search_name.encode("utf-8")) > 255:
        return get_data_error_result(message=f"Search name length is {len(search_name)} which is large than 255.")
    e, _ = TenantService.get_by_id(current_user.id)
    if not e:
        return get_data_error_result(message="Authorized identity.")

    search_name = search_name.strip()
    search_name = duplicate_name(SearchService.query, name=search_name, tenant_id=current_user.id, status=StatusEnum.VALID.value)

    req["id"] = get_uuid()
    req["name"] = search_name
    req["description"] = description
    req["tenant_id"] = current_user.id
    req["created_by"] = current_user.id
    with DB.atomic():
        try:
            if not SearchService.save(**req):
                return get_data_error_result()
            return get_json_result(data={"search_id": req["id"]})
        except Exception as e:
            return server_error_response(e)


@manager.route("/searches", methods=["GET"])  # noqa: F821
@login_required
def list_searches():
    keywords = request.args.get("keywords", "")
    page_number = int(request.args.get("page", 0))
    items_per_page = validate_rest_api_page_size(int(request.args.get("page_size", 0)))
    orderby = request.args.get("orderby", "create_time")
    desc = request.args.get("desc", "true").lower() != "false"
    owner_ids = request.args.getlist("owner_ids")

    try:
        if not owner_ids:
            tenants = []
            search_apps, total = SearchService.get_by_tenant_ids(tenants, current_user.id, page_number, items_per_page, orderby, desc, keywords)
        else:
            search_apps, total = SearchService.get_by_tenant_ids(owner_ids, current_user.id, 0, 0, orderby, desc, keywords)
            search_apps = [s for s in search_apps if s["tenant_id"] in owner_ids]
            total = len(search_apps)
            if page_number and items_per_page:
                search_apps = search_apps[(page_number - 1) * items_per_page: page_number * items_per_page]
        return get_json_result(data={"search_apps": search_apps, "total": total})
    except Exception as e:
        return server_error_response(e)


@manager.route("/searches/<search_id>", methods=["GET"])  # noqa: F821
@login_required
def detail(search_id):
    try:
        tenants = UserTenantService.query(user_id=current_user.id)
        for tenant in tenants:
            if SearchService.query(tenant_id=tenant.tenant_id, id=search_id):
                break
        else:
            return get_json_result(data=False, message="Has no permission for this operation.", code=RetCode.OPERATING_ERROR)

        search = SearchService.get_detail(search_id)
        if not search:
            return get_data_error_result(message="Can't find this Search App!")
        return get_json_result(data=search)
    except Exception as e:
        return server_error_response(e)


@manager.route("/searches/<search_id>", methods=["PUT"])  # noqa: F821
@login_required
@validate_request("name", "search_config")
async def update(search_id):
    req = await get_request_json()
    if not isinstance(req["name"], str):
        return get_data_error_result(message="Search name must be string.")
    if req["name"].strip() == "":
        return get_data_error_result(message="Search name can't be empty.")
    if len(req["name"].encode("utf-8")) > DATASET_NAME_LIMIT:
        return get_data_error_result(message=f"Search name length is {len(req['name'])} which is large than {DATASET_NAME_LIMIT}")
    req["name"] = req["name"].strip()

    e, _ = TenantService.get_by_id(current_user.id)
    if not e:
        return get_data_error_result(message="Authorized identity.")

    if not SearchService.accessible4deletion(search_id, current_user.id):
        return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)

    try:
        search_app = SearchService.query(tenant_id=current_user.id, id=search_id)[0]
        if not search_app:
            return get_json_result(data=False, message=f"Cannot find search {search_id}", code=RetCode.DATA_ERROR)

        if req["name"].lower() != search_app.name.lower() and len(SearchService.query(name=req["name"], tenant_id=current_user.id, status=StatusEnum.VALID.value)) >= 1:
            return get_data_error_result(message="Duplicated search name.")

        current_config = search_app.search_config or {}
        new_config = req["search_config"]
        if not isinstance(new_config, dict):
            return get_data_error_result(message="search_config must be a JSON object")
        req["search_config"] = {**current_config, **new_config}
        logging.debug(
            "Search update weight: search_id=%s user_id=%s "
            "incoming_vector_similarity_weight=%s stored_vector_similarity_weight=%s "
            "stored_full_text_weight=%s",
            search_id,
            current_user.id,
            new_config.get("vector_similarity_weight"),
            req["search_config"].get("vector_similarity_weight"),
            _full_text_weight(req["search_config"].get("vector_similarity_weight", 0.3)),
        )

        for field in ("search_id", "tenant_id", "created_by", "update_time", "id"):
            req.pop(field, None)

        updated = SearchService.update_by_id(search_id, req)
        if not updated:
            return get_data_error_result(message="Failed to update search")

        e, updated_search = SearchService.get_by_id(search_id)
        if not e:
            return get_data_error_result(message="Failed to fetch updated search")

        return get_json_result(data=updated_search.to_dict())

    except Exception as e:
        return server_error_response(e)


@manager.route("/searches/<search_id>", methods=["DELETE"])  # noqa: F821
@login_required
def delete_search(search_id):
    if not SearchService.accessible4deletion(search_id, current_user.id):
        return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)

    try:
        if not SearchService.delete_by_id(search_id):
            return get_data_error_result(message=f"Failed to delete search App {search_id}")
        return get_json_result(data=True)
    except Exception as e:
        return server_error_response(e)


@manager.route("/searches/<search_id>/completion", methods=["POST"])  # noqa: F821
@manager.route("/searches/<search_id>/completions", methods=["POST"])  # noqa: F821
@login_required
@validate_request("question")
async def completion(search_id):
    if not SearchService.accessible4deletion(search_id, current_user.id):
        return get_json_result(
            data=False,
            message="No authorization.",
            code=RetCode.AUTHENTICATION_ERROR,
        )

    req = await get_request_json()
    uid = current_user.id
    search_app = SearchService.get_detail(search_id)
    if not search_app:
        return get_data_error_result(message=f"Cannot find search {search_id}")

    search_config = search_app.get("search_config", {})
    logging.debug(
        "Search completion loaded weight: search_id=%s user_id=%s "
        "stored_vector_similarity_weight=%s stored_full_text_weight=%s",
        search_id,
        uid,
        search_config.get("vector_similarity_weight", 0.3),
        _full_text_weight(search_config.get("vector_similarity_weight", 0.3)),
    )
    kb_ids = search_config.get("kb_ids") or req.get("kb_ids") or []
    if not kb_ids:
        return get_data_error_result(message="`kb_ids` is required.")

    async def stream():
        nonlocal req, uid, kb_ids, search_config
        try:
            async for ans in async_ask(req["question"], kb_ids, uid, search_config=search_config, search_id=search_id):
                yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
        except Exception as ex:
            yield "data:" + json.dumps(
                {"code": 500, "message": str(ex), "data": {"answer": "**ERROR**: " + str(ex), "reference": []}},
                ensure_ascii=False,
            ) + "\n\n"
        yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"

    resp = Response(stream(), mimetype="text/event-stream")
    resp.headers.add_header("Cache-control", "no-cache")
    resp.headers.add_header("Connection", "keep-alive")
    resp.headers.add_header("X-Accel-Buffering", "no")
    resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
    return resp