"""
恢复处理中断的索引任务
用于服务重启后继续处理pending/processing状态的任务
"""

import asyncio
import logging
import sys
from pathlib import Path

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import select
from app.core.database import async_session_maker
from app.models.entities import IndexJob, Document, IndexJobStatus, DocumentStatus
from app.api.admin.kb import _index_document

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


async def resume_pending_jobs():
    """恢复所有pending和processing状态的任务"""
    async with async_session_maker() as session:
        # 查询所有未完成的任务
        result = await session.execute(
            select(IndexJob).where(
                IndexJob.status.in_([IndexJobStatus.PENDING.value, IndexJobStatus.PROCESSING.value])
            )
        )
        pending_jobs = result.scalars().all()

        if not pending_jobs:
            logger.info("没有需要恢复的任务")
            return

        logger.info(f"发现 {len(pending_jobs)} 个未完成的任务")

        for job in pending_jobs:
            try:
                # 获取关联的文档
                doc_result = await session.execute(
                    select(Document).where(Document.id == job.doc_id)
                )
                doc = doc_result.scalar_one_or_none()

                if not doc:
                    logger.error(f"找不到文档: {job.doc_id}")
                    continue

                if not doc.file_path or not Path(doc.file_path).exists():
                    logger.error(f"文档文件不存在: {doc.file_path}")
                    # 标记为失败
                    job.status = IndexJobStatus.FAILED.value
                    job.error_msg = "文档文件不存在"
                    doc.status = DocumentStatus.FAILED.value
                    doc.error_msg = "文档文件不存在"
                    await session.commit()
                    continue

                logger.info(f"恢复处理: job_id={job.id}, doc_id={doc.id}, file={doc.file_name}")

                # 读取文件内容
                with open(doc.file_path, 'rb') as f:
                    file_content = f.read()

                # 重置任务状态为pending
                job.status = IndexJobStatus.PENDING.value
                job.progress = 0
                job.error_msg = None
                await session.commit()

                # 启动后台任务处理
                # 注意：这里我们直接调用，不使用background_tasks
                await process_job(
                    tenant_id=job.tenant_id,
                    kb_id=doc.kb_id,
                    job_id=str(job.id),
                    doc_id=str(doc.id),
                    file_content=file_content,
                    filename=doc.file_name,
                    metadata=doc.doc_metadata or {}
                )

                logger.info(f"任务处理完成: job_id={job.id}")

            except Exception as e:
                logger.error(f"处理任务失败: job_id={job.id}, error={e}")
                # 标记为失败
                job.status = IndexJobStatus.FAILED.value
                job.error_msg = str(e)
                if doc:
                    doc.status = DocumentStatus.FAILED.value
                    doc.error_msg = str(e)
                await session.commit()

        logger.info("所有任务处理完成")


async def process_job(tenant_id: str, kb_id: str, job_id: str, doc_id: str,
                      file_content: bytes, filename: str, metadata: dict):
    """
    处理单个索引任务
    复制自 _index_document 函数
    """
    import tempfile
    from pathlib import Path

    from qdrant_client.models import PointStruct

    from app.core.qdrant_client import get_qdrant_client
    from app.services.document import DocumentParseException, UnsupportedFormatError, parse_document
    from app.services.embedding import get_embedding_provider
    from app.services.kb import KBService
    from app.api.admin.kb import chunk_text_by_lines, TextChunk

    logger.info(f"[RESUME] Starting indexing: tenant={tenant_id}, kb_id={kb_id}, job_id={job_id}, doc_id={doc_id}")

    async with async_session_maker() as session:
        kb_service = KBService(session)
        try:
            await kb_service.update_job_status(
                tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=10
            )
            await session.commit()

            parse_result = None
            text = None
            file_ext = Path(filename or "").suffix.lower()
            logger.info(f"[RESUME] File extension: {file_ext}, content size: {len(file_content)} bytes")

            text_extensions = {".txt", ".md", ".markdown", ".rst", ".log", ".json", ".xml", ".yaml", ".yml"}

            if file_ext in text_extensions or not file_ext:
                logger.info("[RESUME] Treating as text file")
                text = None
                for encoding in ["utf-8", "gbk", "gb2312", "gb18030", "big5", "utf-16", "latin-1"]:
                    try:
                        text = file_content.decode(encoding)
                        logger.info(f"[RESUME] Successfully decoded with encoding: {encoding}")
                        break
                    except (UnicodeDecodeError, LookupError):
                        continue

                if text is None:
                    text = file_content.decode("utf-8", errors="replace")
            else:
                logger.info("[RESUME] Binary file detected, will parse with document parser")
                await kb_service.update_job_status(
                    tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=15
                )
                await session.commit()

                with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp_file:
                    tmp_file.write(file_content)
                    tmp_path = tmp_file.name

                logger.info(f"[RESUME] Temp file created: {tmp_path}")

                try:
                    logger.info(f"[RESUME] Starting document parsing for {file_ext}...")
                    parse_result = parse_document(tmp_path)
                    text = parse_result.text
                    logger.info(
                        f"[RESUME] Parsed document SUCCESS: {filename}, chars={len(text)}"
                    )
                except UnsupportedFormatError as e:
                    logger.error(f"[RESUME] UnsupportedFormatError: {e}")
                    text = file_content.decode("utf-8", errors="ignore")
                except DocumentParseException as e:
                    logger.error(f"[RESUME] DocumentParseException: {e}")
                    text = file_content.decode("utf-8", errors="ignore")
                except Exception as e:
                    logger.error(f"[RESUME] Unexpected parsing error: {type(e).__name__}: {e}")
                    text = file_content.decode("utf-8", errors="ignore")
                finally:
                    Path(tmp_path).unlink(missing_ok=True)

            logger.info(f"[RESUME] Final text length: {len(text)} chars")

            await kb_service.update_job_status(
                tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=20
            )
            await session.commit()

            logger.info("[RESUME] Getting embedding provider...")
            embedding_provider = await get_embedding_provider()
            logger.info(f"[RESUME] Embedding provider: {type(embedding_provider).__name__}")

            all_chunks: list[TextChunk] = []

            if parse_result and parse_result.pages:
                logger.info(f"[RESUME] PDF with {len(parse_result.pages)} pages")
                for page in parse_result.pages:
                    page_chunks = chunk_text_by_lines(
                        page.text,
                        min_line_length=10,
                        source=filename,
                    )
                    for pc in page_chunks:
                        pc.page = page.page
                    all_chunks.extend(page_chunks)
            else:
                logger.info("[RESUME] Using line-based chunking")
                all_chunks = chunk_text_by_lines(
                    text,
                    min_line_length=10,
                    source=filename,
                )

            logger.info(f"[RESUME] Total chunks: {len(all_chunks)}")

            qdrant = await get_qdrant_client()
            await qdrant.ensure_kb_collection_exists(tenant_id, kb_id, use_multi_vector=True)

            from app.services.embedding.nomic_provider import NomicEmbeddingProvider
            use_multi_vector = isinstance(embedding_provider, NomicEmbeddingProvider)
            logger.info(f"[RESUME] Using multi-vector format: {use_multi_vector}")

            import uuid
            points = []
            total_chunks = len(all_chunks)
            doc_metadata = metadata or {}

            for i, chunk in enumerate(all_chunks):
                payload = {
                    "text": chunk.text,
                    "source": doc_id,
                    "kb_id": kb_id,
                    "chunk_index": i,
                    "start_token": chunk.start_token,
                    "end_token": chunk.end_token,
                    "metadata": doc_metadata,
                }
                if chunk.page is not None:
                    payload["page"] = chunk.page
                if chunk.source:
                    payload["filename"] = chunk.source

                if use_multi_vector:
                    embedding_result = await embedding_provider.embed_document(chunk.text)
                    points.append({
                        "id": str(uuid.uuid4()),
                        "vector": {
                            "full": embedding_result.embedding_full,
                            "dim_256": embedding_result.embedding_256,
                            "dim_512": embedding_result.embedding_512,
                        },
                        "payload": payload,
                    })
                else:
                    embedding = await embedding_provider.embed(chunk.text)
                    points.append(
                        PointStruct(
                            id=str(uuid.uuid4()),
                            vector=embedding,
                            payload=payload,
                        )
                    )

                progress = 20 + int((i + 1) / total_chunks * 70)
                if i % 10 == 0 or i == total_chunks - 1:
                    await kb_service.update_job_status(
                        tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=progress
                    )
                    await session.commit()

            if points:
                logger.info(f"[RESUME] Upserting {len(points)} vectors to Qdrant...")
                if use_multi_vector:
                    await qdrant.upsert_multi_vector(tenant_id, points, kb_id=kb_id)
                else:
                    await qdrant.upsert_vectors(tenant_id, points, kb_id=kb_id)

            await kb_service.update_job_status(
                tenant_id, job_id, IndexJobStatus.COMPLETED.value, progress=100
            )
            await session.commit()

            logger.info(
                f"[RESUME] COMPLETED: tenant={tenant_id}, kb_id={kb_id}, "
                f"job_id={job_id}, chunks={len(all_chunks)}"
            )

        except Exception as e:
            import traceback
            logger.error(f"[RESUME] FAILED: {e}\n{traceback.format_exc()}")
            await session.rollback()
            async with async_session_maker() as error_session:
                kb_service = KBService(error_session)
                await kb_service.update_job_status(
                    tenant_id, job_id, IndexJobStatus.FAILED.value,
                    progress=0, error_msg=str(e)
                )
                await error_session.commit()


if __name__ == "__main__":
    logger.info("开始恢复索引任务...")
    asyncio.run(resume_pending_jobs())
    logger.info("恢复脚本执行完成")
