"""数据合成服务 - 核心业务逻辑层"""
import asyncio
from contextlib import asynccontextmanager

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.db.models.base_entity import LineageNode, LineageEdge
from app.db.models.data_synthesis import (
    DataSynthInstance,
    DataSynthesisFileInstance,
    DataSynthesisChunkInstance,
    SynthesisData,
)
from app.db.models.dataset_management import DatasetFiles, Dataset
from app.db.session import logger, AsyncSessionLocal
from app.module.generation.schema.generation import Config, SyntheConfig
from app.module.generation.service.chunk_processor import ChunkProcessor
from app.module.generation.service.qa_generator import QAGenerator
from app.module.generation.service.qa_generator import _filter_docs_by_size
from app.module.generation.service.task_executor import run_in_thread
from app.module.shared.common.document_loaders import load_documents
from app.module.shared.common.text_split import DocumentSplitter
from app.module.shared.llm import LLMFactory
from app.module.system.service.common_service import get_model_by_id
from app.module.shared.common.lineage import LineageService
from app.module.shared.schema import NodeType, EdgeType


class GenerationService:
    """数据合成服务 - 使用模块化架构"""

    def __init__(self, db: AsyncSession):
        self.db = db
        self.chunk_processor = ChunkProcessor(db)
        self.qa_generator = QAGenerator(db)

    @asynccontextmanager
    async def _session_scope(self):
        """
        创建独立的数据库会话上下文管理器
        用于长时间运行的任务，避免长时间占用连接
        """
        session = AsyncSessionLocal()
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()

    def _load_and_split_sync(
        self, file_path: str, chunk_size: int, chunk_overlap: int
    ):
        """
        同步方法：加载并切分文件
        在线程池中执行，不阻塞事件循环
        """
        try:
            docs = load_documents(file_path)
            split_docs = DocumentSplitter.auto_split(docs, chunk_size, chunk_overlap)
            return _filter_docs_by_size(split_docs, chunk_size)
        except Exception as e:
            logger.error(f"Error loading or splitting file {file_path}: {e}")
            raise

    async def process_task(self, task_id: str):
        """
        处理数据合成任务入口 - 使用独立会话和线程池

        Args:
            task_id: 合成任务ID
        """
        # 使用独立的数据库会话，不与API请求共享
        async with self._session_scope() as session:
            synth_task: DataSynthInstance | None = await session.get(
                DataSynthInstance, task_id
            )
            if not synth_task:
                logger.error(f"Synthesis task {task_id} not found, abort processing")
                return

            logger.info(f"Start processing synthesis task {task_id}")

            # 更新任务状态为处理中
            synth_task.status = "processing"
            await session.commit()

            # 从 synth_config 中读取 max_qa_pairs
            max_qa_pairs = self._parse_max_qa_pairs(synth_task)

            # 获取任务关联的文件ID列表
            file_ids = await self._get_file_ids_for_task(session, task_id)
            if not file_ids:
                logger.warning(
                    f"No files associated with task {task_id}, abort processing"
                )
                await self._mark_task_failed(session, task_id, "no_files_associated")
                return

            # 逐个文件处理
            processed_count = 0
            total_task_chunks = 0
            processed_task_chunks = 0
            for file_id in file_ids:
                try:
                    file_task = await self._get_or_create_file_instance(
                        session, str(synth_task.id), file_id
                    )
                    success = await self._process_single_file(
                        session, synth_task, file_id, max_qa_pairs=max_qa_pairs
                    )
                    # 累加任务级别的切片统计
                    if file_task:
                        total_task_chunks += file_task.total_chunks or 0
                        processed_task_chunks += file_task.processed_chunks or 0
                except Exception as e:
                    logger.exception(
                        f"Unexpected error when processing file {file_id} for task {task_id}: {e}"
                    )
                    await self._mark_file_failed(
                        session, str(synth_task.id), file_id, str(e)
                    )
                    success = False

                if success:
                    processed_count += 1
                    synth_task.processed_files = processed_count
                    # 更新任务级别的切片统计
                    synth_task.total_chunks = total_task_chunks
                    synth_task.processed_chunks = processed_task_chunks
                    await session.commit()

            # 更新最终任务状态
            if processed_count == len(file_ids):
                synth_task.status = "completed"
            else:
                synth_task.status = "partially_completed"
            # 确保任务完成时切片统计准确
            if synth_task.status in ("completed", "partially_completed"):
                synth_task.total_chunks = total_task_chunks
                synth_task.processed_chunks = processed_task_chunks
            await session.commit()

            logger.info(f"Finished processing synthesis task {synth_task.id}")

    def _parse_max_qa_pairs(self, synth_task: DataSynthInstance) -> int | None:
        """解析最大QA对数量配置"""
        try:
            cfg = Config(**(synth_task.synth_config or {}))
            max_qa_pairs = (
                cfg.max_qa_pairs
                if (cfg and cfg.max_qa_pairs and cfg.max_qa_pairs > 0)
                else None
            )
        except Exception:
            max_qa_pairs = None
        return max_qa_pairs

    async def _process_single_file(
        self,
        session: AsyncSession,
        synth_task: DataSynthInstance,
        file_id: str,
        max_qa_pairs: int | None = None,
    ) -> bool:
        """
        处理单个源文件 - 使用线程池执行阻塞操作

        Args:
            session: 数据库会话
            synth_task: 合成任务实例
            file_id: 源文件ID
            max_qa_pairs: 最大QA对数量限制

        Returns:
            处理是否成功
        """
        # 解析文件路径
        file_path = await self._resolve_file_path(session, file_id)
        if not file_path:
            logger.warning(f"File path not found for file_id={file_id}, skip")
            await self._mark_file_failed(
                session, str(synth_task.id), file_id, "file_path_not_found"
            )
            return False

        logger.info(f"Processing file_id={file_id}, path={file_path}")

        # 解析配置
        try:
            config = Config(**(synth_task.synth_config or {}))
        except Exception as e:
            logger.error(f"Invalid synth_config for task={synth_task.id}: {e}")
            await self._mark_file_failed(
                session, str(synth_task.id), file_id, "invalid_synth_config"
            )
            return False

        # 1. 加载并切片 - 在线程池中执行阻塞操作
        try:
            chunks = await run_in_thread(
                self._load_and_split_sync,
                file_path,
                config.text_split_config.chunk_size,
                config.text_split_config.chunk_overlap,
            )
        except Exception as e:
            logger.error(f"Failed to load and split file {file_id}: {e}")
            await self._mark_file_failed(
                session, str(synth_task.id), file_id, f"load_split_failed: {e}"
            )
            return False

        if not chunks:
            logger.warning(f"No chunks generated for file_id={file_id}")
            await self._mark_file_failed(
                session, str(synth_task.id), file_id, "no_chunks_generated"
            )
            return False

        logger.info(f"File {file_id} split into {len(chunks)} chunks")

        # 2. 获取文件实例并持久化切片
        file_task = await self._get_or_create_file_instance(
            session,
            synthesis_task_id=str(synth_task.id),
            source_file_id=file_id,
        )
        if not file_task:
            logger.error(
                f"DataSynthesisFileInstance not found for task={synth_task.id}, file_id={file_id}"
            )
            await self._mark_file_failed(
                session, str(synth_task.id), file_id, "file_instance_not_found"
            )
            return False

        # 使用 ChunkProcessor 持久化切片
        chunk_processor = ChunkProcessor(session)
        await chunk_processor.persist_chunks(synth_task, file_task, file_id, chunks)
        total_chunks = len(chunks)
        del chunks  # 释放内存

        # 3. 验证问答配置
        question_cfg: SyntheConfig | None = config.question_synth_config
        answer_cfg: SyntheConfig | None = config.answer_synth_config
        if not question_cfg or not answer_cfg:
            logger.error(
                f"Question/Answer synth config missing for task={synth_task.id}, file={file_id}"
            )
            await self._mark_file_failed(
                session, str(synth_task.id), file_id, "qa_config_missing"
            )
            return False

        logger.info(
            f"Start QA generation for task={synth_task.id}, file={file_id}, total_chunks={total_chunks}"
        )

        # 4. 创建模型客户端
        question_model = await get_model_by_id(session, question_cfg.model_id)
        answer_model = await get_model_by_id(session, answer_cfg.model_id)
        question_chat = LLMFactory.create_chat(
            question_model.model_name, question_model.base_url, question_model.api_key
        )
        answer_chat = LLMFactory.create_chat(
            answer_model.model_name, answer_model.base_url, answer_model.api_key
        )

        # 5. 分批次处理切片
        await self._process_chunks_in_batches(
            session=session,
            file_task=file_task,
            total_chunks=total_chunks,
            question_cfg=question_cfg,
            answer_cfg=answer_cfg,
            question_chat=question_chat,
            answer_chat=answer_chat,
            synth_task_id=str(synth_task.id),
            max_qa_pairs=max_qa_pairs,
        )

        # 6. 标记文件任务完成
        file_task.status = "completed"
        await session.commit()
        await session.refresh(file_task)
        return True

    async def _process_chunks_in_batches(
        self,
        session: AsyncSession,
        file_task: DataSynthesisFileInstance,
        total_chunks: int,
        question_cfg: SyntheConfig,
        answer_cfg: SyntheConfig,
        question_chat,
        answer_chat,
        synth_task_id: str,
        max_qa_pairs: int | None,
    ) -> None:
        """分批次处理切片

        Args:
            session: 数据库会话
            file_task: 文件任务实例
            total_chunks: 总切片数
            question_cfg: 问题生成配置
            answer_cfg: 答案生成配置
            question_chat: 问题生成模型
            answer_chat: 答案生成模型
            synth_task_id: 合成任务ID
            max_qa_pairs: 最大QA对数量限制
        """
        batch_size = 50  # 减小批次大小，降低内存占用
        current_index = 1

        chunk_processor = ChunkProcessor(session)
        qa_generator = QAGenerator(session)

        while current_index <= total_chunks:
            end_index = min(current_index + batch_size - 1, total_chunks)

            # 使用 ChunkProcessor 加载切片批次
            chunk_batch = await chunk_processor.load_chunk_batch(
                file_task_id=file_task.id,
                start_index=current_index,
                end_index=end_index,
            )

            if not chunk_batch:
                logger.warning(
                    f"Empty chunk batch loaded for file={file_task.id}, range=[{current_index}, {end_index}]"
                )
                current_index = end_index + 1
                continue

            # 顺序处理每个切片，避免并发过高
            for chunk in chunk_batch:
                try:
                    success = await self._process_single_chunk(
                        session=session,
                        file_task=file_task,
                        chunk=chunk,
                        question_cfg=question_cfg,
                        answer_cfg=answer_cfg,
                        question_chat=question_chat,
                        answer_chat=answer_chat,
                        synth_task_id=synth_task_id,
                        max_qa_pairs=max_qa_pairs,
                        qa_generator=qa_generator,
                    )
                    if success:
                        # 更新已处理的切片计数
                        file_task.processed_chunks = (file_task.processed_chunks or 0) + 1
                        # 实时刷新进度到数据库
                        await session.commit()
                        await session.refresh(file_task)
                        logger.info(f"File {file_task.id} progress: {file_task.processed_chunks}/{file_task.total_chunks} chunks processed")
                except Exception as e:
                    logger.error(f"Error processing chunk {chunk.id}: {e}")
                    # 即使处理失败，也要更新进度
                    file_task.processed_chunks = (file_task.processed_chunks or 0) + 1
                    await session.commit()
                    await session.refresh(file_task)

                # 短暂释放控制权，让出事件循环，处理其他请求
                await asyncio.sleep(0)

            current_index = end_index + 1

    async def _process_single_chunk(
        self,
        session: AsyncSession,
        file_task: DataSynthesisFileInstance,
        chunk: DataSynthesisChunkInstance,
        question_cfg: SyntheConfig,
        answer_cfg: SyntheConfig,
        question_chat,
        answer_chat,
        synth_task_id: str,
        max_qa_pairs: int | None,
        qa_generator: QAGenerator,
    ) -> bool:
        """处理单个切片

        Args:
            session: 数据库会话
            file_task: 文件任务实例
            chunk: 切片实例
            question_cfg: 问题生成配置
            answer_cfg: 答案生成配置
            question_chat: 问题生成模型
            answer_chat: 答案生成模型
            synth_task_id: 合成任务ID
            max_qa_pairs: 最大QA对数量限制
            qa_generator: QA生成器

        Returns:
            处理是否成功
        """
        # 检查QA对数量上限
        if max_qa_pairs and max_qa_pairs > 0:
            if await qa_generator.check_qa_limit(synth_task_id, max_qa_pairs):
                logger.info(f"max_qa_pairs reached, skipping chunk {chunk.chunk_index}")
                file_task.status = "completed"
                if file_task.total_chunks is not None:
                    file_task.processed_chunks = file_task.total_chunks
                await session.commit()
                await session.refresh(file_task)
                return False

        # 使用 QAGenerator 处理问答生成
        success_count = await qa_generator.process_chunk_qa(
            file_task=file_task,
            chunk=chunk,
            question_cfg=question_cfg,
            answer_cfg=answer_cfg,
            question_chat=question_chat,
            answer_chat=answer_chat,
        )

        return success_count > 0

    async def _resolve_file_path(self, session: AsyncSession, file_id: str) -> str | None:
        """根据文件ID获取文件路径

        Args:
            session: 数据库会话
            file_id: 文件ID

        Returns:
            文件路径
        """
        result = await session.execute(
            select(DatasetFiles).where(DatasetFiles.id == file_id)
        )
        file_obj = result.scalar_one_or_none()
        if not file_obj:
            return None
        return file_obj.file_path

    async def _get_or_create_file_instance(
        self,
        session: AsyncSession,
        synthesis_task_id: str,
        source_file_id: str,
    ) -> DataSynthesisFileInstance | None:
        """获取或创建文件任务实例

        Args:
            session: 数据库会话
            synthesis_task_id: 合成任务ID
            source_file_id: 源文件ID

        Returns:
            文件任务实例
        """
        result = await session.execute(
            select(DataSynthesisFileInstance).where(
                DataSynthesisFileInstance.synthesis_instance_id == synthesis_task_id,
                DataSynthesisFileInstance.source_file_id == source_file_id,
            )
        )
        return result.scalar_one_or_none()

    async def _mark_file_failed(
        self,
        session: AsyncSession,
        synth_task_id: str,
        file_id: str,
        reason: str | None = None,
    ) -> None:
        """标记文件任务失败

        Args:
            session: 数据库会话
            synth_task_id: 合成任务ID
            file_id: 文件ID
            reason: 失败原因
        """
        try:
            result = await session.execute(
                select(DataSynthesisFileInstance).where(
                    DataSynthesisFileInstance.synthesis_instance_id == synth_task_id,
                    DataSynthesisFileInstance.source_file_id == file_id,
                )
            )
            file_task = result.scalar_one_or_none()
            if not file_task:
                logger.warning(
                    f"Failed to mark file as failed: no DataSynthesisFileInstance found for task={synth_task_id}, file_id={file_id}, reason={reason}"
                )
                return

            file_task.status = "failed"
            await session.commit()
            await session.refresh(file_task)
            logger.info(
                f"Marked file task as failed for task={synth_task_id}, file_id={file_id}, reason={reason}"
            )
        except Exception as e:
            logger.exception(
                f"Unexpected error when marking file failed for task={synth_task_id}, file_id={file_id}, original_reason={reason}, error={e}"
            )

    async def _get_file_ids_for_task(
        self, session: AsyncSession, synth_task_id: str
    ):
        """获取任务关联的文件ID列表

        Args:
            session: 数据库会话
            synth_task_id: 合成任务ID

        Returns:
            文件ID列表
        """
        result = await session.execute(
            select(DataSynthesisFileInstance.source_file_id).where(
                DataSynthesisFileInstance.synthesis_instance_id == synth_task_id
            )
        )
        return result.scalars().all()

    async def _mark_task_failed(
        self,
        session: AsyncSession,
        task_id: str,
        reason: str | None = None,
    ) -> None:
        """标记任务失败

        Args:
            session: 数据库会话
            task_id: 任务ID
            reason: 失败原因
        """
        try:
            task = await session.get(DataSynthInstance, task_id)
            if task:
                task.status = "failed"
                await session.commit()
                logger.info(
                    f"Marked task {task_id} as failed, reason={reason}"
                )
        except Exception as e:
            logger.exception(f"Error marking task {task_id} as failed: {e}")

    async def add_synthesis_to_graph(
        self, db: AsyncSession, task_id: str, dest_dataset_id: str
    ) -> None:
        """记录数据合成血缘关系

        Args:
            db: 数据库会话
            task_id: 任务ID
            dest_dataset_id: 目标数据集ID
        """
        try:
            task = await db.get(DataSynthInstance, task_id)
            src_dataset_result = await db.execute(
                select(DatasetFiles.dataset_id)
                .join(
                    DataSynthesisFileInstance,
                    DatasetFiles.id == DataSynthesisFileInstance.source_file_id,
                )
                .where(DataSynthesisFileInstance.synthesis_instance_id == task_id)
                .limit(1)
            )
            src_dataset_id = src_dataset_result.scalar_one_or_none()
            src_dataset = await db.get(Dataset, src_dataset_id)
            dst_dataset = await db.get(Dataset, dest_dataset_id)

            if not task or not dst_dataset:
                logger.warning("Missing task or destination dataset for lineage graph")
                return

            src_node = LineageNode(
                id=src_dataset.id,
                node_type=NodeType.DATASET.value,
                name=src_dataset.name,
                description=src_dataset.description,
            )
            dest_node = LineageNode(
                id=dst_dataset.id,
                node_type=NodeType.DATASET.value,
                name=dst_dataset.name,
                description=dst_dataset.description,
            )
            synthesis_edge = LineageEdge(
                process_id=task_id,
                name=task.name,
                edge_type=EdgeType.DATA_SYNTHESIS.value,
                description=task.description,
                from_node_id=src_node.id,
                to_node_id=dst_dataset.id,
            )

            lineage_service = LineageService(db=db)
            await lineage_service.generate_graph(src_node, synthesis_edge, dest_node)
            await db.commit()

            logger.info(f"Added synthesis lineage: {src_node.name} -> {dest_dataset.name}")
        except Exception as exc:
            logger.error(f"Failed to add synthesis lineage: {exc}")