rag/automation_sources.py

from typing import Any, Dict, List, Tuple
import time

from core.books.storage import fetch_document_metadata

# =========================
# Helpers
# =========================
def _as_dict(x: Any) -> Dict[str, Any]:
    return x if isinstance(x, dict) else {}


def _as_list(x: Any) -> List[Any]:
    return x if isinstance(x, list) else []


def _get_planning_obj(sub: Dict[str, Any]) -> Dict[str, Any]:
    """
    OLD schema: sub["planning"] موجود
    NEW schema: الخطة نفسها هي sub مباشرة (plan object)
    """
    planning = sub.get("planning")
    if isinstance(planning, dict):
        return planning
    return sub


def _normalize_author(author_val: Any) -> str:
    """
    يحاول يطلع اسم مؤلف مناسب للاقتباس.
    """
    if isinstance(author_val, list):
        for a in author_val:
            if isinstance(a, str) and a.strip():
                return a.strip()
        return "مؤلف"
    if isinstance(author_val, str) and author_val.strip():
        return author_val.strip()
    return "مؤلف"


def _normalize_year(year_val: Any) -> Any:
    """
    يرجّع int لو ينفع، وإلا None.
    """
    if isinstance(year_val, int):
        return year_val
    if isinstance(year_val, str):
        y = year_val.strip()
        if y.isdigit():
            try:
                return int(y)
            except Exception:
                return None
    return None


# =========================
# Query Bundle Builder
# =========================
def build_query_bundle(
    sub: dict,
    chapter_title: str,
    section_title: str,
    subsection_title: str = "",
) -> List[str]:
    planning = _get_planning_obj(sub)
    queries: List[str] = []

    # 1) core_idea
    core_idea = planning.get("core_idea")
    if isinstance(core_idea, str) and core_idea.strip():
        queries.append(core_idea.strip())

    # 2) key_points (top 3)
    for kp in _as_list(planning.get("key_points"))[:3]:
        if isinstance(kp, str) and kp.strip():
            queries.append(kp.strip())

    # 3) suggested_queries (لو موجودة في أي شكل)
    for q in _as_list(sub.get("suggested_queries"))[:2]:
        if isinstance(q, str) and q.strip():
            queries.append(q.strip())

    # 4) subsection title
    if isinstance(subsection_title, str) and subsection_title.strip():
        queries.append(subsection_title.strip())
    elif isinstance(sub.get("title"), str) and sub.get("title", "").strip():
        # fallback للقديم
        queries.append(sub["title"].strip())

    # 5) context query
    context_q = f"{chapter_title} - {section_title}".strip(" -")
    if context_q:
        queries.append(context_q)

    # de-dup مع الحفاظ على الترتيب
    seen = set()
    final = []
    for q in queries:
        qq = q.strip()
        if qq and qq not in seen:
            seen.add(qq)
            final.append(qq)

    return final


# =========================
# Citations
# =========================
def make_in_text_citation(author, year, page_start, title=None):
    a = _normalize_author(author)
    y = year if isinstance(year, int) else "د.ت"
    p = page_start if isinstance(page_start, int) else "؟"
    t = title.strip() if isinstance(title, str) and title.strip() else "مصدر بدون عنوان"

    if a == "مؤلف":
        return f"({t}، {y}، ص. {p})"
    citation = f"({a}، {y}، ص. {p})"
    
    return citation


def make_reference_apa(author, year, title, source_url=None,publisher_or_journal="", normalized_source_type="pdf"):
    a = _normalize_author(author)
    y = year if isinstance(year, int) else "د.ت"
    t = title.strip() if isinstance(title, str) and title.strip() else "مصدر بدون عنوان"
    p = publisher_or_journal.strip() if isinstance(publisher_or_journal, str) and publisher_or_journal.strip() else ""
    n= normalized_source_type.strip() if isinstance(normalized_source_type, str) and normalized_source_type.strip() else ""

    if source_url:
        base = f"المؤلفون:{a}. التاريخ:({y}). العنوان:{t}. الجامعة/الناشر/المجلة:{p}. نوع المصدر:{n}."
        return f"{base} متاح على لينك: {source_url}"
    return base


# =========================
# RAG for Subsection / Plan
# =========================
# def build_rag_context_for_subsection(
#     rag_engines: List,
#     sub: dict,
#     chapter_title: str,
#     section_title: str,
#     subsection_title: str = "",
#     top_k: int = 5,
# ) -> dict:
#     queries = build_query_bundle(
#         sub=sub,
#         chapter_title=chapter_title,
#         section_title=section_title,
#         subsection_title=subsection_title,
#     )

#     all_hits = []

#     # 🔍 search in all collections
#     for rag in rag_engines:
#         try:
#             print("🔍 Searching collection:", getattr(rag, "collection", "unknown"))
#             print("🔍 Queries:", queries)

#             # count اختياري (لو متاح في engine)
#             try:
#                 count = rag.qdrant.count(rag.collection, exact=True)
#                 print("📦 Collection count:", count)
#             except Exception:
#                 pass

#             hits = rag.retrieve(queries=queries) or []
#             all_hits.extend(hits)
#         except Exception as e:
#             print(
#                 f"⚠️ RAG engine failed on collection {getattr(rag, 'collection', 'unknown')}: {e}"
#             )
#             continue

#     if not all_hits:
#         return {
#             "query_bundle": queries,
#             "selected_k": 0,
#             "chunks": [],
#             "coverage_note": "لم يتم العثور على مراجع مناسبة لهذا المحور.",
#         }

#     # sort globally + dedupe (doc_id, chunk_id)
#     all_hits = sorted(all_hits, key=lambda x: getattr(x, "score", 0.0), reverse=True)

#     unique: List[Any] = []
#     seen_keys = set()
#     for h in all_hits:
#         payload = _as_dict(getattr(h, "payload", {}))
#         key = (payload.get("doc_id"), getattr(h, "id", None))
#         if key in seen_keys:
#             continue
#         seen_keys.add(key)
#         unique.append(h)
#         if len(unique) >= top_k:
#             break

#     chunks = []
#     for h in unique:
#         payload = _as_dict(getattr(h, "payload", {}))
#         doc_id = payload.get("doc_id")
#         source_url = payload.get("source_url")

#         md = fetch_document_metadata(doc_id) if doc_id else None
#         md = _as_dict(md)

#         # metadata fallback names
#         author = md.get("authors") or md.get("author")
#         year = _normalize_year(md.get("year"))
#         title = md.get("title") or payload.get("title")

#         page_start = payload.get("page_start")
#         if isinstance(page_start, str) and page_start.isdigit():
#             page_start = int(page_start)

#         page_end = payload.get("page_end")
#         if isinstance(page_end, str) and page_end.isdigit():
#             page_end = int(page_end)

#         chunks.append(
#             {
#                 "chunk_id": getattr(h, "id", None),
#                 "score": float(getattr(h, "score", 0.0) or 0.0),
#                 "doc_id": doc_id,
#                 "source_url": source_url,
#                 "page_start": page_start,
#                 "page_end": page_end,
#                 "text": payload.get("text"),
#                 "title": title,
#                 "author": _normalize_author(author),
#                 "year": year,
#                 "in_text_citation": make_in_text_citation(
#                     author=author,
#                     year=year,
#                     page_start=page_start,
#                     source_url=source_url,
#                 ),
#                 "reference_apa": make_reference_apa(
#                     author=author, year=year, title=title, source_url=source_url
#                 ),
#             }
#         )

#     return {
#         "query_bundle": queries,
#         "selected_k": len(chunks),
#         "chunks": chunks,
#         "coverage_note": "تم اختيار أفضل المراجع المتوافقة مع فكرة المحور وحدوده.",
#     }


def build_rag_context_for_subsection(
    rag_engines: List,
    sub: dict,
    chapter_title: str,
    section_title: str,
    subsection_title: str = "",
    top_k: int = 5,
) -> dict:

    queries = build_query_bundle(
        sub=sub,
        chapter_title=chapter_title,
        section_title=section_title,
        subsection_title=subsection_title,
    )

    all_hits = []

    # 🔍 search in all collections
    for rag in rag_engines:
        try:
            print("🔍 Searching collection:", getattr(rag, "collection", "unknown"))
            print("🔍 Queries:", queries)

            try:
                count = rag.qdrant.count(rag.collection, exact=True)
                print("📦 Collection count:", count)
            except Exception:
                pass

            hits = rag.retrieve(queries=queries) or []
            all_hits.extend(hits)

        except Exception as e:
            print(
                f"⚠️ RAG engine failed on collection {getattr(rag, 'collection', 'unknown')}: {e}"
            )
            continue

    if not all_hits:
        return {
            "query_bundle": queries,
            "selected_k": 0,
            "chunks": [],
            "coverage_note": "لم يتم العثور على مراجع مناسبة لهذا المحور.",
        }

    # ==============================
    # Diversity + Source-Aware Ranking
    # ==============================

    # sort by base score
    all_hits = sorted(
        all_hits,
        key=lambda x: float(getattr(x, "score", 0.0) or 0.0),
        reverse=True,
    )

    max_per_doc = 2
    min_unique_docs = 3

    doc_counter = {}
    rescored_hits = []

    for h in all_hits:
        payload = _as_dict(getattr(h, "payload", {}))
        doc_id = payload.get("doc_id")

        if not doc_id:
            continue

        base_score = float(getattr(h, "score", 0.0) or 0.0)
        current_count = doc_counter.get(doc_id, 0)

        if current_count >= max_per_doc:
            continue

        diversity_penalty = 1 - (0.25 * current_count)
        adjusted_score = base_score * diversity_penalty

        doc_counter[doc_id] = current_count + 1

        rescored_hits.append((h, adjusted_score))

    # نرتب بالـ adjusted_score
    rescored_hits = sorted(rescored_hits, key=lambda x: x[1], reverse=True)

    # ناخد top_k
    selected = rescored_hits[:top_k]

    unique_doc_ids = {
        _as_dict(getattr(h, "payload", {})).get("doc_id")
        for h, _ in selected
        if _as_dict(getattr(h, "payload", {})).get("doc_id")
    }

    if len(unique_doc_ids) < min_unique_docs:
        print("⚠️ Warning: Low source diversity detected")

    # ==============================
    # Build Final Chunks
    # ==============================

    chunks = []

    for h, adjusted_score in selected:
        payload = _as_dict(getattr(h, "payload", {}))
        doc_id = payload.get("doc_id")
        source_url = payload.get("source_url")

        md = fetch_document_metadata(doc_id) if doc_id else None
        md = _as_dict(md)

        author = md.get("authors") or md.get("author")
        year = _normalize_year(md.get("year"))
        title = md.get("title") or payload.get("title")
        publisher_or_journal= md.get("publisher_or_journal") or ""
        normalized_source_type= md.get("normalized_source_type", "pdf")

        page_start = payload.get("page_start")
        if isinstance(page_start, str) and page_start.isdigit():
            page_start = int(page_start)

        page_end = payload.get("page_end")
        if isinstance(page_end, str) and page_end.isdigit():
            page_end = int(page_end)

        chunks.append(
            {
                "chunk_id": getattr(h, "id", None),
                "score": float(getattr(h, "score", 0.0) or 0.0),
                "adjusted_score": float(adjusted_score),
                "doc_id": doc_id,
                "source_url": source_url,
                "page_start": page_start,
                "page_end": page_end,
                "text": payload.get("text"),
                "title": title,
                "author": _normalize_author(author),
                "year": year,
                "in_text_citation": make_in_text_citation(
                    author=author,
                    year=year,
                    page_start=page_start,
                    title=title,
                ),
                "reference_apa": make_reference_apa(
                    author=author,
                    year=year,
                    title=title,
                    source_url=source_url,
                    publisher_or_journal=publisher_or_journal,
                    normalized_source_type=normalized_source_type,
                ),
            }
        )

    return {
        "query_bundle": queries,
        "selected_k": len(chunks),
        "chunks": chunks,
        "coverage_note": "تم اختيار أفضل المراجع المتوافقة مع فكرة المحور وحدوده.",
    }


# =========================
# Book-Level Automation
# Supports NEW + OLD schema
# =========================
def run_rag_automation_on_book_sources(
    planned_book: dict,
    rag_engines: List,
    top_k: int = 5,
    sleep_s: float = 0.1,
) -> dict:
    for chapter in _as_list(planned_book.get("chapters")):
        ch_title = chapter.get("chapter_title", "")

        for section in _as_list(chapter.get("sections")):
            # NEW: section_title | OLD: title
            sec_title = section.get("section_title", section.get("title", ""))

            # -------------------------
            # NEW SCHEMA: plans (dict)
            # -------------------------
            plans = section.get("plans")
            if isinstance(plans, dict):
                for sub_title, plan_obj in plans.items():
                    if not isinstance(plan_obj, dict):
                        continue

                    rag_context = build_rag_context_for_subsection(
                        rag_engines=rag_engines,
                        sub=plan_obj,  # plan نفسه
                        chapter_title=ch_title,
                        section_title=sec_title,
                        subsection_title=str(sub_title),
                        top_k=top_k,
                    )

                    section["plans"][sub_title]["rag_context"] = rag_context
                    time.sleep(sleep_s)
                continue

            # -------------------------
            # OLD SCHEMA fallback: subsections (list)
            # -------------------------
            for sub in _as_list(section.get("subsections")):
                if not isinstance(sub, dict):
                    continue

                rag_context = build_rag_context_for_subsection(
                    rag_engines=rag_engines,
                    sub=sub,
                    chapter_title=ch_title,
                    section_title=sec_title,
                    subsection_title=sub.get("title", ""),
                    top_k=top_k,
                )

                sub["rag_context"] = rag_context
                time.sleep(sleep_s)

    return planned_book