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

from app.schemas.context_workspace import (
    ContextCapsule,
    ResearchWorkspace,
    normalize_context_capsules_payload,
    normalize_research_workspace_payload,
)

from .shared import *


class AdaptiveRAGSessionStateMixin:
    @staticmethod
    def _normalize_active_workspace(payload: Any) -> dict[str, Any]:
        return normalize_research_workspace_payload(payload)

    @staticmethod
    def _normalize_context_capsules(payload: Any) -> list[dict[str, Any]]:
        return normalize_context_capsules_payload(payload)

    @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 {}
        )
        customer_interaction_context = (
            response_payload.get("customer_interaction_context_packet")
            if isinstance(response_payload.get("customer_interaction_context_packet"), dict)
            else {}
        )
        current_focus_object = (
            customer_interaction_context.get("current_focus_object")
            if isinstance(customer_interaction_context.get("current_focus_object"), 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
        if project_id is None:
            try:
                project_id = int(current_focus_object.get("object_id")) if current_focus_object.get("object_id") is not None else None
            except Exception:
                project_id = None
        target_ref = (
            str(current_focus_object.get("object_ref") or "").strip()
            or 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(current_focus_object.get("object_type") or "").strip().lower()
                    or
                    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]:
        if anchor_resolution is None:
            resolution = type(
                "_FollowupResolution",
                (),
                {"followup_mode": "none", "anchor_mode": "none"},
            )()
        else:
            resolution = anchor_resolution
        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=[],
                active_workspace=None,
                context_capsules=[],
                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"}

        if anchor_resolution is None:
            resolution = type(
                "_FollowupResolution",
                (),
                {"followup_mode": "none", "anchor_mode": "none"},
            )()
        else:
            resolution = anchor_resolution
        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:
            active_workspace = self._normalize_active_workspace(getattr(session_state, "active_workspace", None))
            if active_workspace:
                primary_title = str(active_workspace.get("primary_object_title") or "").strip()
                workspace_topic = str(active_workspace.get("topic") or "").strip()
                candidate_names = [
                    str(item.get("display_name") or "").strip()
                    for item in list(active_workspace.get("candidate_items") or [])[:3]
                    if isinstance(item, dict) and str(item.get("display_name") or "").strip()
                ]
                scope_bits = [bit for bit in [primary_title, workspace_topic] if bit]
                if candidate_names:
                    scope_bits.append("候选对象：" + "、".join(candidate_names))
                if scope_bits:
                    rewritten = f"围绕当前研究工作区（{'；'.join(scope_bits[:3])}）回答：{current}"
                    return rewritten[:1200], {
                        "mode": "active_workspace",
                        "workspace_id": str(active_workspace.get("workspace_id") or "").strip() or None,
                        "workspace_topic": workspace_topic or None,
                        "primary_object_ref": str(active_workspace.get("primary_object_ref") or "").strip() or None,
                        "candidate_set_binding": str(active_workspace.get("candidate_set_binding") or "").strip() or 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]],
        response_payload: Optional[dict[str, Any]] = None,
        route_decision: Optional[dict[str, Any]] = None,
    ) -> 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":
            next_workspace = self._workspace_from_turn(
                response_payload=response_payload,
                route_decision=route_decision,
                question=question,
                answer=answer,
                sources=sources,
            )
            if next_workspace:
                session_state.active_workspace = next_workspace
                session_state.context_capsules = self._merge_context_capsules(
                    existing=getattr(session_state, "context_capsules", None),
                    workspace=next_workspace,
                )
            subject = self._extract_state_subject(question=question, answer=answer, sources=sources)
            if not subject and next_workspace:
                subject = (
                    str(next_workspace.get("primary_object_title") or "").strip()
                    or str(next_workspace.get("topic") or "").strip()
                )
            if subject:
                session_state.last_subject = subject[:255]
            collected_entity_ids = self._collect_entity_ids_from_sources(sources)
            if not collected_entity_ids and next_workspace:
                workspace_project_id = next_workspace.get("primary_object_id")
                try:
                    if workspace_project_id is not None:
                        collected_entity_ids = [int(workspace_project_id)]
                except Exception:
                    collected_entity_ids = []
            if collected_entity_ids:
                session_state.last_entity_ids = collected_entity_ids
            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 _workspace_from_turn(
        self,
        *,
        response_payload: Optional[dict[str, Any]],
        route_decision: Optional[dict[str, Any]],
        question: str,
        answer: str,
        sources: list[dict[str, Any]],
    ) -> dict[str, Any]:
        payload = response_payload if isinstance(response_payload, dict) else {}
        workspace = self._normalize_active_workspace(payload.get("research_workspace"))
        if workspace:
            return workspace
        normalized_route = route_decision if isinstance(route_decision, dict) else {}
        effective_source_policy = str(
            normalized_route.get("effective_source_policy") or normalized_route.get("route_mode") or ""
        ).strip().lower()
        interaction_context = (
            dict(payload.get("customer_interaction_context_packet") or {})
            if isinstance(payload.get("customer_interaction_context_packet"), dict)
            else {}
        )
        current_focus = (
            dict(interaction_context.get("current_focus_object") or {})
            if isinstance(interaction_context.get("current_focus_object"), dict)
            else {}
        )
        if effective_source_policy not in {"external", "hybrid"}:
            return {}
        object_ref = str(current_focus.get("object_ref") or "").strip()
        object_title = str(current_focus.get("title") or "").strip()
        object_id = current_focus.get("object_id")
        try:
            object_id = int(object_id) if object_id is not None else None
        except Exception:
            object_id = None
        if not object_ref and not object_title:
            return {}
        topic = str(question or "").strip()[:120] or "外部研究"
        summary = str(answer or "").strip()[:240]
        return ResearchWorkspace(
            workspace_id=f"ws_{hashlib.sha1(topic.encode('utf-8'), usedforsecurity=False).hexdigest()[:12]}",
            workspace_kind="hybrid_research" if object_ref else "external_research",
            topic=topic,
            primary_object_ref=object_ref or None,
            primary_object_id=object_id,
            primary_object_title=object_title or None,
            research_goal=summary or None,
            analysis_dimensions=[],
            evidence_refs=[],
            confidence_state="low",
            open_questions=[],
            next_recommended_actions=[],
            status="active",
        ).model_dump()

    def _merge_context_capsules(
        self,
        *,
        existing: Any,
        workspace: dict[str, Any],
    ) -> list[dict[str, Any]]:
        normalized_workspace = self._normalize_active_workspace(workspace)
        if not normalized_workspace:
            return self._normalize_context_capsules(existing)
        capsule = ContextCapsule(
            workspace_id=str(normalized_workspace.get("workspace_id") or "").strip(),
            workspace_kind=str(normalized_workspace.get("workspace_kind") or "").strip(),
            topic=str(normalized_workspace.get("topic") or "").strip(),
            primary_object_ref=str(normalized_workspace.get("primary_object_ref") or "").strip() or None,
            candidate_set_binding=str(normalized_workspace.get("candidate_set_binding") or "").strip() or None,
            summary=str(normalized_workspace.get("research_goal") or normalized_workspace.get("topic") or "").strip(),
            open_questions=[
                str(item).strip()
                for item in list(normalized_workspace.get("open_questions") or [])[:4]
                if str(item).strip()
            ],
            next_action=(
                str((normalized_workspace.get("next_recommended_actions") or [None])[0] or "").strip()
                or None
            ),
        ).model_dump()
        capsules = self._normalize_context_capsules(existing)
        merged = [capsule]
        for item in capsules:
            if not isinstance(item, dict):
                continue
            if str(item.get("workspace_id") or "").strip() == str(capsule.get("workspace_id") or "").strip():
                continue
            merged.append(item)
        return merged[:8]

    def _collect_entity_ids_from_sources(self, sources: list[dict[str, Any]]) -> list[int]:
        explicit_ids: list[int] = []
        graph_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"):
                value = meta.get(key)
                try:
                    if value is not None:
                        explicit_ids.append(int(value))
                except Exception:
                    continue
            matched = meta.get("matched_project_ids")
            if isinstance(matched, list):
                for value in matched:
                    try:
                        explicit_ids.append(int(value))
                    except Exception:
                        continue
            for key in ("entity_id", "subject_entity_id", "object_entity_id"):
                value = meta.get(key)
                try:
                    if value is not None:
                        graph_ids.append(int(value))
                except Exception:
                    continue
        candidates = explicit_ids or graph_ids
        uniq: list[int] = []
        for value in candidates:
            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]
            source_type = str(meta.get("source_type") or "").strip().lower()
            object_ref = str(meta.get("object_ref") or "").strip().lower()
            if source_type == "customer_visible_object" or object_ref.startswith("project:"):
                subject_name = str(meta.get("subject_name") or "").strip()
                if subject_name:
                    return subject_name[:255]

        for source in sources:
            if not isinstance(source, dict):
                continue
            meta = source.get("metadata") if isinstance(source.get("metadata"), dict) else {}
            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:],
            "active_workspace": self._normalize_active_workspace(getattr(state, "active_workspace", None)),
            "context_capsules": self._normalize_context_capsules(getattr(state, "context_capsules", None)),
            "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,
        }