"""Mine exercises from indexed reference books and persist them as
``interactive_workbook`` materials on the plan.

Pipeline
--------
1. Walk all chunks belonging to the plan's indexed sources
   (``encoach.embedding`` rows with ``content_type='course_plan_source'``
   and ``entity_id=plan.id``).
2. Group chunks into N-chunk batches and ask the LLM to identify
   exercises in the text — gap-fills, multiple-choice items, matching
   pairs, word-reorder, transformation, short-answer.
3. Validate / clean each exercise with cheap heuristics. Drop anything
   without a usable answer key. De-dupe against previously-extracted
   stems (same plan).
4. For each accepted batch, attach the workbook to the most relevant
   week — picked by keyword overlap between the exercise stems and
   each week's ``unit + focus``. Falls back to the first week.
5. Persist as one ``encoach.course.plan.material`` per batch with
   ``material_type='interactive_workbook'`` and an
   ``extracted_from`` provenance record so the UI can show "Extracted
   from <book>".
"""

from __future__ import annotations

import json
import logging
import re

_logger = logging.getLogger(__name__)


# Per LLM call: how many ~2K-char chunks to feed in one go. 6 keeps the
# prompt under ~12K chars, well within a 16K context budget while still
# letting the model see enough surrounding context to identify rubrics.
CHUNKS_PER_BATCH = 6
# Hard cap on total batches per run — protects against runaway cost on a
# very large book. Operators can lift this via ``max_batches`` kwarg.
DEFAULT_MAX_BATCHES = 8

ALLOWED_TYPES = {
    'gap_fill', 'multiple_choice', 'match_pairs',
    'reorder_words', 'transformation', 'short_answer',
}

_EXTRACT_HINT = """
You are reading raw text extracted from a printed English exercise book.
Identify any drills, controlled-practice exercises, or workbook tasks in
the passage and extract them as machine-checkable JSON.

Return JSON with EXACTLY this shape:
{
  "exercises": [
    {"id":"e1", "type":"gap_fill",        "stem":"I ___ (be) a student.", "answer":"am", "alt":["am"], "hint":"present simple of 'be'"},
    {"id":"e2", "type":"multiple_choice", "stem":"...", "options":["A","B","C","D"], "answer":"B", "rationale":"..."},
    {"id":"e3", "type":"match_pairs",     "left":["dog","cat"], "right":["barks","meows"], "answer":[[0,0],[1,1]]},
    {"id":"e4", "type":"reorder_words",   "tokens":["I","am","a","student"], "answer":"I am a student"},
    {"id":"e5", "type":"transformation",  "from":"He plays football.", "instruction":"Make it negative.", "answer":"He doesn't play football.", "alt":["He does not play football."]},
    {"id":"e6", "type":"short_answer",    "stem":"What's your name?", "answer":"My name is …", "accepted":["My name is *","I'm *","I am *"]}
  ],
  "topic_keywords": ["routine","present simple", "..."],
  "page_hint": "if the source mentions 'p. 12' or 'Unit 1', echo it here"
}

Strict rules
------------
- Skip narrative prose, scope-and-sequence pages, answer keys.
- Skip exercises whose answer cannot be inferred from the source.
- Use ONLY the six types listed above; never invent a new type.
- Stems must be self-contained (no "look at exercise 3").
- Answers must be a string, a boolean, or a list of [int,int] pairs as
  shown — never a sentence like "answer key on p. 99".
- Output MUST be valid JSON. If the passage contains no usable
  exercises, return ``{"exercises": [], "topic_keywords": [], "page_hint": ""}``.
"""


def _normalize_stem(stem: str) -> str:
    """Lowercase + collapse whitespace — used for de-dup keys."""
    return re.sub(r'\s+', ' ', (stem or '').strip().lower())


def _is_clean_exercise(ex: dict) -> bool:
    """Cheap structural validation. Drops anything we can't grade."""
    t = (ex.get('type') or '').lower()
    if t not in ALLOWED_TYPES:
        return False
    if not (ex.get('stem') or ex.get('from') or ex.get('left') or ex.get('tokens')):
        return False
    ans = ex.get('answer')
    if ans is None or ans == '' or ans == []:
        return False
    if t == 'multiple_choice':
        opts = ex.get('options') or []
        if not isinstance(opts, list) or len(opts) < 2:
            return False
        if isinstance(ans, str) and ans not in opts:
            # Some authors use "B" as the answer; that's fine.
            if not (len(ans) == 1 and ans.isalpha()):
                return False
    if t == 'match_pairs':
        if not (isinstance(ex.get('left'), list)
                and isinstance(ex.get('right'), list)
                and isinstance(ans, list)):
            return False
        if not all(isinstance(p, list) and len(p) == 2 for p in ans):
            return False
    if t == 'reorder_words':
        toks = ex.get('tokens')
        if not (isinstance(toks, list) and len(toks) >= 2):
            return False
    return True


class ExerciseExtractor:
    """Walk the plan's indexed corpus and persist interactive workbooks."""

    CONTENT_TYPE = 'course_plan_source'

    def __init__(self, env, *, language: str = 'en'):
        self.env = env
        self.language = language
        try:
            from odoo.addons.encoach_ai.services.openai_service import (
                OpenAIService,
            )
        except ImportError:
            OpenAIService = None
        self._OpenAIService = OpenAIService
        self._ai = None  # lazy

    # ------------------------------------------------------------------
    def _service(self):
        if self._ai is not None:
            return self._ai
        if self._OpenAIService is None:
            raise RuntimeError('encoach_ai.OpenAIService not available')
        self._ai = self._OpenAIService(self.env, language=self.language)
        return self._ai

    # ------------------------------------------------------------------
    def _load_corpus(self, plan) -> list[dict]:
        """Return all indexed chunks for the plan, oldest-first."""
        Embedding = self.env['encoach.embedding'].sudo()
        rows = Embedding.search([
            ('content_type', '=', self.CONTENT_TYPE),
            ('entity_id', '=', plan.id),
        ], order='content_id asc, chunk_index asc')
        chunks = []
        for r in rows:
            metadata = {}
            try:
                metadata = json.loads(r.metadata_json or '{}')
            except (TypeError, ValueError):
                pass
            chunks.append({
                'source_id': int(metadata.get('source_id') or r.content_id),
                'title': metadata.get('title') or '',
                'chunk_index': r.chunk_index,
                'text': r.content_text or '',
            })
        return chunks

    # ------------------------------------------------------------------
    def _existing_stems(self, plan) -> set[str]:
        """Already-stored normalized stems → avoid duplicate workbooks."""
        Material = self.env['encoach.course.plan.material'].sudo()
        rows = Material.search([
            ('plan_id', '=', plan.id),
            ('material_type', '=', 'interactive_workbook'),
        ])
        seen: set[str] = set()
        for m in rows:
            try:
                body = json.loads(m.body_json or '{}')
            except (TypeError, ValueError):
                continue
            for ex in body.get('exercises') or []:
                key = _normalize_stem(ex.get('stem') or ex.get('from') or '')
                if key:
                    seen.add(key)
        return seen

    # ------------------------------------------------------------------
    @staticmethod
    def _pick_week(plan, topic_keywords: list[str], stems: list[str]):
        """Return the week most relevant to a batch's topic keywords.

        Falls back to week 1 when the plan has no overlap signal — that's
        better than dropping the workbook entirely.
        """
        weeks = list(plan.week_ids.sorted('week_number'))
        if not weeks:
            return None
        haystack = ' '.join(topic_keywords + stems).lower()
        if not haystack.strip():
            return weeks[0]

        def score(week):
            text = f'{week.unit or ""} {week.focus or ""}'.lower()
            terms = [t for t in re.split(r'\W+', text) if len(t) > 3]
            return sum(1 for t in terms if t in haystack)

        ranked = sorted(weeks, key=score, reverse=True)
        # If nothing scored at all, prefer week 1 for stability.
        return ranked[0] if score(ranked[0]) > 0 else weeks[0]

    # ------------------------------------------------------------------
    def _llm_extract_batch(self, chunks: list[dict]) -> dict:
        """Run one extraction LLM call against a batch of chunks."""
        ai = self._service()
        joined = '\n\n---\n\n'.join(
            f'[chunk {c["chunk_index"]} | source #{c["source_id"]}]\n{c["text"]}'
            for c in chunks
        )
        messages = [
            {
                'role': 'system',
                'content': (
                    'You extract original printed exercises from English '
                    'workbook scans. You return STRICT JSON. You never '
                    'invent exercises that are not on the page.'
                ),
            },
            {
                'role': 'user',
                'content': (
                    f'Extract the exercises from the passage below.\n\n'
                    f'{_EXTRACT_HINT}\n\nPassage:\n{joined}'
                ),
            },
        ]
        try:
            return ai.chat_json(
                messages,
                temperature=0.1,
                max_tokens=4000,
                action='course_plan.extract_exercises',
            ) or {}
        except Exception as exc:
            _logger.warning('Exercise extraction LLM call failed: %s', exc)
            return {}

    # ------------------------------------------------------------------
    def run(self, plan, *, max_batches: int = DEFAULT_MAX_BATCHES) -> dict:
        """Extract exercises for ``plan`` and persist workbook materials.

        :returns: ``{materials_created, exercises_total, batches_run, skipped}``.
        """
        chunks = self._load_corpus(plan)
        if not chunks:
            return {
                'materials_created': 0,
                'exercises_total': 0,
                'batches_run': 0,
                'skipped': 0,
                'reason': 'no indexed chunks for this plan',
            }

        seen_stems = self._existing_stems(plan)
        Material = self.env['encoach.course.plan.material'].sudo()
        materials_created = 0
        exercises_total = 0
        batches_run = 0
        skipped_dup = 0

        # Group chunks per source so a single workbook stays cohesive.
        by_source: dict[int, list[dict]] = {}
        for ch in chunks:
            by_source.setdefault(ch['source_id'], []).append(ch)

        for source_id, source_chunks in by_source.items():
            for i in range(0, len(source_chunks), CHUNKS_PER_BATCH):
                if batches_run >= max_batches:
                    break
                batch = source_chunks[i:i + CHUNKS_PER_BATCH]
                batches_run += 1
                payload = self._llm_extract_batch(batch)
                if not payload or 'error' in payload:
                    continue
                raw_exercises = payload.get('exercises') or []
                clean: list[dict] = []
                for ex in raw_exercises:
                    if not isinstance(ex, dict):
                        continue
                    if not _is_clean_exercise(ex):
                        continue
                    key = _normalize_stem(
                        ex.get('stem') or ex.get('from') or '',
                    )
                    if not key or key in seen_stems:
                        skipped_dup += 1
                        continue
                    seen_stems.add(key)
                    clean.append(ex)
                if not clean:
                    continue

                topic_keywords = [
                    str(t).strip() for t in (payload.get('topic_keywords') or [])
                    if str(t).strip()
                ]
                page_hint = (payload.get('page_hint') or '').strip()
                stems = [ex.get('stem') or ex.get('from') or '' for ex in clean]
                week = self._pick_week(plan, topic_keywords, stems)
                if week is None:
                    continue

                title = self._pick_title(batch, topic_keywords, week)
                summary = (
                    f'Extracted from {batch[0].get("title") or "uploaded book"}'
                    + (f' ({page_hint})' if page_hint else '')
                    + f' — {len(clean)} interactive exercise(s).'
                )
                body = {
                    'lesson_plan': {
                        'warmup': '',
                        'presentation': (
                            'Originally printed exercises mined from the '
                            'reference book. Use them as in-class drills '
                            'or homework.'
                        ),
                        'controlled_practice': '',
                        'freer_practice': '',
                        'exit_ticket': '',
                    },
                    'exercises': clean,
                }
                provenance = {
                    'source_id': int(source_id),
                    'source_title': batch[0].get('title') or '',
                    'page_hint': page_hint,
                    'topic_keywords': topic_keywords,
                    'chunk_indices': [c['chunk_index'] for c in batch],
                }
                Material.create({
                    'plan_id': plan.id,
                    'week_id': week.id,
                    'skill': 'integrated',
                    'material_type': 'interactive_workbook',
                    'title': title,
                    'summary': summary,
                    'body_json': json.dumps(body, ensure_ascii=False),
                    'body_text': '',
                    'extracted_from_json': json.dumps(
                        provenance, ensure_ascii=False,
                    ),
                })
                materials_created += 1
                exercises_total += len(clean)

            if batches_run >= max_batches:
                break

        return {
            'materials_created': materials_created,
            'exercises_total': exercises_total,
            'batches_run': batches_run,
            'skipped': skipped_dup,
        }

    # ------------------------------------------------------------------
    @staticmethod
    def _pick_title(batch, topic_keywords, week) -> str:
        """Compose a friendly material title for the extracted workbook."""
        topic = ''
        if topic_keywords:
            topic = topic_keywords[0]
        if not topic and week and (week.unit or week.focus):
            topic = (week.unit or week.focus or '').strip()
        if not topic:
            topic = 'Practice'
        return f'Workbook — {topic.title()}'
