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

from .shared import *


class AdaptiveRAGRetrievalGraphMixin:
    async def _retrieve_graph_context(
        self,
        question: str,
        *,
        project_id: Optional[int],
        resolved_entities: Optional[dict[str, Any]] = None,
    ) -> list[RetrievedItem]:
        try:
            resolution = resolved_entities or {}
            matched_entities = (
                resolution.get("matched_entities") if isinstance(resolution.get("matched_entities"), list) else []
            )
            anchor_ids_raw = (
                resolution.get("anchor_entity_ids") if isinstance(resolution.get("anchor_entity_ids"), list) else []
            )
            anchor_ids: list[int] = []
            for value in anchor_ids_raw:
                try:
                    anchor_ids.append(int(value))
                except Exception:
                    continue

            seed_entities: list[KGEntity] = []
            if matched_entities:
                ids: list[int] = []
                for row in matched_entities:
                    if not isinstance(row, dict):
                        continue
                    try:
                        ids.append(int(row.get("entity_id")))
                    except Exception:
                        continue
                if ids:
                    seed_entities = (
                        await self.db.execute(
                            select(KGEntity).where(
                                and_(
                                    KGEntity.tenant_id == self.tenant_id,
                                    KGEntity.id.in_(ids),
                                )
                            )
                        )
                    ).scalars().all()

            if not seed_entities:
                term_collector = getattr(self, "_collect_lookup_terms", None)
                terms = term_collector(question=question, limit=6) if callable(term_collector) else []
                if not terms:
                    normalizer = getattr(self, "_normalize_lookup_text", None)
                    normalized = normalizer(question) if callable(normalizer) else str(question or "").strip()
                    if normalized:
                        terms = [normalized[:32]]
                if terms:
                    conds = []
                    for term in terms[:6]:
                        like = f"%{term}%"
                        conds.append(KGEntity.display_name.ilike(like))
                        conds.append(KGEntity.entity_key.ilike(like))
                    seed_entities = (
                        await self.db.execute(
                            select(KGEntity)
                            .where(
                                and_(
                                    KGEntity.tenant_id == self.tenant_id,
                                    or_(*conds),
                                )
                            )
                            .order_by(KGEntity.last_seen_at.desc(), KGEntity.id.desc())
                            .limit(15)
                        )
                    ).scalars().all()
            if not seed_entities and anchor_ids:
                seed_entities = (
                    await self.db.execute(
                        select(KGEntity).where(
                            and_(
                                KGEntity.tenant_id == self.tenant_id,
                                KGEntity.id.in_(anchor_ids[:10]),
                            )
                        )
                    )
                ).scalars().all()
            if not seed_entities:
                return []

            seed_ids = [int(x.id) for x in seed_entities]
            relation_rows = (
                await self.db.execute(
                    select(KGRelation)
                    .where(
                        and_(
                            KGRelation.tenant_id == self.tenant_id,
                            or_(
                                KGRelation.subject_entity_id.in_(seed_ids),
                                KGRelation.object_entity_id.in_(seed_ids),
                            ),
                        )
                    )
                    .order_by(KGRelation.weight.desc(), KGRelation.id.desc())
                    .limit(120)
                )
            ).scalars().all()

            entity_ids: set[int] = set(seed_ids)
            for rel in relation_rows:
                entity_ids.add(int(rel.subject_entity_id))
                entity_ids.add(int(rel.object_entity_id))
            entities = (
                await self.db.execute(
                    select(KGEntity).where(
                        and_(
                            KGEntity.tenant_id == self.tenant_id,
                            KGEntity.id.in_(list(entity_ids)),
                        )
                    )
                )
            ).scalars().all()
            entity_map = {int(row.id): row for row in entities}

            relation_ids = [int(row.id) for row in relation_rows]
            evidence_rows: list[KGEvidence] = []
            if relation_ids or entity_ids:
                evidence_rows = (
                    await self.db.execute(
                        select(KGEvidence)
                        .where(
                            and_(
                                KGEvidence.tenant_id == self.tenant_id,
                                or_(
                                    KGEvidence.relation_id.in_(relation_ids or [-1]),
                                    KGEvidence.entity_id.in_(list(entity_ids) or [-1]),
                                ),
                            )
                        )
                        .order_by(KGEvidence.occurred_at.desc(), KGEvidence.id.desc())
                        .limit(220)
                    )
                ).scalars().all()
            evidence_by_relation: dict[int, str] = {}
            evidence_by_entity: dict[int, str] = {}
            evidence_updated_at_by_relation: dict[int, str] = {}
            evidence_updated_at_by_entity: dict[int, str] = {}
            for ev in evidence_rows:
                excerpt = str(ev.content_excerpt or "").strip()
                occurred_at = self._isoformat_or_none(getattr(ev, "occurred_at", None))
                if not excerpt:
                    excerpt = ""
                if ev.relation_id is not None and int(ev.relation_id) not in evidence_by_relation:
                    evidence_by_relation[int(ev.relation_id)] = excerpt[:320]
                if occurred_at and ev.relation_id is not None and int(ev.relation_id) not in evidence_updated_at_by_relation:
                    evidence_updated_at_by_relation[int(ev.relation_id)] = occurred_at
                if ev.entity_id is not None and int(ev.entity_id) not in evidence_by_entity:
                    evidence_by_entity[int(ev.entity_id)] = excerpt[:320]
                if occurred_at and ev.entity_id is not None and int(ev.entity_id) not in evidence_updated_at_by_entity:
                    evidence_updated_at_by_entity[int(ev.entity_id)] = occurred_at

            out: list[RetrievedItem] = []
            for rel in relation_rows:
                subject = entity_map.get(int(rel.subject_entity_id))
                target = entity_map.get(int(rel.object_entity_id))
                if subject is None or target is None:
                    continue
                subject_payload = self._safe_dict(subject.payload)
                target_payload = self._safe_dict(target.payload)
                if project_id is not None:
                    subject_pid = subject_payload.get("project_id")
                    target_pid = target_payload.get("project_id")
                    if subject_pid is not None and int(subject_pid) != int(project_id):
                        if target_pid is not None and int(target_pid) != int(project_id):
                            continue
                evidence_text = evidence_by_relation.get(int(rel.id))
                if not evidence_text:
                    evidence_text = evidence_by_entity.get(int(subject.id)) or evidence_by_entity.get(int(target.id)) or ""
                content = (
                    f"图路径: {str(subject.display_name or subject.entity_key)} --[{str(rel.relation_type)}]--> "
                    f"{str(target.display_name or target.entity_key)}\n"
                    f"证据: {evidence_text or '无显式证据摘录'}"
                )
                weight = self._to_float(rel.weight, 1.0)
                score = min(0.99, max(0.45, 0.62 + min(weight, 1.5) * 0.18))
                relation_project_id = (
                    subject_payload.get("project_id")
                    or target_payload.get("project_id")
                    or project_id
                )
                metadata = self._normalize_source_registry_metadata(
                    source_type="graph",
                    metadata={
                        "relation_id": int(rel.id),
                        "relation_type": str(rel.relation_type or ""),
                        "subject_entity_id": int(subject.id),
                        "object_entity_id": int(target.id),
                        "subject_name": str(subject.display_name or subject.entity_key),
                        "object_name": str(target.display_name or target.entity_key),
                        "path": f"{subject.id}->{str(rel.relation_type)}->{target.id}",
                        "project_id": relation_project_id,
                        "source_ref": getattr(rel, "source_ref", None),
                    },
                    explicit_updated_at=(
                        evidence_updated_at_by_relation.get(int(rel.id))
                        or evidence_updated_at_by_entity.get(int(subject.id))
                        or evidence_updated_at_by_entity.get(int(target.id))
                    ),
                )
                out.append(
                    RetrievedItem(
                        source="graph",
                        content=content[:1200],
                        score=score,
                        source_label="Knowledge Graph",
                        metadata=metadata,
                    )
                )

            if not out:
                for seed in seed_entities[:DEFAULT_TOP_K]:
                    payload = self._safe_dict(seed.payload)
                    if project_id is not None and payload.get("project_id") is not None:
                        if int(payload.get("project_id")) != int(project_id):
                            continue
                    excerpt = evidence_by_entity.get(int(seed.id)) or ""
                    metadata = self._normalize_source_registry_metadata(
                        source_type="graph",
                        metadata={
                            "entity_id": int(seed.id),
                            "entity_type": str(seed.entity_type or ""),
                            "entity_key": str(seed.entity_key or ""),
                            "project_id": payload.get("project_id") or project_id,
                            "source_ref": getattr(seed, "source_ref", None),
                        },
                        explicit_updated_at=(
                            evidence_updated_at_by_entity.get(int(seed.id))
                            or getattr(seed, "last_seen_at", None)
                        ),
                    )
                    out.append(
                        RetrievedItem(
                            source="graph",
                            content=f"实体: {str(seed.display_name or seed.entity_key)}\n证据: {excerpt or '无显式证据摘录'}",
                            score=0.58,
                            source_label="Knowledge Graph",
                            metadata=metadata,
                        )
                    )
            out.sort(
                key=lambda item: (
                    item.score,
                    self._source_priority_sort_value(item.metadata),
                ),
                reverse=True,
            )
            return out[: max(DEFAULT_TOP_K, 6)]
        except Exception:
            return []