from __future__ import annotations

import hashlib
import uuid
from typing import Any

from .models import VALID_TASK_TYPES


class OrchestratorIngestionMixin:
    """Knowledge ingestion, memory chunk storage, and vector searching."""

    async def ingest_knowledge(
        self,
        task_type: str,
        chunks: list[dict[str, Any]],
        default_metadata: dict[str, Any] | None = None,
    ) -> int:
        """Embed and store knowledge chunks."""

        self._validate_task_type(task_type)
        texts = [str(chunk.get("text", "")).strip() for chunk in chunks]
        embeddings = self._embedding_service.embed_texts(texts)
        points: list[dict[str, Any]] = []
        defaults = dict(default_metadata or {})
        defaults["task_type"] = task_type

        for index, (chunk, vector) in enumerate(zip(chunks, embeddings), start=1):
            metadata = dict(defaults)
            metadata.update(dict(chunk.get("metadata", {}) or {}))
            metadata["task_type"] = task_type
            raw_id = metadata.get("id")
            if isinstance(raw_id, int):
                point_id = raw_id
            elif raw_id:
                point_id = str(
                    uuid.UUID(hashlib.md5(str(raw_id).encode("utf-8")).hexdigest())
                )
            else:
                point_id = str(uuid.uuid4())
            metadata["source_index"] = index
            points.append(
                {
                    "id": point_id,
                    "vector": vector,
                    "payload": {
                        "text": chunk.get("text", ""),
                        "metadata": metadata,
                    },
                }
            )
        await self._qdrant_store.upsert(points)
        return len(points)

    async def ingest_review_memory(
        self,
        review_id: str,
        task_type: str | None = None,
        default_metadata: dict[str, Any] | None = None,
    ) -> int:
        """Load one dossier review session from PostgreSQL and ingest memory chunks."""

        sessions = self._postgres_client.fetch_dossier_review_sessions(
            review_id=review_id
        )
        if not sessions:
            raise ValueError(
                f"Không tìm thấy dossier review session với review_id={review_id}."
            )

        session = sessions[0]
        chunks = self._review_memory_service.build_review_session_chunks(session)
        if not chunks:
            raise ValueError("Không tạo được review memory chunk phù hợp để ingest.")

        resolved_task_type = task_type or self._resolve_task_type_from_memory_chunks(
            chunks
        )
        if not resolved_task_type:
            raise ValueError("Không xác định được task_type phù hợp cho review memory.")

        merged_defaults = dict(default_metadata or {})
        merged_defaults.setdefault("project_code", session.get("project_code"))
        merged_defaults.setdefault("agent_scope", session.get("lead_agent_code"))
        merged_defaults.setdefault("review_id", session.get("review_id"))
        return await self.ingest_knowledge(
            task_type=resolved_task_type,
            chunks=chunks,
            default_metadata=merged_defaults,
        )

    async def ingest_project_memory(
        self,
        task_type: str,
        events: list[dict[str, Any]],
        default_metadata: dict[str, Any] | None = None,
    ) -> int:
        """Convert project events into knowledge chunks and ingest them."""

        self._validate_task_type(task_type)
        all_chunks: list[dict[str, Any]] = []

        for event in events:
            event_payload = dict(event)
            event_payload.setdefault("metadata", {})
            event_payload["metadata"].setdefault("task_type", task_type)
            all_chunks.extend(
                self._review_memory_service.build_project_event_chunks(event_payload)
            )

        if not all_chunks:
            raise ValueError("Không có project memory chunk hợp lệ để ingest.")

        return await self.ingest_knowledge(
            task_type=task_type,
            chunks=all_chunks,
            default_metadata=default_metadata,
        )

    async def search_knowledge(
        self,
        task_type: str,
        query: str,
        top_k: int = 3,
        retrieval_filters: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Search knowledge chunks for a task type."""

        self._validate_task_type(task_type)
        effective_filters = self._merge_retrieval_filters(
            task_type=task_type,
            agent_code=None,
            retrieval_filters=retrieval_filters,
        )
        return await self._retrieve(
            task_type=task_type,
            query=query,
            top_k=top_k,
            retrieval_filters=effective_filters,
        )

    def _resolve_task_type_from_memory_chunks(
        chunks: list[dict[str, Any]],
    ) -> str | None:
        """Resolve task type from generated memory chunk metadata."""
        for chunk in chunks:
            metadata = dict(chunk.get("metadata", {}) or {})
            task_type = metadata.get("task_type")
            if isinstance(task_type, str) and task_type in VALID_TASK_TYPES:
                return task_type
        return None
