#
#  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
import os
from datetime import datetime, timedelta

from peewee import fn

from api.db import VALID_PIPELINE_TASK_TYPES
from api.db.db_models import DB, Document, PipelineOperationLog
from api.db.services.canvas_service import UserCanvasService
from api.db.services.common_service import CommonService
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID, TaskService
from common.constants import PipelineTaskType, TaskStatus
from common.misc_utils import get_uuid
from common.time_utils import current_timestamp, datetime_format


# KB-level fan-out pipeline task types (task row carries a fake doc_id; the real
# participants live in task["doc_ids"]) → the KB ``<type>_task_finish_at`` column
# stamped when the task completes. Membership also marks a task as KB-scoped so
# the per-document progress update is skipped.
_PIPELINE_TASK_TYPE_TO_FINISH_FIELD = {
    PipelineTaskType.GRAPH_RAG: "graphrag_task_finish_at",
    PipelineTaskType.RAPTOR: "raptor_task_finish_at",
    PipelineTaskType.MINDMAP: "mindmap_task_finish_at",
    PipelineTaskType.ARTIFACT: "artifact_task_finish_at",
    PipelineTaskType.SKILL: "skill_task_finish_at",
    PipelineTaskType.STRUCTURE_GRAPH: "structure_graph_task_finish_at",
    PipelineTaskType.STRUCTURE_MINDMAP: "structure_mindmap_task_finish_at",
    PipelineTaskType.TIMELINE: "timeline_task_finish_at",
    PipelineTaskType.SESSION_GRAPH: "session_graph_task_finish_at",
    PipelineTaskType.SESSION_ESSENCE: "session_essence_task_finish_at",
    PipelineTaskType.STRUCTURE: "structure_task_finish_at",
}


class PipelineOperationLogService(CommonService):
    model = PipelineOperationLog

    @classmethod
    def _final_operation_statuses(cls, operation_status):
        final_statuses = [TaskStatus.DONE.value, TaskStatus.FAIL.value]
        if not operation_status:
            return final_statuses
        requested = {status.value if isinstance(status, TaskStatus) else str(status) for status in operation_status}
        return [status for status in final_statuses if status in requested]

    @classmethod
    def _is_final_progress(cls, progress):
        return progress == 1 or progress == -1

    @classmethod
    def get_file_logs_fields(cls):
        return [
            cls.model.id,
            cls.model.document_id,
            cls.model.tenant_id,
            cls.model.kb_id,
            cls.model.pipeline_id,
            cls.model.pipeline_title,
            cls.model.parser_id,
            cls.model.document_name,
            cls.model.document_suffix,
            cls.model.document_type,
            cls.model.source_from,
            cls.model.progress,
            cls.model.progress_msg,
            cls.model.process_begin_at,
            cls.model.process_duration,
            cls.model.dsl,
            cls.model.task_type,
            cls.model.operation_status,
            cls.model.avatar,
            cls.model.status,
            cls.model.create_time,
            cls.model.create_date,
            cls.model.update_time,
            cls.model.update_date,
        ]

    @classmethod
    def get_dataset_logs_fields(cls):
        return [
            cls.model.id,
            cls.model.tenant_id,
            cls.model.kb_id,
            cls.model.progress,
            cls.model.progress_msg,
            cls.model.process_begin_at,
            cls.model.process_duration,
            cls.model.task_type,
            cls.model.operation_status,
            cls.model.avatar,
            cls.model.status,
            cls.model.create_time,
            cls.model.create_date,
            cls.model.update_time,
            cls.model.update_date,
        ]

    @classmethod
    def save(cls, **kwargs):
        """
        wrap this function in a transaction
        """
        sample_obj = cls.model(**kwargs).save(force_insert=True)
        return sample_obj

    @classmethod
    @DB.connection_context()
    def create(cls, document_id, pipeline_id, task_type, task_id=None, referred_document_id=None, dsl: str = "{}"):
        if document_id != GRAPH_RAPTOR_FAKE_DOC_ID:
            referred_document_id = document_id

        # no need to update document for KB-level fan-out tasks
        if task_type not in _PIPELINE_TASK_TYPE_TO_FINISH_FIELD:
            ok, document = DocumentService.get_by_id(referred_document_id)
            if not ok:
                logging.warning(f"Document for referred_document_id {referred_document_id} not found")
                return None
            DocumentService.update_progress_immediately([document.to_dict()])

        ok, document = DocumentService.get_by_id(referred_document_id)
        if not ok:
            logging.warning(f"Document for referred_document_id {referred_document_id} not found")
            return None

        # From document
        title = document.parser_id
        avatar = document.thumbnail
        document_name = document.name
        operation_status = document.run
        progress = document.progress
        progress_msg = document.progress_msg
        process_begin_at = document.process_begin_at
        process_duration = document.process_duration

        if pipeline_id:
            ok, user_pipeline = UserCanvasService.get_by_id(pipeline_id)
            if not ok:
                raise RuntimeError(f"Pipeline {pipeline_id} not found")
            tenant_id = user_pipeline.user_id
            title = user_pipeline.title
            avatar = user_pipeline.avatar
        else:
            ok, kb_info = KnowledgebaseService.get_by_id(document.kb_id)
            if not ok:
                raise RuntimeError(f"Cannot find dataset {document.kb_id} for referred_document {referred_document_id}")
            tenant_id = kb_info.tenant_id

        if task_type not in VALID_PIPELINE_TASK_TYPES:
            raise ValueError(f"Invalid task type: {task_type}")

        if task_type in _PIPELINE_TASK_TYPE_TO_FINISH_FIELD:
            # query task to get progress information from task
            ok, task = TaskService.get_by_id(task_id)
            if not ok:
                raise RuntimeError(f"Task not found for dataset {document.kb_id}")
            title = task_type
            document_name = task_type
            operation_status = TaskStatus.DONE if task.progress == 1 else TaskStatus.FAIL
            progress = task.progress
            progress_msg = task.progress_msg
            process_begin_at = task.begin_at
            process_duration = task.process_duration

            if not cls._is_final_progress(progress):
                logging.info("Skip non-final dataset pipeline operation log task_id=%s task_type=%s progress=%s", task_id, task_type, progress)
                return None

            finish_at = process_begin_at + timedelta(seconds=process_duration)
            KnowledgebaseService.update_by_id(
                document.kb_id,
                {_PIPELINE_TASK_TYPE_TO_FINISH_FIELD[task_type]: finish_at},
            )
        elif not cls._is_final_progress(progress):
            logging.info("Skip non-final file pipeline operation log document_id=%s task_type=%s progress=%s", document_id, task_type, progress)
            return None

        log = dict(
            id=get_uuid(),
            document_id=document_id,  # GRAPH_RAPTOR_FAKE_DOC_ID or real document_id
            tenant_id=tenant_id,
            kb_id=document.kb_id,
            pipeline_id=pipeline_id,
            pipeline_title=title,
            parser_id=document.parser_id,
            document_name=document_name,
            document_suffix=document.suffix,
            document_type=document.type,
            source_from=document.source_type.split("/")[0],
            progress=progress,
            progress_msg=progress_msg,
            process_begin_at=process_begin_at,
            process_duration=process_duration,
            dsl=json.loads(dsl),
            task_type=task_type,
            operation_status=operation_status,
            avatar=avatar,
        )
        timestamp = current_timestamp()
        datetime_now = datetime_format(datetime.now())
        log["create_time"] = timestamp
        log["create_date"] = datetime_now
        log["update_time"] = timestamp
        log["update_date"] = datetime_now
        with DB.atomic():
            obj = cls.save(**log)

            limit = int(os.getenv("PIPELINE_OPERATION_LOG_LIMIT", 1000))
            total = cls.model.select().where(cls.model.kb_id == document.kb_id).count()

            if total > limit:
                keep_ids = [m.id for m in cls.model.select(cls.model.id).where(cls.model.kb_id == document.kb_id).order_by(cls.model.create_time.desc()).limit(limit)]

                deleted = cls.model.delete().where(cls.model.kb_id == document.kb_id, cls.model.id.not_in(keep_ids)).execute()
                logging.info(f"[PipelineOperationLogService] Cleaned {deleted} old logs, kept latest {limit} for {document.kb_id}")

        return obj

    @classmethod
    @DB.connection_context()
    def record_pipeline_operation(cls, document_id, pipeline_id, task_type, task_id=None, referred_document_id=None):
        return cls.create(document_id=document_id, pipeline_id=pipeline_id, task_type=task_type, task_id=task_id, referred_document_id=referred_document_id)

    @classmethod
    @DB.connection_context()
    def get_file_logs_by_kb_id(cls, kb_id, page_number, items_per_page, orderby, desc, keywords, operation_status, types, suffix, create_date_from=None, create_date_to=None):
        fields = cls.get_file_logs_fields()
        if keywords:
            logs = cls.model.select(*fields).where((cls.model.kb_id == kb_id), (fn.LOWER(cls.model.document_name).contains(keywords.lower())))
        else:
            logs = cls.model.select(*fields).where(cls.model.kb_id == kb_id)

        logs = logs.where(cls.model.document_id != GRAPH_RAPTOR_FAKE_DOC_ID)

        logs = logs.where(cls.model.operation_status.in_(cls._final_operation_statuses(operation_status)))
        if types:
            logs = logs.where(cls.model.document_type.in_(types))
        if suffix:
            logs = logs.where(cls.model.document_suffix.in_(suffix))
        if create_date_from:
            logs = logs.where(cls.model.create_date >= create_date_from)
        if create_date_to:
            logs = logs.where(cls.model.create_date <= create_date_to)

        count = logs.count()
        if desc:
            logs = logs.order_by(cls.model.getter_by(orderby).desc())
        else:
            logs = logs.order_by(cls.model.getter_by(orderby).asc())

        if page_number and items_per_page:
            logs = logs.paginate(page_number, items_per_page)

        return list(logs.dicts()), count

    @classmethod
    @DB.connection_context()
    def get_documents_info(cls, id):
        fields = [Document.id, Document.name, Document.progress, Document.kb_id]
        return cls.model.select(*fields).join(Document, on=(cls.model.document_id == Document.id)).where(cls.model.id == id).dicts()

    @classmethod
    @DB.connection_context()
    def get_dataset_logs_by_kb_id(cls, kb_id, page_number, items_per_page, orderby, desc, operation_status, create_date_from=None, create_date_to=None, keywords=None):
        fields = cls.get_dataset_logs_fields()
        if keywords:
            logs = cls.model.select(*fields).where((cls.model.kb_id == kb_id), (cls.model.document_id == GRAPH_RAPTOR_FAKE_DOC_ID), (fn.LOWER(cls.model.document_name).contains(keywords.lower())))
        else:
            logs = cls.model.select(*fields).where((cls.model.kb_id == kb_id), (cls.model.document_id == GRAPH_RAPTOR_FAKE_DOC_ID))

        logs = logs.where(cls.model.operation_status.in_(cls._final_operation_statuses(operation_status)))
        if create_date_from:
            logs = logs.where(cls.model.create_date >= create_date_from)
        if create_date_to:
            logs = logs.where(cls.model.create_date <= create_date_to)

        count = logs.count()
        if desc:
            logs = logs.order_by(cls.model.getter_by(orderby).desc())
        else:
            logs = logs.order_by(cls.model.getter_by(orderby).asc())

        if page_number and items_per_page:
            logs = logs.paginate(page_number, items_per_page)

        return list(logs.dicts()), count