"""Build RAG context for course-plan generation.

Wraps :class:`encoach_vector.services.embedding_service.EmbeddingService` to
return the top-k chunks from a plan's indexed reference sources, scoped to
the canonical ``content_type='course_plan_source'`` and trimmed to a
character budget so the LLM prompt stays within context limits.

Indexer convention
------------------
:class:`SourceIndexer` stores each ``encoach.course.plan.source`` chunk
with metadata ``{plan_id, source_id, kind, title, entity_id=plan_id}``
(see :mod:`source_indexer`). This builder uses that convention to scope
retrieval to a single plan: it filters embeddings on ``entity_id=plan.id``
and ``content_type='course_plan_source'``.

Usage
-----
.. code-block:: python

    builder = RAGContextBuilder(env)
    ctx = builder.for_week(plan, week, skill='reading', k=6)
    if ctx['passages']:
        prompt += "\\n\\nReference passages from your uploaded book(s):\\n"
        prompt += ctx['as_prompt_block']
"""

from __future__ import annotations

import logging

_logger = logging.getLogger(__name__)


# Soft char budget for the assembled passage block. ~3500 chars is roughly
# 800-1000 tokens, leaving room for the rest of the prompt (~2-3 K tokens)
# under a 16 K context window.
DEFAULT_CHAR_BUDGET = 3500
DEFAULT_K = 6


class RAGContextBuilder:
    """Produce RAG passages for a (plan, week, skill) tuple."""

    CONTENT_TYPE = 'course_plan_source'

    def __init__(self, env, *, char_budget: int = DEFAULT_CHAR_BUDGET):
        self.env = env
        self.char_budget = int(char_budget)
        self._svc = None  # lazily created on first call

    # ------------------------------------------------------------------
    def _service(self):
        """Lazy-load the embedding service so import never fails fatally."""
        if self._svc is not None:
            return self._svc
        try:
            from odoo.addons.encoach_vector.services.embedding_service import (
                EmbeddingService,
            )
        except ImportError:
            _logger.warning(
                'encoach_vector not installed; RAG context disabled.'
            )
            return None
        self._svc = EmbeddingService(self.env)
        return self._svc

    # ------------------------------------------------------------------
    def has_indexed_sources(self, plan) -> bool:
        """Return True iff the plan has at least one indexed (status='indexed') source."""
        try:
            count = self.env['encoach.course.plan.source'].sudo().search_count([
                ('plan_id', '=', plan.id),
                ('status', '=', 'indexed'),
            ])
        except Exception:
            return False
        return bool(count)

    # ------------------------------------------------------------------
    def _build_query(self, plan, week, skill: str) -> str:
        """Synthesize a focused retrieval query for (week, skill).

        The query is intentionally dense: skill, week unit + focus, and the
        descriptions of the outcomes targeted by this week's items. We feed
        the descriptions (not just codes) so semantic search can match book
        passages that talk about, say, "personal introductions" even when
        no outcome code appears in the source text.
        """
        outcomes = plan._loads(plan.outcomes_json, {})
        items = week._loads(week.items_json, []) if week else []
        codes_for_skill: list[str] = []
        for it in items:
            if (it.get('skill') or '').lower() == (skill or '').lower():
                codes_for_skill.extend(it.get('outcome_codes') or [])

        outcome_descs: list[str] = []
        skill_outcomes = outcomes.get((skill or '').lower()) or []
        for o in skill_outcomes:
            if o.get('code') in set(codes_for_skill):
                outcome_descs.append((o.get('description') or '').strip())

        parts = [
            (skill or '').strip(),
            (week.unit or '').strip() if week else '',
            (week.focus or '').strip() if week else '',
            ' '.join(d for d in outcome_descs if d),
        ]
        query = ' — '.join(p for p in parts if p)
        return query or (plan.name or 'general english')

    # ------------------------------------------------------------------
    def _trim_to_budget(self, results: list[dict]) -> tuple[list[dict], int]:
        """Cut chunks until total chars <= self.char_budget.

        Returns the truncated list and the actual char count used.
        """
        out: list[dict] = []
        used = 0
        for r in results:
            text = (r.get('text') or '').strip()
            if not text:
                continue
            # Per-chunk soft cap so one giant page can't eat the whole budget.
            if len(text) > 1200:
                text = text[:1200].rsplit(' ', 1)[0] + '…'
            if used + len(text) > self.char_budget:
                # If we have nothing yet, take a hard slice so the prompt is
                # never empty when the budget is tiny.
                if not out:
                    text = text[: max(self.char_budget // 2, 200)]
                else:
                    break
            out.append({**r, 'text': text})
            used += len(text)
        return out, used

    # ------------------------------------------------------------------
    def for_week(self, plan, week, skill: str, *, k: int = DEFAULT_K) -> dict:
        """Return RAG context for a (plan, week, skill).

        :returns: ``{passages, sources, char_budget_used, query, as_prompt_block}``.
                  Empty dict-shape when the plan has no indexed sources.
        """
        empty = {
            'passages': [],
            'sources': [],
            'char_budget_used': 0,
            'query': '',
            'as_prompt_block': '',
        }
        svc = self._service()
        if svc is None:
            return empty
        if not self.has_indexed_sources(plan):
            return empty

        query = self._build_query(plan, week, skill)
        try:
            raw_results = svc.search(
                query,
                content_type=self.CONTENT_TYPE,
                limit=int(k),
                entity_id=plan.id,  # SourceIndexer stores plan_id under entity_id
            )
        except Exception:
            _logger.exception('RAG search failed for plan=%s skill=%s', plan.id, skill)
            return empty

        # Decorate with source titles for citation.
        Source = self.env['encoach.course.plan.source'].sudo()
        passages: list[dict] = []
        for r in raw_results or []:
            md = r.get('metadata') or {}
            source_id = int(md.get('source_id') or 0) or None
            title = md.get('title') or ''
            if source_id and not title:
                rec = Source.browse(source_id)
                if rec.exists():
                    title = rec.name or rec.file_name or rec.url or f'Source #{source_id}'
            passages.append({
                'text': r.get('text') or '',
                'similarity': float(r.get('score') or r.get('similarity') or 0.0),
                'source_id': source_id,
                'source_title': title,
                'chunk_index': md.get('chunk_index'),
            })

        passages, used = self._trim_to_budget(passages)
        sources = self._collect_sources(passages)
        block = self._render_prompt_block(passages)
        return {
            'passages': passages,
            'sources': sources,
            'char_budget_used': used,
            'query': query,
            'as_prompt_block': block,
        }

    # ------------------------------------------------------------------
    def for_plan(self, plan, *, k: int = 12, query: str | None = None) -> dict:
        """Plan-wide retrieval — used by ExerciseExtractor and the workbook
        cross-week pass. Falls back to ``plan.name`` when no query is given.
        """
        empty = {
            'passages': [],
            'sources': [],
            'char_budget_used': 0,
            'query': '',
            'as_prompt_block': '',
        }
        svc = self._service()
        if svc is None or not self.has_indexed_sources(plan):
            return empty
        q = (query or plan.name or 'exercises').strip()
        try:
            raw_results = svc.search(
                q,
                content_type=self.CONTENT_TYPE,
                limit=int(k),
                entity_id=plan.id,
            )
        except Exception:
            _logger.exception('RAG plan-wide search failed for plan=%s', plan.id)
            return empty
        Source = self.env['encoach.course.plan.source'].sudo()
        passages: list[dict] = []
        for r in raw_results or []:
            md = r.get('metadata') or {}
            source_id = int(md.get('source_id') or 0) or None
            title = md.get('title') or ''
            if source_id and not title:
                rec = Source.browse(source_id)
                if rec.exists():
                    title = rec.name or rec.file_name or rec.url or f'Source #{source_id}'
            passages.append({
                'text': r.get('text') or '',
                'similarity': float(r.get('score') or r.get('similarity') or 0.0),
                'source_id': source_id,
                'source_title': title,
                'chunk_index': md.get('chunk_index'),
            })
        passages, used = self._trim_to_budget(passages)
        sources = self._collect_sources(passages)
        block = self._render_prompt_block(passages)
        return {
            'passages': passages,
            'sources': sources,
            'char_budget_used': used,
            'query': q,
            'as_prompt_block': block,
        }

    # ------------------------------------------------------------------
    @staticmethod
    def _collect_sources(passages: list[dict]) -> list[dict]:
        """Group passages by source_id → ``[{source_id, title, chunks_used}]``."""
        agg: dict[int, dict] = {}
        for p in passages:
            sid = p.get('source_id') or 0
            entry = agg.setdefault(sid, {
                'source_id': sid,
                'title': p.get('source_title') or '',
                'chunks_used': 0,
            })
            entry['chunks_used'] += 1
            if not entry['title']:
                entry['title'] = p.get('source_title') or ''
        return [v for v in agg.values() if v['source_id']]

    # ------------------------------------------------------------------
    @staticmethod
    def _render_prompt_block(passages: list[dict]) -> str:
        """Format passages as a citation-ready prompt block."""
        if not passages:
            return ''
        lines: list[str] = []
        for i, p in enumerate(passages, start=1):
            title = p.get('source_title') or 'Reference'
            lines.append(f'[{i}] ({title})\n{p.get("text", "")}\n')
        return '\n'.join(lines).strip()
