"""Adaptive RAG mixin module."""
from __future__ import annotations

from .shared import *


class AdaptiveRAGAnswerPayloadMixin:
    @staticmethod
    def _safe_source_metadata(raw: Any) -> dict[str, Any]:
        return raw if isinstance(raw, dict) else {}

    @staticmethod
    def _source_ref_value(value: Any) -> Any:
        if value is None:
            return None
        if isinstance(value, (str, int, float, bool)):
            text = str(value).strip()
            return text or None
        if isinstance(value, list):
            rows = [str(item).strip() for item in value if str(item).strip()]
            return rows or None
        if isinstance(value, dict):
            try:
                return json.loads(json.dumps(value, ensure_ascii=False))
            except Exception:
                return dict(value)
        text = str(value).strip()
        return text or None

    @staticmethod
    def _normalize_answer_mode(value: Optional[str]) -> str:
        raw = str(value or "").strip().lower()
        if raw in {
            ANSWER_MODE_FAST_PATH,
            ANSWER_MODE_RETRIEVAL_NORMAL,
            ANSWER_MODE_RETRIEVAL_DEEP,
            ANSWER_MODE_CLARIFY,
            ANSWER_MODE_BLOCKED,
        }:
            return raw
        return ANSWER_MODE_RETRIEVAL_NORMAL

    def _normalize_grounding_status(
        self,
        value: Optional[str],
        *,
        sources: Optional[list[dict[str, Any]]] = None,
        missing_fields: Optional[list[str]] = None,
        answer_mode: Optional[str] = None,
        reason_code: Optional[str] = None,
    ) -> str:
        raw = str(value or "").strip().lower()
        if raw in {
            GROUNDING_STATUS_GROUNDED,
            GROUNDING_STATUS_PARTIALLY_GROUNDED,
            GROUNDING_STATUS_INSUFFICIENT_EVIDENCE,
            GROUNDING_STATUS_BLOCKED,
        }:
            return raw

        normalized_mode = self._normalize_answer_mode(answer_mode)
        if normalized_mode == ANSWER_MODE_BLOCKED:
            return GROUNDING_STATUS_BLOCKED
        if normalized_mode == ANSWER_MODE_CLARIFY:
            return GROUNDING_STATUS_INSUFFICIENT_EVIDENCE

        source_rows = sources if isinstance(sources, list) else []
        snapshot_only = False
        if hasattr(self, "_is_snapshot_only_sources"):
            try:
                snapshot_only = bool(self._is_snapshot_only_sources(source_rows))
            except Exception:
                snapshot_only = False
        if not snapshot_only and source_rows:
            snapshot_only = all(str(item.get("source") or "").strip().lower() == "system" for item in source_rows if isinstance(item, dict))

        missing = [str(item).strip() for item in (missing_fields or []) if str(item).strip()]
        normalized_reason = str(reason_code or "").strip().lower()
        if normalized_reason == "insufficient_evidence":
            return GROUNDING_STATUS_INSUFFICIENT_EVIDENCE
        if snapshot_only or missing:
            return GROUNDING_STATUS_PARTIALLY_GROUNDED
        if source_rows:
            return GROUNDING_STATUS_GROUNDED
        return GROUNDING_STATUS_PARTIALLY_GROUNDED

    def _build_citations_from_sources(self, sources: list[dict[str, Any]]) -> list[dict[str, Any]]:
        normalized: list[dict[str, Any]] = []
        seen: set[tuple[str, str, str, str]] = set()
        for item in (sources or [])[:8]:
            if not isinstance(item, dict):
                continue
            metadata = self._safe_source_metadata(item.get("metadata"))
            source_type = self._source_ref_value(
                metadata.get("source_type")
                or item.get("source")
                or metadata.get("table")
                or metadata.get("mode")
            )
            source_label = self._source_ref_value(item.get("label") or metadata.get("source_label"))
            source_id = self._source_ref_value(
                metadata.get("source_id")
                or metadata.get("doc_id")
                or metadata.get("knowledge_doc_id")
                or metadata.get("message_id")
                or metadata.get("url")
                or metadata.get("table")
                or metadata.get("matched_project_ids")
            )
            chunk_id = self._source_ref_value(metadata.get("chunk_id") or metadata.get("chunk_index"))
            target_url = self._source_ref_value(
                metadata.get("target_url")
                or metadata.get("doc_url")
                or metadata.get("url")
                or metadata.get("permalink")
            )
            project_id = self._source_ref_value(metadata.get("project_id"))
            updated_at = self._source_ref_value(
                metadata.get("updated_at")
                or metadata.get("doc_updated_at")
                or metadata.get("source_updated_at")
            )
            citation = {
                "source_type": source_type,
                "source_label": source_label,
                "source_id": source_id,
                "chunk_id": chunk_id,
                "target_url": target_url,
                "project_id": project_id,
                "updated_at": updated_at,
                "confidence": round(self._to_float(item.get("score"), 0.0), 4),
            }
            key = (
                str(source_type or ""),
                str(source_id or ""),
                str(chunk_id or ""),
                str(target_url or ""),
            )
            if key in seen:
                continue
            seen.add(key)
            normalized.append(citation)
        return normalized

    @staticmethod
    def _build_sources_used_from_citations(citations: list[dict[str, Any]]) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for item in citations:
            if not isinstance(item, dict):
                continue
            rows.append(
                {
                    "source_type": item.get("source_type"),
                    "source_label": item.get("source_label"),
                    "source_id": item.get("source_id"),
                    "chunk_id": item.get("chunk_id"),
                    "target_url": item.get("target_url"),
                    "project_id": item.get("project_id"),
                    "updated_at": item.get("updated_at"),
                }
            )
        return rows

    def _build_evidence_pack(
        self,
        *,
        answer: str,
        citations: list[dict[str, Any]],
        tool_trace: list[dict[str, Any]],
        missing_fields: Optional[list[str]] = None,
        grounding_status: Optional[str] = None,
        reason_code: Optional[str] = None,
    ) -> dict[str, Any]:
        evidence_labels: list[str] = []
        for item in citations[:5]:
            if not isinstance(item, dict):
                continue
            label = str(item.get("source_label") or item.get("source_id") or item.get("source_type") or "").strip()
            if label:
                evidence_labels.append(label)
        uncertainties = [str(item).strip() for item in (missing_fields or []) if str(item).strip()]
        normalized_grounding = str(grounding_status or "").strip().lower()
        normalized_reason = str(reason_code or "").strip()
        if normalized_grounding in {GROUNDING_STATUS_INSUFFICIENT_EVIDENCE, GROUNDING_STATUS_BLOCKED} and normalized_reason:
            if normalized_reason not in uncertainties:
                uncertainties.append(normalized_reason)
        return {
            "summary": {
                "conclusion": str(answer or "").strip(),
                "evidence": evidence_labels,
                "uncertainties": uncertainties,
            },
            "tool_trace": tool_trace if isinstance(tool_trace, list) else [],
        }

    def _build_response_payload(
        self,
        *,
        answer: str,
        sources: list[dict[str, Any]],
        domains_used: list[str],
        response_time_ms: int,
        db_query_ms: int,
        vector_search_ms: int,
        graph_query_ms: int = 0,
        db_timeout: bool = False,
        vector_timeout: bool = False,
        graph_timeout: bool = False,
        qa_id: Optional[int] = None,
        session_id: Optional[int] = None,
        turn_index: int = 1,
        relevance_max: float = 0.0,
        query_type: str = QUERY_TYPE_UNKNOWN,
        confidence: float = 0.0,
        missing_fields: Optional[list[str]] = None,
        tool_trace: Optional[list[dict[str, Any]]] = None,
        thinking_mode: str = THINKING_MODE_NORMAL,
        action_taken: Optional[str] = None,
        signal_id: Optional[int] = None,
        reply_required: bool = False,
        reply_text: Optional[str] = None,
        ack_latency_ms: Optional[int] = None,
        trace_id: Optional[str] = None,
        processing_stage: Optional[str] = None,
        clarification_code: Optional[str] = None,
        clarification_candidates: Optional[list[dict[str, Any]]] = None,
        anchor_mode: Optional[str] = None,
        anchor_kind: Optional[str] = None,
        meta_followup_kind: Optional[str] = None,
        reason_code: Optional[str] = None,
        answer_mode: Optional[str] = None,
        grounding_status: Optional[str] = None,
        policy_version: Optional[str] = None,
        citations: Optional[list[dict[str, Any]]] = None,
        sources_used: Optional[list[dict[str, Any]]] = None,
        evidence_pack: Optional[dict[str, Any]] = None,
        prompt_variant_id: Optional[str] = None,
    ) -> dict[str, Any]:
        evidence: list[dict[str, Any]] = []
        for item in sources[:8]:
            if not isinstance(item, dict):
                continue
            evidence.append(
                {
                    "source": str(item.get("source") or "unknown"),
                    "label": str(item.get("label") or "unknown"),
                    "score": round(self._to_float(item.get("score"), 0.0), 4),
                    "metadata": item.get("metadata") if isinstance(item.get("metadata"), dict) else {},
                }
            )
        normalized_tool_trace = tool_trace if isinstance(tool_trace, list) else []
        normalized_citations = citations if isinstance(citations, list) else self._build_citations_from_sources(sources)
        normalized_sources_used = (
            sources_used
            if isinstance(sources_used, list)
            else self._build_sources_used_from_citations(normalized_citations)
        )
        normalized_answer_mode = self._normalize_answer_mode(answer_mode)
        normalized_grounding_status = self._normalize_grounding_status(
            grounding_status,
            sources=sources,
            missing_fields=missing_fields,
            answer_mode=normalized_answer_mode,
            reason_code=reason_code,
        )
        normalized_policy_version = str(
            policy_version
            or self._build_policy_version(prompt_variant_id=prompt_variant_id)
        ).strip() or self._build_policy_version(prompt_variant_id=prompt_variant_id)
        normalized_evidence_pack = (
            evidence_pack
            if isinstance(evidence_pack, dict)
            else self._build_evidence_pack(
                answer=answer,
                citations=normalized_citations,
                tool_trace=normalized_tool_trace,
                missing_fields=missing_fields,
                grounding_status=normalized_grounding_status,
                reason_code=reason_code,
            )
        )
        return {
            "answer": str(answer or "").strip(),
            "evidence": evidence,
            "confidence": round(max(0.0, min(1.0, self._to_float(confidence, 0.0))), 4),
            "thinking_mode": self._normalize_thinking_mode(thinking_mode),
            "missing_fields": [str(x) for x in (missing_fields or []) if str(x).strip()],
            "query_type": str(query_type or QUERY_TYPE_UNKNOWN),
            "sources": sources,
            "domains_used": domains_used,
            "response_time_ms": int(response_time_ms or 0),
            "retrieval_metrics": {
                "db_query_ms": int(db_query_ms or 0),
                "vector_search_ms": int(vector_search_ms or 0),
                "graph_query_ms": int(graph_query_ms or 0),
                "db_timeout": bool(db_timeout),
                "vector_timeout": bool(vector_timeout),
                "graph_timeout": bool(graph_timeout),
            },
            "tool_trace": normalized_tool_trace,
            "qa_id": int(qa_id) if qa_id is not None else None,
            "session_id": int(session_id) if session_id is not None else None,
            "turn_index": int(turn_index or 1),
            "relevance_max": round(max(0.0, self._to_float(relevance_max, 0.0)), 4),
            "action_taken": str(action_taken).strip() if action_taken else None,
            "signal_id": int(signal_id) if signal_id is not None else None,
            "reply_required": bool(reply_required),
            "reply_text": str(reply_text or "").strip() if reply_text else None,
            "ack_latency_ms": int(ack_latency_ms) if ack_latency_ms is not None else None,
            "trace_id": str(trace_id or "").strip() or None,
            "processing_stage": str(processing_stage or "").strip() or None,
            "clarification_code": str(clarification_code or "").strip() or None,
            "clarification_candidates": (
                clarification_candidates[:3]
                if isinstance(clarification_candidates, list)
                else []
            ),
            "anchor_mode": str(anchor_mode or "").strip() or None,
            "anchor_kind": str(anchor_kind or "").strip() or None,
            "meta_followup_kind": str(meta_followup_kind or "").strip() or None,
            "reason_code": str(reason_code or "").strip() or None,
            "answer_mode": normalized_answer_mode,
            "grounding_status": normalized_grounding_status,
            "policy_version": normalized_policy_version,
            "citations": normalized_citations,
            "sources_used": normalized_sources_used,
            "evidence_pack": normalized_evidence_pack,
        }

    @staticmethod
    def _extract_ack_metadata_from_sources(
        sources: Any,
        *,
        answer: str = "",
    ) -> dict[str, Any]:
        rows = sources if isinstance(sources, list) else []
        action_taken = ""
        signal_id: Optional[int] = None
        reply_required = False
        ack_latency_ms: Optional[int] = None
        trace_id = ""
        processing_stage = ""
        for source in rows:
            if not isinstance(source, dict):
                continue
            meta = source.get("metadata") if isinstance(source.get("metadata"), dict) else {}
            candidate_action = str(meta.get("action_taken") or "").strip()
            if not candidate_action:
                continue
            action_taken = candidate_action
            try:
                signal_id = int(meta.get("signal_id")) if meta.get("signal_id") is not None else None
            except Exception:
                signal_id = None
            trace_id = str(meta.get("trace_id") or "").strip()
            processing_stage = str(meta.get("processing_stage") or meta.get("action_mode") or "").strip()
            try:
                ack_latency_ms = (
                    int(meta.get("ack_latency_ms")) if meta.get("ack_latency_ms") is not None else None
                )
            except Exception:
                ack_latency_ms = None
            reply_required = True
            break
        reply_text = str(answer or "").strip() if reply_required else ""
        normalized_stage = processing_stage
        if normalized_stage == "write":
            normalized_stage = "ack_only"
        return {
            "action_taken": action_taken or None,
            "signal_id": signal_id,
            "reply_required": bool(reply_required),
            "reply_text": reply_text or None,
            "ack_latency_ms": ack_latency_ms,
            "trace_id": trace_id or None,
            "processing_stage": normalized_stage or None,
        }