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

from .shared import *


class AdaptiveRAGSessionStateMixin:
    @staticmethod
    def _normalize_state_scope(payload: Optional[dict[str, Any]]) -> dict[str, Any]:
        row = payload if isinstance(payload, dict) else {}
        project_id = row.get("project_id")
        try:
            normalized_project_id = int(project_id) if project_id is not None else None
        except Exception:
            normalized_project_id = None
        target_ref = str(row.get("target_ref") or "").strip() or None
        target_entity_type = str(row.get("target_entity_type") or "").strip().lower() or None
        return {
            "project_id": normalized_project_id,
            "target_ref": target_ref,
            "target_entity_type": target_entity_type,
        }

    def _state_scope_from_qa_record(self, qa_record: Optional[QARecord]) -> dict[str, Any]:
        if qa_record is None:
            return self._normalize_state_scope(None)
        snapshot = qa_record.routing_snapshot if isinstance(qa_record.routing_snapshot, dict) else {}
        query_semantics = snapshot.get("query_semantics") if isinstance(snapshot.get("query_semantics"), dict) else {}
        response_payload = snapshot.get("response_payload") if isinstance(snapshot.get("response_payload"), dict) else {}
        target_resolution = (
            query_semantics.get("target_resolution")
            if isinstance(query_semantics.get("target_resolution"), dict)
            else {}
        )
        analysis_packet = (
            response_payload.get("analysis_packet")
            if isinstance(response_payload.get("analysis_packet"), dict)
            else {}
        )
        fallback_target_ref = ""
        fallback_target_entity_type = ""
        for row in list(analysis_packet.get("focus_objects") or []):
            if not isinstance(row, dict):
                continue
            fallback_target_ref = str(row.get("object_ref") or row.get("target_ref") or "").strip()
            fallback_target_entity_type = str(row.get("object_type") or row.get("target_entity_type") or "").strip().lower()
            if fallback_target_ref:
                break

        project_id = qa_record.project_id
        target_ref = str(target_resolution.get("target_ref") or "").strip() or fallback_target_ref or None
        if project_id is None and target_ref and str(target_ref).startswith("project:"):
            try:
                project_id = int(str(target_ref).split(":", 1)[1])
            except Exception:
                project_id = None
        return self._normalize_state_scope(
            {
                "project_id": project_id,
                "target_ref": target_ref,
                "target_entity_type": (
                    str(target_resolution.get("target_entity_type") or "").strip().lower()
                    or fallback_target_entity_type
                    or None
                ),
            }
        )

    async def _load_last_turn_scope(self, *, session_id: Optional[int]) -> dict[str, Any]:
        if session_id is None:
            return self._normalize_state_scope(None)
        row = (
            await self.db.execute(
                select(QARecord)
                .where(
                    and_(
                        QARecord.tenant_id == self.tenant_id,
                        QARecord.session_id == int(session_id),
                    )
                )
                .order_by(QARecord.turn_index.desc(), QARecord.id.desc())
                .limit(1)
            )
        ).scalar_one_or_none()
        return self._state_scope_from_qa_record(row)

    async def _resolve_current_turn_scope(
        self,
        *,
        question: str,
        project_id: Optional[int],
        source_policy_hint: Optional[str],
        query_semantics: Optional[dict[str, Any]],
    ) -> dict[str, Any]:
        semantics = dict(query_semantics or {}) if isinstance(query_semantics, dict) else {}
        if not semantics:
            resolver = getattr(self, "_resolve_query_semantics", None)
            if callable(resolver):
                try:
                    semantics = await resolver(
                        message=question,
                        source_system="adaptive_rag_state_rewrite_gate",
                        source_policy_hint=source_policy_hint,
                    )
                except Exception:
                    semantics = {}

        scoped_project_id = project_id
        scope_resolver = getattr(self, "_resolve_project_scope_for_question", None)
        if callable(scope_resolver):
            try:
                scoped_project_id, _, _ = await scope_resolver(
                    question=question,
                    project_id=project_id,
                    query_semantics=semantics,
                    internal_context_packet=None,
                )
            except TypeError:
                scoped_project_id, _, _ = await scope_resolver(
                    question=question,
                    project_id=project_id,
                    query_semantics=semantics,
                )
            except Exception:
                scoped_project_id = project_id

        target_resolution = (
            semantics.get("target_resolution")
            if isinstance(semantics.get("target_resolution"), dict)
            else {}
        )
        target_ref = str(target_resolution.get("target_ref") or "").strip() or None
        if scoped_project_id is None and target_ref and target_ref.startswith("project:"):
            try:
                scoped_project_id = int(target_ref.split(":", 1)[1])
            except Exception:
                scoped_project_id = None
        return self._normalize_state_scope(
            {
                "project_id": scoped_project_id,
                "target_ref": target_ref,
                "target_entity_type": str(target_resolution.get("target_entity_type") or "").strip().lower() or None,
            }
        )

    async def _build_state_rewrite_governance(
        self,
        *,
        question: str,
        project_id: Optional[int],
        session_id: Optional[int],
        source_policy_hint: Optional[str],
        query_semantics: Optional[dict[str, Any]],
        session_state: Optional[QASessionState],
        anchor_resolution: Optional[FollowupAnchorResolution],
    ) -> dict[str, Any]:
        resolution = anchor_resolution or FollowupAnchorResolution()
        followup_mode = str(resolution.followup_mode or "").strip().lower()
        anchor_mode = str(resolution.anchor_mode or "").strip().lower()
        governance: dict[str, Any] = {
            "allow_state_rewrite": True,
            "reason": "followup_scope_continuity",
            "previous_scope": self._normalize_state_scope(None),
            "current_scope": self._normalize_state_scope(None),
        }
        if followup_mode != "referential":
            governance["reason"] = "non_referential_turn"
            return governance
        if anchor_mode == "explicit":
            governance["reason"] = "explicit_anchor"
            return governance
        if session_state is None or session_id is None:
            governance["reason"] = "missing_session_state"
            return governance

        previous_scope = await self._load_last_turn_scope(session_id=session_id)
        current_scope = await self._resolve_current_turn_scope(
            question=question,
            project_id=project_id,
            source_policy_hint=source_policy_hint,
            query_semantics=query_semantics,
        )
        governance["previous_scope"] = previous_scope
        governance["current_scope"] = current_scope

        previous_project_id = previous_scope.get("project_id")
        current_project_id = current_scope.get("project_id")
        previous_target_ref = str(previous_scope.get("target_ref") or "").strip() or None
        current_target_ref = str(current_scope.get("target_ref") or "").strip() or None

        if current_project_id is not None and previous_project_id is not None and current_project_id != previous_project_id:
            governance["allow_state_rewrite"] = False
            governance["reason"] = "new_project_scope"
            return governance
        if current_project_id is not None and previous_project_id is None:
            governance["allow_state_rewrite"] = False
            governance["reason"] = "new_project_scope"
            return governance
        if current_project_id is not None and previous_project_id == current_project_id and current_target_ref and previous_target_ref and current_target_ref != previous_target_ref:
            governance["allow_state_rewrite"] = False
            governance["reason"] = "new_business_scope"
            return governance
        return governance

    async def _ensure_session_state(self, *, session_id: Optional[int]) -> Optional[QASessionState]:
        if session_id is None:
            return None
        row = (
            await self.db.execute(
                select(QASessionState).where(
                    and_(
                        QASessionState.tenant_id == self.tenant_id,
                        QASessionState.session_id == int(session_id),
                    )
                )
            )
        ).scalar_one_or_none()
        if row is None:
            row = QASessionState(
                tenant_id=self.tenant_id,
                session_id=int(session_id),
                last_subject=None,
                last_entity_ids=[],
                last_intent=None,
                ambiguity="unknown",
                rewrite_history=[],
                pending_clarification=None,
                pending_clarification_expires_at=None,
            )
            self.db.add(row)
            await self.db.flush()
        return row

    def _rewrite_question_with_state(
        self,
        *,
        question: str,
        conversation: list[dict[str, str]],
        session_state: Optional[QASessionState],
        anchor_resolution: Optional[FollowupAnchorResolution] = None,
        rewrite_governance: Optional[dict[str, Any]] = None,
    ) -> tuple[str, dict[str, Any]]:
        _ = conversation
        current = str(question or "").strip()
        if not current:
            return current, {"mode": "empty"}

        resolution = anchor_resolution or FollowupAnchorResolution()
        if str(resolution.followup_mode or "").strip().lower() != "referential":
            return current, {"mode": "direct"}
        governance = rewrite_governance if isinstance(rewrite_governance, dict) else {}
        if governance and governance.get("allow_state_rewrite") is False:
            return current, {
                "mode": "scope_reset",
                "reason": str(governance.get("reason") or "scope_changed").strip() or "scope_changed",
                "current_scope": self._normalize_state_scope(governance.get("current_scope")),
                "previous_scope": self._normalize_state_scope(governance.get("previous_scope")),
            }

        explicit_anchor_mode = str(resolution.anchor_mode or "").strip().lower()
        if explicit_anchor_mode == "explicit":
            anchored_question, anchor_meta = self._rewrite_question_with_followup_anchor(
                question=current,
                anchor_resolution=resolution,
            )
            if str(anchor_meta.get("mode") or "").strip() not in {"", "direct"}:
                return anchored_question, anchor_meta

        state_subject = ""
        state_entity_ids: list[int] = []
        if session_state is not None:
            state_subject = str(session_state.last_subject or "").strip()
            if isinstance(session_state.last_entity_ids, list):
                for value in session_state.last_entity_ids:
                    try:
                        state_entity_ids.append(int(value))
                    except Exception:
                        continue
        if state_subject:
            rewritten = f"围绕“{state_subject}”回答：{current}"
            return rewritten[:1200], {"mode": "state_subject", "subject": state_subject}
        if state_entity_ids:
            anchor = ",".join(str(x) for x in state_entity_ids[:6])
            rewritten = f"围绕会话实体ID[{anchor}]回答：{current}"
            return rewritten[:1200], {"mode": "state_entities", "entity_ids": state_entity_ids[:6]}
        anchored_question, anchor_meta = self._rewrite_question_with_followup_anchor(
            question=current,
            anchor_resolution=resolution,
        )
        if str(anchor_meta.get("mode") or "").strip() not in {"", "direct"}:
            return anchored_question, anchor_meta
        return current, {"mode": "referential_without_subject"}

    async def _update_session_state_after_turn(
        self,
        *,
        session_state: Optional[QASessionState],
        question: str,
        effective_question: str,
        answer: str,
        query_type: str,
        ambiguity: str,
        rewrite_meta: dict[str, Any],
        sources: list[dict[str, Any]],
    ) -> None:
        if session_state is None:
            return
        session_state.last_intent = "question"
        session_state.ambiguity = str(ambiguity or "clear")[:20]
        if str(ambiguity or "").strip() != "needs_clarification":
            subject = self._extract_state_subject(question=question, answer=answer, sources=sources)
            if subject:
                session_state.last_subject = subject[:255]
            session_state.last_entity_ids = self._collect_entity_ids_from_sources(sources)
            self._clear_pending_clarification(session_state=session_state)

        history = session_state.rewrite_history if isinstance(session_state.rewrite_history, list) else []
        history.append(
            {
                "at": now_cn_aware().isoformat(),
                "question": question[:500],
                "effective_question": effective_question[:500],
                "query_type": str(query_type or QUERY_TYPE_UNKNOWN),
                "ambiguity": str(ambiguity or "clear"),
                "rewrite_meta": rewrite_meta if isinstance(rewrite_meta, dict) else {},
            }
        )
        session_state.rewrite_history = history[-20:]
        await self.db.flush()

    def _collect_entity_ids_from_sources(self, sources: list[dict[str, Any]]) -> list[int]:
        ids: list[int] = []
        for source in sources:
            if not isinstance(source, dict):
                continue
            meta = source.get("metadata") if isinstance(source.get("metadata"), dict) else {}
            for key in (
                "project_id",
                "task_id",
                "signal_id",
                "user_id",
                "entity_id",
                "subject_entity_id",
                "object_entity_id",
            ):
                value = meta.get(key)
                try:
                    if value is not None:
                        ids.append(int(value))
                except Exception:
                    continue
            matched = meta.get("matched_project_ids")
            if isinstance(matched, list):
                for value in matched:
                    try:
                        ids.append(int(value))
                    except Exception:
                        continue
        uniq: list[int] = []
        for value in ids:
            if value not in uniq:
                uniq.append(value)
        return uniq[:20]

    @staticmethod
    def _extract_project_subject_from_answer(answer: str) -> str:
        text = str(answer or "").strip()
        if not text:
            return ""
        patterns = (
            r"项目名(?:称)?[：: ]*[“\"]?([^”\"，。；、\n]{2,80})[”\"]?",
            r"项目名[“\"]([^”\"\n]{2,80})[”\"]",
            r"项目\s*#\d+[《“\"]([^》”\"\n]{2,80})[》”\"]",
            r"#\d+\s+([^\n（(]{2,80})[（(]",
            r"项目[《“\"]([^》”\"\n]{2,80})[》”\"]",
        )
        for pattern in patterns:
            match = re.search(pattern, text)
            if match:
                return str(match.group(1)).strip()[:255]
        return ""

    def _extract_state_subject(
        self,
        *,
        question: str,
        answer: str,
        sources: list[dict[str, Any]],
    ) -> str:
        for source in sources:
            if not isinstance(source, dict):
                continue
            meta = source.get("metadata") if isinstance(source.get("metadata"), dict) else {}
            project_keyword = str(meta.get("query_keyword") or "").strip()
            if project_keyword:
                return project_keyword[:255]
            subject_name = str(meta.get("subject_name") or "").strip()
            object_name = str(meta.get("object_name") or "").strip()
            if subject_name and object_name:
                return f"{subject_name}->{object_name}"[:255]

        answer_subject = self._extract_project_subject_from_answer(answer)
        if answer_subject:
            return answer_subject[:255]

        text = str(question or "").strip()
        normalizer = getattr(self, "_normalize_lookup_text", None)
        normalized_question = normalizer(text) if callable(normalizer) else text
        if 2 <= len(normalized_question) <= 64:
            return normalized_question[:255]

        return ""

    async def get_session_state_replay(self, *, session_id: int, user_id: int) -> dict[str, Any]:
        session = (
            await self.db.execute(
                select(QASession).where(
                    and_(
                        QASession.tenant_id == self.tenant_id,
                        QASession.id == int(session_id),
                        QASession.user_id == int(user_id),
                    )
                )
            )
        ).scalar_one_or_none()
        if session is None:
            raise ValueError("qa session 不存在")

        state = await self._ensure_session_state(session_id=int(session_id))
        history = state.rewrite_history if isinstance(state.rewrite_history, list) else []
        return {
            "session_id": int(session_id),
            "last_subject": str(state.last_subject or "") if state else "",
            "last_entity_ids": state.last_entity_ids if state else [],
            "last_intent": str(state.last_intent or "") if state else "",
            "ambiguity": str(state.ambiguity or "") if state else "",
            "rewrite_history": history[-20:],
            "pending_clarification": state.pending_clarification if state else None,
            "pending_clarification_expires_at": (
                state.pending_clarification_expires_at.isoformat()
                if state and state.pending_clarification_expires_at
                else None
            ),
            "updated_at": state.updated_at.isoformat() if state and state.updated_at else None,
        }