from __future__ import annotations

import asyncio
import hashlib
import logging
import os
from collections.abc import Awaitable, Callable
from inspect import isawaitable
from typing import TYPE_CHECKING

from app.core.rag.contracts import RagDocument
from app.core.rag.indexing.common.path_filter import filter_snapshot_files
from app.core.rag.indexing.code.pipeline import CodeIndexingPipeline
from app.core.rag.indexing.common.report import IndexReport
from app.core.rag.indexing.docs.pipeline import DocsIndexingPipeline
from app.core.rag.embedding.gigachat_embedder import GigaChatEmbedder

LOGGER = logging.getLogger(__name__)

if TYPE_CHECKING:
    from app.core.rag.persistence.repository import RagRepository


class RagService:
    def __init__(
        self,
        embedder: GigaChatEmbedder,
        repository: "RagRepository",
        chunker=None,
    ) -> None:
        self._embedder = embedder
        self._repo = repository
        self._docs = DocsIndexingPipeline()
        self._code = CodeIndexingPipeline()
        self._cache_enabled = os.getenv("RAG_DOCUMENT_CACHE_ENABLED", "false").strip().lower() in {"1", "true", "yes", "on"}

    async def index_snapshot(
        self,
        rag_session_id: str,
        files: list[dict],
        progress_cb: Callable[[int, int, str], Awaitable[None] | None] | None = None,
    ) -> tuple[int, int, int, int]:
        LOGGER.warning("rag index snapshot started: rag_session_id=%s files=%s", rag_session_id, len(files))
        report = await self._index_files(rag_session_id, files, progress_cb=progress_cb)
        self._repo.replace_documents(rag_session_id, report.documents_list)
        LOGGER.warning(
            "rag index snapshot persisted: rag_session_id=%s indexed=%s failed=%s docs=%s",
            rag_session_id,
            report.indexed_files,
            report.failed_files,
            len(report.documents_list),
        )
        return report.as_tuple()

    async def index_changes(
        self,
        rag_session_id: str,
        changed_files: list[dict],
        progress_cb: Callable[[int, int, str], Awaitable[None] | None] | None = None,
    ) -> tuple[int, int, int, int]:
        LOGGER.warning("rag index changes started: rag_session_id=%s changes=%s", rag_session_id, len(changed_files))
        delete_paths: list[str] = []
        upserts: list[dict] = []
        for item in changed_files:
            if str(item.get("op")) == "delete":
                delete_paths.append(str(item.get("path", "")))
            else:
                upserts.append(item)
        report = await self._index_files(rag_session_id, upserts, progress_cb=progress_cb)
        self._repo.apply_document_changes(rag_session_id, delete_paths, report.documents_list)
        LOGGER.warning(
            "rag index changes persisted: rag_session_id=%s indexed=%s failed=%s docs=%s delete_paths=%s",
            rag_session_id,
            report.indexed_files,
            report.failed_files,
            len(report.documents_list),
            len(delete_paths),
        )
        return report.as_tuple()

    async def _index_files(
        self,
        rag_session_id: str,
        files: list[dict],
        progress_cb: Callable[[int, int, str], Awaitable[None] | None] | None = None,
    ) -> "_PipelineReport":
        indexable_files = filter_snapshot_files(files)
        skipped_files = len(files) - len(indexable_files)
        total_files = len(indexable_files)
        report = _PipelineReport()
        repo_id = self._resolve_repo_id(rag_session_id)
        if skipped_files > 0:
            LOGGER.warning(
                "rag ingest filter: rag_session_id=%s skipped_files=%s",
                rag_session_id,
                skipped_files,
            )
        for index, file in enumerate(indexable_files, start=1):
            path = str(file.get("path", ""))
            try:
                LOGGER.warning(
                    "rag index file started: rag_session_id=%s file=%s/%s path=%s",
                    rag_session_id,
                    index,
                    total_files,
                    path,
                )
                blob_sha = self._blob_sha(file)
                pipelines = self._resolve_pipeline_names(path)
                cached = []
                if self._cache_enabled:
                    cached = await asyncio.to_thread(self._repo.get_cached_documents, repo_id, blob_sha)
                if cached:
                    self._report_missing_or_partial_docs(path, cached)
                    report.documents_list.extend(self._with_file_metadata(cached, file, repo_id, blob_sha))
                    report.cache_hit_files += 1
                    LOGGER.warning(
                        "rag ingest file: rag_session_id=%s path=%s processing=cache pipeline=%s",
                        rag_session_id,
                        path,
                        ",".join(pipelines),
                    )
                else:
                    built = self._build_documents(repo_id, path, file)
                    LOGGER.warning(
                        "rag index file built docs: rag_session_id=%s path=%s docs=%s",
                        rag_session_id,
                        path,
                        len(built),
                    )
                    self._report_missing_or_partial_docs(path, built)
                    embed_timeout_sec = max(1, int(os.getenv("RAG_EMBED_FILE_TIMEOUT_SEC", "8")))
                    LOGGER.warning(
                        "rag index file embedding started: rag_session_id=%s path=%s timeout_sec=%s",
                        rag_session_id,
                        path,
                        embed_timeout_sec,
                    )
                    embedded = await asyncio.wait_for(
                        asyncio.to_thread(self._embed_documents, built, file, repo_id, blob_sha),
                        timeout=embed_timeout_sec,
                    )
                    LOGGER.warning(
                        "rag index file embedded docs: rag_session_id=%s path=%s docs=%s",
                        rag_session_id,
                        path,
                        len(embedded),
                    )
                    report.documents_list.extend(embedded)
                    if self._cache_enabled:
                        await asyncio.to_thread(self._repo.cache_documents, repo_id, path, blob_sha, embedded)
                    report.cache_miss_files += 1
                    LOGGER.warning(
                        "rag ingest file: rag_session_id=%s path=%s processing=embed pipeline=%s",
                        rag_session_id,
                        path,
                        ",".join(pipelines),
                    )
                report.indexed_files += 1
                LOGGER.warning(
                    "rag index file completed: rag_session_id=%s file=%s/%s path=%s",
                    rag_session_id,
                    index,
                    total_files,
                    path,
                )
            except Exception as exc:
                report.failed_files += 1
                report.warnings.append(f"{path}: {exc}")
                LOGGER.error(
                    "rag ingest failed file: rag_session_id=%s path=%s error=%s",
                    rag_session_id,
                    path,
                    exc,
                )
                # Fail-fast: stop indexing immediately so caller can expose the exact error to plugin.
                raise RuntimeError(f"RAG indexing failed for '{path}': {exc}") from exc
            await self._notify_progress(progress_cb, index, total_files, path)
        report.documents = len(report.documents_list)
        return report

    def _build_documents(self, repo_id: str, path: str, file: dict) -> list[RagDocument]:
        content = str(file.get("content") or "")
        commit_sha = file.get("commit_sha")
        docs: list[RagDocument] = []
        if self._docs.supports(path):
            docs.extend(self._docs.index_file(repo_id=repo_id, commit_sha=commit_sha, path=path, content=content))
        if self._code.supports(path):
            docs.extend(self._code.index_file(repo_id=repo_id, commit_sha=commit_sha, path=path, content=content))
        if not docs:
            docs.extend(self._docs.index_file(repo_id=repo_id, commit_sha=commit_sha, path=path, content=content))
        return docs

    def _resolve_pipeline_names(self, path: str) -> list[str]:
        names: list[str] = []
        if self._docs.supports(path):
            names.append("DOCS")
        if self._code.supports(path):
            names.append("CODE")
        if not names:
            names.append("DOCS")
        return names

    def _report_missing_or_partial_docs(self, path: str, docs: list[RagDocument]) -> None:
        if not docs:
            LOGGER.error("rag ingest file not indexed: path=%s reason=no_documents_built", path)
            return
        if not str(path).startswith("docs/"):
            return
        has_d1 = any(str(doc.layer or "") == "D1_DOCUMENT_CATALOG" for doc in docs)
        if not has_d1:
            LOGGER.error("rag ingest file partially indexed: path=%s reason=missing_d1_document_catalog", path)

    def _embed_documents(self, docs: list[RagDocument], file: dict, repo_id: str, blob_sha: str) -> list[RagDocument]:
        if not docs:
            return []
        batch_size = max(1, int(os.getenv("RAG_EMBED_BATCH_SIZE", "16")))
        request_timeout_sec = max(1, int(os.getenv("RAG_EMBED_REQUEST_TIMEOUT_SEC", "5")))
        request_retries = max(1, int(os.getenv("RAG_EMBED_REQUEST_MAX_RETRIES", "1")))
        max_chars = max(200, int(os.getenv("RAG_EMBED_MAX_CHARS", "1200")))
        overlap_chars = max(0, int(os.getenv("RAG_EMBED_OVERLAP_CHARS", "120")))
        prepared = self._prepare_docs_for_embedding(docs, max_chars=max_chars, overlap_chars=overlap_chars)
        if len(prepared) != len(docs):
            LOGGER.warning(
                "rag embed doc split: path=%s original_docs=%s prepared_docs=%s max_chars=%s overlap_chars=%s",
                file.get("path", ""),
                len(docs),
                len(prepared),
                max_chars,
                overlap_chars,
            )
        metadata = self._document_metadata(file, repo_id, blob_sha)
        for doc in prepared:
            doc.metadata.update(metadata)
        for start in range(0, len(prepared), batch_size):
            batch = prepared[start : start + batch_size]
            LOGGER.warning(
                "rag embed batch start: path=%s batch_start=%s batch_size=%s timeout_sec=%s retries=%s",
                file.get("path", ""),
                start,
                len(batch),
                request_timeout_sec,
                request_retries,
            )
            vectors = self._embedder.embed(
                [doc.text for doc in batch],
                timeout_sec=request_timeout_sec,
                max_retries=request_retries,
            )
            LOGGER.warning(
                "rag embed batch done: path=%s batch_start=%s vectors=%s",
                file.get("path", ""),
                start,
                len(vectors),
            )
            for doc, vector in zip(batch, vectors):
                doc.embedding = vector
        return prepared

    def _prepare_docs_for_embedding(
        self,
        docs: list[RagDocument],
        *,
        max_chars: int,
        overlap_chars: int,
    ) -> list[RagDocument]:
        prepared: list[RagDocument] = []
        for doc in docs:
            text = str(doc.text or "")
            if len(text) <= max_chars:
                prepared.append(doc)
                continue
            parts = self._split_text_for_embedding(text, max_chars=max_chars, overlap_chars=overlap_chars)
            if len(parts) <= 1:
                prepared.append(doc)
                continue
            for idx, part in enumerate(parts, start=1):
                metadata = dict(doc.metadata)
                metadata["embed_part_index"] = idx
                metadata["embed_part_total"] = len(parts)
                if doc.doc_id:
                    metadata["embed_original_doc_id"] = doc.doc_id
                prepared.append(
                    RagDocument(
                        layer=doc.layer,
                        source=doc.source,
                        title=f"{doc.title} [part {idx}/{len(parts)}]",
                        text=part,
                        metadata=metadata,
                        links=list(doc.links),
                        span=doc.span,
                        doc_id=f"{doc.doc_id}:part:{idx}" if doc.doc_id else None,
                        lang=doc.lang,
                    )
                )
        return prepared

    def _split_text_for_embedding(self, text: str, *, max_chars: int, overlap_chars: int) -> list[str]:
        normalized = text.replace("\r\n", "\n").replace("\r", "\n")
        parts: list[str] = []
        start = 0
        length = len(normalized)
        while start < length:
            hard_end = min(length, start + max_chars)
            end = hard_end
            if hard_end < length:
                soft_start = min(length, start + max_chars // 2)
                pivot = normalized.rfind("\n", soft_start, hard_end)
                if pivot > start:
                    end = pivot + 1
            chunk = normalized[start:end].strip()
            if chunk:
                parts.append(chunk)
            if end >= length:
                break
            next_start = max(start + 1, end - overlap_chars)
            start = next_start
        return parts or [normalized]

    def _with_file_metadata(self, docs: list[RagDocument], file: dict, repo_id: str, blob_sha: str) -> list[RagDocument]:
        metadata = self._document_metadata(file, repo_id, blob_sha)
        for doc in docs:
            doc.metadata.update(metadata)
            doc.source.repo_id = repo_id
            doc.source.path = str(file.get("path", doc.source.path))
        return docs

    def _document_metadata(self, file: dict, repo_id: str, blob_sha: str) -> dict:
        metadata = {
            "blob_sha": blob_sha,
            "repo_id": repo_id,
            "artifact_type": file.get("artifact_type"),
            "section": file.get("section"),
            "doc_id": file.get("doc_id"),
            "doc_version": file.get("doc_version"),
            "owner": file.get("owner"),
            "system_component": file.get("system_component"),
            "last_modified": file.get("last_modified"),
            "staleness_score": file.get("staleness_score"),
        }
        return {key: value for key, value in metadata.items() if value not in (None, "", [], {})}

    def _resolve_repo_id(self, rag_session_id: str) -> str:
        session = self._repo.get_session(rag_session_id)
        if not session:
            return rag_session_id
        return str(session.get("project_id") or rag_session_id)

    def _blob_sha(self, file: dict) -> str:
        raw = str(file.get("content_hash") or "").strip()
        if raw:
            return raw
        content = str(file.get("content") or "")
        return hashlib.sha256(content.encode("utf-8")).hexdigest()

    async def _notify_progress(
        self,
        progress_cb: Callable[[int, int, str], Awaitable[None] | None] | None,
        current_file_index: int,
        total_files: int,
        current_file_name: str,
    ) -> None:
        if not progress_cb:
            return
        result = progress_cb(current_file_index, total_files, current_file_name)
        if isawaitable(result):
            await result

class _PipelineReport(IndexReport):
    def __init__(self) -> None:
        super().__init__()
        self.documents_list: list[RagDocument] = []
