src/summarization/note_generator.py

import json
import os
import re
from typing import Dict, List

try:
    from langdetect import detect as _detect_lang
except ImportError:
    _detect_lang = None

from pydantic import ValidationError

from ..utils.logger import setup_logger
from ..utils.model_loader import generate_text, INFERENCE_MODE, get_groq_client, _generate_text_local
from .segmenter import TranscriptSegmenter
from .schemas import SummarySchema, SegmentSchema

logger = setup_logger(__name__)

# Segment titles cycled for the timeline (Arabic, matches existing UX copy)
SEGMENT_TITLES = [
    "المقدمة والأفكار الأولى",
    "الشرح المتعمق والتحليل",
    "نقاط إضافية ومناقشة",
    "أمثلة وتفاصيل تطبيقية",
    "الشرح المتقدم",
    "ملاحظات ختامية",
    "الخلاصة والتوصيات النهائية",
]

WHITESPACE_RE = re.compile(r"\s+")

# ── Strict language-matching instruction for language consistency ─────────────
_LANG_MATCH_INSTRUCTION = (
    "Identify the primary language of the input text. You MUST generate your entire response strictly in that exact same language. DO NOT mix languages under any circumstances."
)




class NoteGenerator:
    """AI Summarization Engine using Qwen2.5-0.5B-Instruct (local, CPU).

    Pipeline (map-reduce):
      1. MAP    -> split transcript into word-bounded chunks (TranscriptSegmenter)
                   and summarize each chunk individually with the local model.
      2. REDUCE -> summarize the concatenation of all chunk-summaries to produce
                   one coherent overall summary.
      3. VALIDATE -> assemble everything into SummarySchema via model_validate()
                     so malformed/missing fields are caught before reaching
                     the formatting layer.
    """

    def __init__(self, chunk_size_words: int = 200):
        self.segmenter = TranscriptSegmenter(max_segment_words=chunk_size_words)
        self.chunk_size_words = chunk_size_words
        if INFERENCE_MODE == "groq":
            logger.info("📝 NoteGenerator ready (Groq Cloud API, llama-3.3-70b-versatile).")
        else:
            logger.info("📝 NoteGenerator ready (Qwen2.5-0.5B-Instruct, CPU).")

    @staticmethod
    def _detect_language(text: str) -> str:
        """Detect whether the transcript is Arabic or English.

        Uses langdetect on the first ~500 chars for speed.
        Returns 'Arabic' or 'English' (defaults to 'Arabic' on failure).
        """
        if not text or not text.strip():
            return "Arabic"

        if _detect_lang is None:
            # langdetect not installed — simple heuristic: check Arabic char ratio
            arabic_chars = sum(1 for c in text[:500] if '\u0600' <= c <= '\u06FF')
            return "Arabic" if arabic_chars > len(text[:500]) * 0.2 else "English"

        try:
            lang_code = _detect_lang(text[:500])
            if lang_code == "ar":
                return "Arabic"
            return "English"
        except Exception:
            return "Arabic"

    # ── Core summarization via local Qwen model ──────────────────────────

    def _run_model_summarization(
        self,
        text: str,
        max_new_tokens: int = 100,
    ) -> str:
        """Run the local Qwen model on a single piece of text and return its summary."""
        try:
            clean_text = WHITESPACE_RE.sub(" ", text.strip())
            if not clean_text:
                logger.warning("⚠️ _run_model_summarization received empty text after cleaning.")
                return ""

            logger.info(f"🔎 Input to model (len={len(clean_text)} chars): {clean_text[:120]!r}...")

            messages = [
                {
                    "role": "system",
                    "content": (
                        "You are a professional summarizer.\n"
                        f"{_LANG_MATCH_INSTRUCTION}\n\n"
                        "Summarize the following text concisely. "
                        "Focus on the main ideas and key points. "
                        "Keep the summary short (3-5 sentences maximum). "
                        "Output ONLY the summary, nothing else."
                    ),
                },
                {"role": "user", "content": clean_text[:2000]},
            ]

            if INFERENCE_MODE == "groq":
                logger.info("🟢 Generating summary chunk via Groq API (llama-3.3-70b-versatile)...")
                groq_client = get_groq_client()
                chat_completion = groq_client.chat.completions.create(
                    model="llama-3.3-70b-versatile",
                    messages=messages,
                    max_tokens=max_new_tokens,
                    temperature=0.0,
                )
                result = chat_completion.choices[0].message.content or ""
            else:
                logger.info("🤖 Generating summary chunk via local Qwen pipeline...")
                result = _generate_text_local(messages, max_new_tokens=max_new_tokens)

            logger.info(f"🔎 Model output (len={len(result)} chars): {result[:200]!r}")

            if not result or len(result.strip()) < 5:
                logger.warning(
                    f"⚠️ Model returned a suspiciously short/empty summary: {result!r}. "
                    f"Falling back to placeholder text."
                )

            return result
        except Exception as e:
            logger.error(f"❌ Error during model generation: {str(e)}", exc_info=True)
            return ""

    def _map_chunks(self, transcript_text: str, video_title: str) -> List[Dict]:
        """MAP step: split transcript into chunks and summarize each one."""
        chunks = self.segmenter.segment_text_by_words(
            transcript_text, chunk_size=self.chunk_size_words
        )

        if not chunks:
            chunks = [f"هذا الفيديو يتحدث عن {video_title}."]

        segments_list = []
        for i, chunk in enumerate(chunks):
            chunk_summary = self._run_model_summarization(chunk)
            if not chunk_summary or len(chunk_summary.strip()) < 5:
                chunk_summary = f"شرح وتحليل للأفكار الواردة في الجزء رقم {i + 1} من الفيديو."

            title = SEGMENT_TITLES[i % len(SEGMENT_TITLES)]
            segments_list.append(
                {
                    "title": title,
                    "summary": chunk_summary,
                    "key_insight": f"الاستنتاج الجوهري من هذا الجزء يدعم فهم سياق {video_title}.",
                    "why_it_matters": "يساعد هذا التقسيم في استيعاب النقاط الرئيسية بشكل منظم ومتسلسل.",
                }
            )

        # Keep the timeline within the schema's allowed range (3-7 segments).
        if len(segments_list) > 7:
            segments_list = self._merge_to_limit(segments_list, limit=7)
        elif len(segments_list) < 3:
            segments_list = self._pad_to_minimum(segments_list, video_title, minimum=3)

        return segments_list

    def _merge_to_limit(self, segments_list: List[Dict], limit: int) -> List[Dict]:
        """Merge extra segments into the last allowed slot to respect max_length."""
        kept = segments_list[: limit - 1]
        overflow = segments_list[limit - 1:]
        merged_summary = " ".join(s["summary"] for s in overflow)
        kept.append(
            {
                "title": SEGMENT_TITLES[-1],
                "summary": merged_summary,
                "key_insight": overflow[-1]["key_insight"],
                "why_it_matters": overflow[-1]["why_it_matters"],
            }
        )
        return kept

    def _pad_to_minimum(
        self, segments_list: List[Dict], video_title: str, minimum: int
    ) -> List[Dict]:
        """Pad segments (very short transcripts) to satisfy min_length=3."""
        while len(segments_list) < minimum:
            idx = len(segments_list)
            segments_list.append(
                {
                    "title": SEGMENT_TITLES[idx % len(SEGMENT_TITLES)],
                    "summary": f"تتمة الشرح والنقاط المرتبطة بـ {video_title}.",
                    "key_insight": "نقطة تكميلية تدعم الفهم العام للموضوع.",
                    "why_it_matters": "تضيف سياقًا إضافيًا لاستيعاب الفيديو بالكامل.",
                }
            )
        return segments_list

    # ── Valid categories for the reduce prompt ────────────────────────────
    VALID_CATEGORIES = [
        "Technology & AI",
        "Business & Finance",
        "Education",
        "Science",
        "Productivity & Self-Growth",
        "Health & Wellness",
        "Sports & Fitness",
        "Entertainment",
        "History",
        "Philosophy",
        "Arts & Culture",
    ]

    def _parse_reduce_json(self, raw_text: str, video_title: str) -> Dict:
        """Safely parse the strict JSON returned by the reduce LLM call.

        Returns a dict with keys: markdown_summary, key_points, category, language.
        Falls back gracefully if parsing fails.
        """
        text = raw_text.strip()
        
        # Use regex to find the outermost JSON object to bypass any conversational filler
        json_match = re.search(r"(\{.*\})", text, re.DOTALL)
        if json_match:
            text = json_match.group(1)

        try:
            data = json.loads(text)
        except json.JSONDecodeError as e:
            logger.warning(f"⚠️ Failed to parse reduce JSON: {e}. Raw text preview: {raw_text[:200]!r}")
            return {
                "markdown_summary": raw_text,
                "key_points": [],
                "category": "Education",
                "language": "ar",
            }

        # Validate and sanitize each field
        markdown_summary = data.get("markdown_summary", "").strip()
        if not markdown_summary:
            markdown_summary = raw_text

        key_points = data.get("key_points", [])
        if not isinstance(key_points, list):
            key_points = []
        key_points = [str(p).strip() for p in key_points if str(p).strip()][:5]

        category = data.get("category", "").strip()
        if category not in self.VALID_CATEGORIES:
            # Attempt fuzzy match
            category_lower = category.lower()
            matched = False
            for valid_cat in self.VALID_CATEGORIES:
                if valid_cat.lower() in category_lower or category_lower in valid_cat.lower():
                    category = valid_cat
                    matched = True
                    break
            if not matched:
                category = "Education"

        language = data.get("language", "ar").strip().lower()
        if language not in ("ar", "en"):
            language = "ar"

        return {
            "markdown_summary": markdown_summary,
            "key_points": key_points,
            "category": category,
            "language": language,
        }

    def _reduce_summary(self, segments_list: List[Dict], video_title: str) -> Dict:
        """REDUCE step: combine chunk-summaries into a strict JSON with
        markdown_summary, key_points, category, and language.

        Returns a dict (parsed JSON), NOT a raw string.
        """
        combined_text = " ".join(seg["summary"] for seg in segments_list)
        clean_combined = WHITESPACE_RE.sub(" ", combined_text.strip())[:3000]

        if not clean_combined:
            return {
                "markdown_summary": f"Summary of: {video_title}.",
                "key_points": [],
                "category": "Education",
                "language": "ar",
            }

        categories_str = ", ".join(f'"{c}"' for c in self.VALID_CATEGORIES)

        if INFERENCE_MODE == "groq":
            messages = [
                {
                    "role": "system",
                    "content": (
                        "You are a professional summarizer compiling a final high-density, detailed study guide of a video.\n\n"
                        "STRICT RULES — follow every instruction exactly:\n\n"
                        "1. HIGH-DENSITY, DEEP SUMMARY REQUIREMENT:\n"
                        "   - Short, generic, or vague sentences are strictly forbidden. The summary must be detailed, rich, and highly informative.\n"
                        "   - Extract actual specifications, names, data, key parameters, technical metrics, and specific details mentioned in the text.\n"
                        "   - For example, if it's a car review/experience, specify the engine model, horsepower, layout, driving sensations, torque, and other specs instead of saying 'it was a nice experience'.\n\n"
                        "2. LANGUAGE:\n"
                        "   - Detect the primary language of the text. If Arabic, write EVERYTHING (headers, overview, questions, answers, key points) in Arabic. If English, write EVERYTHING in English. NO MIXING. If the video language is Arabic, ensure NO English words leak into questions, answers, or headers.\n\n"
                        "3. OUTPUT FORMAT:\n"
                        "   - You MUST return a single valid JSON object with this exact schema — no conversational filler, ONLY the raw JSON:\n"
                        '{\n'
                        '  "language": "ar" or "en",\n'
                        '  "category": "one of the 11 categories listed below",\n'
                        '  "markdown_summary": "the formatted markdown string",\n'
                        '  "key_points": ["point 1", "point 2", "point 3", "point 4", "point 5"]\n'
                        '}\n\n'
                        '4. MARKDOWN_SUMMARY TYPOGRAPHY & LAYOUT:\n'
                        '   - Start with: ## 📋 الملخص العام  (or ## 📋 General Summary for English)\n'
                        '   - Add a double newline (\\n\\n)\n'
                        '   - Write a rich, comprehensive narrative paragraph (at least 4-6 sentences) mapping out the core thesis, background, and conclusion of the video. It must be highly detailed and dense.\n'
                        '   - Add a separator: \\n\\n---\\n\\n\n'
                        '   - Then write: ## ❓ أبرز الأسئلة والأجوبة  (or ## ❓ Key Questions & Answers for English)\n'
                        '   - Add a double newline (\\n\\n)\n'
                        '   - Generate EXACTLY 7 Q&A pairs — no more, no fewer. Each question must be unique, insightful, and cover a different aspect of the video content. '
                        'Each answer must be detailed, structured, and data-rich (at least 2-3 sentences with specific facts, names, numbers, or examples from the content). '
                        'Format each Q&A block EXACTLY as follows, using a clean double newline between the question and answer, and a horizontal rule as a separator:\n\n'
                        '     ❓ **[السؤال هنا بخط عريض]؟**\n\n'
                        '     💡 [الإجابة المفصلة العميقة والمليئة بالبيانات هنا — يجب أن تكون 2-3 جمل على الأقل مع تفاصيل محددة]\n\n'
                        '     ---\n\n'
                        '     For English:\n'
                        '     ❓ **[Question text here in bold]?**\n\n'
                        '     💡 [Detailed, structured answer here — must be at least 2-3 sentences with specific details]\n\n'
                        '     ---\n\n'
                        '   - You MUST generate EXACTLY 7 Q&A pairs. Count them carefully before outputting.\n'
                        '   - QUESTIONS AND ANSWERS MUST NEVER BE ON THE SAME LINE. ALWAYS use a double newline (\\n\\n) to separate them.\n'
                        '   - Place a horizontal line (---) between each Q&A block. Do NOT place a horizontal line after the final Q&A block.\n'
                        '   - Use relevant modern emojis sparingly.\n'
                        '   - DO NOT include the key points, bullet lists, or any other sections in markdown_summary.\n\n'
                        '5. KEY_POINTS:\n'
                        '   - Exactly 5 concise, factual, and data-driven strings in a JSON array. These must be concrete insights from the content, NOT generic text. Do NOT repeat these key points inside markdown_summary.\n\n'
                        '6. CATEGORY CLASSIFICATION (ZERO WRONG DEFAULTING BIAS):\n'
                        f'   - The `"category"` value must strictly be one of: [{categories_str}].\n'
                        '   - Perform a strict semantic reasoning step before selecting: Analyze the core human domain of the video.\n'
                        '   - For example, if the video is about reviewing cars, driving experiences, car history, racing, or vehicle design, it belongs strictly under "Entertainment" or "Sports & Fitness" or "Science" depending on context, NOT "Technology & AI" unless it specifically focuses on software engineering, computer hardware, or artificial intelligence algorithms.\n'
                        '   - Do not default to "Technology & AI" or "Education" unless the text explicitly falls under those domains.\n\n'
                        "7. Return ONLY raw JSON starting with '{' and ending with '}'."
                    ),
                },
                {"role": "user", "content": clean_combined},
            ]
            logger.info("🟢 Reducing summaries via Groq API (strict JSON mode)...")
            groq_client = get_groq_client()
            chat_completion = groq_client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=messages,
                max_tokens=2500,
                temperature=0.0,
                response_format={"type": "json_object"},
            )
            raw_response = chat_completion.choices[0].message.content or ""
            logger.info(f"🔎 Reduce raw response (len={len(raw_response)}): {raw_response[:300]!r}")
            return self._parse_reduce_json(raw_response, video_title)
        else:
            # Local Qwen mode — keep simple text-based reduce (no JSON)
            messages = [
                {
                    "role": "system",
                    "content": (
                        "You are a professional summarizer compiling a final overview.\n"
                        f"{_LANG_MATCH_INSTRUCTION}\n\n"
                        "Below are partial summaries from different sections of a video. "
                        "Combine them into ONE coherent overall summary (5-8 sentences). "
                        "Then generate EXACTLY 7 detailed Q&A pairs covering different aspects of the content. "
                        "Each answer must be structured and detailed (at least 2-3 sentences). "
                        "Output ONLY the final summary followed by the 7 Q&A pairs, nothing else."
                    ),
                },
                {"role": "user", "content": clean_combined},
            ]
            logger.info("🤖 Reducing summaries via local Qwen pipeline...")
            overall = _generate_text_local(messages, max_new_tokens=500)

            if not overall or len(overall.strip()) < 5:
                overall = f"استعراض شامل ومناقشة تفصيلية لموضوع: {video_title}."

            return {
                "markdown_summary": overall,
                "key_points": [],
                "category": "Education",
                "language": "ar",
            }

    def generateSummary(self, transcript_text: str, video_title: str) -> Dict:
        """Generates a structured AI summary, validated against SummarySchema."""
        if INFERENCE_MODE == "groq":
            logger.info("📝 AI summary generation triggered (map-reduce pipeline, Groq Cloud API).")
        else:
            logger.info("📝 AI summary generation triggered (map-reduce pipeline, Qwen local).")
        logger.info(
            f"🔎 Received video_title={video_title!r}, "
            f"transcript_text length={len(transcript_text) if transcript_text else 0} chars, "
            f"preview={transcript_text[:150] if transcript_text else None!r}"
        )

        if not transcript_text or len(transcript_text.strip()) < 10:
            logger.warning(
                "⚠️ transcript_text is empty or too short (<10 chars). "
                "Using a generic placeholder transcript instead of real content."
            )
            transcript_text = (
                f"هذا الفيديو يتحدث عن {video_title} وشرح تفصيلي للمحتوى المكتوب في العنوان."
            )

        # 1. MAP — summarize each chunk independently
        segments_list = self._map_chunks(transcript_text, video_title)

        # 2. REDUCE — summarize the chunk-summaries into one overall summary
        overall_summary = self._reduce_summary(segments_list, video_title)

        # 3. Detect language from the transcript (used by UI and metadata)
        detected_language = self._detect_language(transcript_text)
        
        # Override detected language based on the REDUCE step if available
        llm_lang = overall_summary.get("language")
        if llm_lang:
            detected_language = "Arabic" if llm_lang == "ar" else "English"

        if INFERENCE_MODE == "groq":
            conclusion_text = (
                f"تم التلخيص بنجاح باستخدام نموذج Groq (llama-3.3-70b) لـ {video_title}."
                if detected_language == "Arabic"
                else f"Summary generated successfully via Groq (llama-3.3-70b) for: {video_title}."
            )
            topics_list = (
                ["تلخيص تلقائي", "ذكاء اصطناعي سحابي", "Groq API", "Llama3.3"]
                if detected_language == "Arabic"
                else ["Auto-summary", "Cloud AI", "Groq API", "Llama3.3"]
            )
        else:
            conclusion_text = (
                f"تم التلخيص بنجاح باستخدام نموذج Qwen المحلي لـ {video_title}."
                if detected_language == "Arabic"
                else f"Summary generated successfully for: {video_title}."
            )
            topics_list = (
                ["تلخيص تلقائي", "ذكاء اصطناعي محلي", "Qwen2.5"]
                if detected_language == "Arabic"
                else ["Auto-summary", "Local AI", "Qwen2.5"]
            )

        candidate = {
            "title": video_title,
            "detected_language": detected_language,
            "summary": overall_summary.get("markdown_summary", ""),
            "segments": segments_list,
            "suggested_category": overall_summary.get("category", "Education"),
            "conclusion": conclusion_text,
            "topics": topics_list,
            "key_points": overall_summary.get("key_points", []),
        }

        # 3. VALIDATE — enforce the schema contract before returning
        try:
            validated = SummarySchema.model_validate(candidate)
            return validated.model_dump()
        except ValidationError as e:
            logger.error(f"SummarySchema validation failed: {e}")
            return self._fallback_summary(video_title)

    def _fallback_summary(self, video_title: str) -> Dict:
        """Safe, schema-valid fallback used only if validation fails."""
        if INFERENCE_MODE == "groq":
            conclusion_text = f"تم التلخيص بنجاح باستخدام نموذج Groq (llama-3.3-70b) لـ {video_title}."
            topics_list = ["تلخيص تلقائي", "ذكاء اصطناعي سحابي", "Groq API"]
        else:
            conclusion_text = f"تم التلخيص بنجاح باستخدام نموذج Qwen المحلي لـ {video_title}."
            topics_list = ["تلخيص تلقائي", "ذكاء اصطناعي محلي", "Qwen2.5"]

        fallback = SummarySchema(
            title=video_title,
            detected_language="Arabic",
            summary=f"استعراض شامل ومناقشة تفصيلية لموضوع: {video_title}.",
            segments=[
                SegmentSchema(
                    title=SEGMENT_TITLES[i],
                    summary=f"شرح وتحليل للأفكار الواردة في الجزء رقم {i + 1} من الفيديو.",
                    key_insight=f"الاستنتاج الجوهري من هذا الجزء يدعم فهم سياق {video_title}.",
                    why_it_matters="يساعد هذا التقسيم في استيعاب النقاط الرئيسية بشكل منظم ومتسلسل.",
                )
                for i in range(3)
            ],
            conclusion=conclusion_text,
            topics=topics_list,
        )
        return fallback.model_dump()

    def format_notes_to_markdown(self, json_notes: Dict) -> str:
        """Convert JSON notes to clean Markdown."""
        if INFERENCE_MODE == "groq":
            return json_notes.get("summary", "")

        lines = []
        lines.append("## 📋 الملخص العام")
        lines.append("")
        lines.append(json_notes.get("summary", ""))
        lines.append("")
        lines.append("---")
        lines.append("")
        lines.append("## 🕐 التسلسل الزمني للأفكار")
        lines.append("")

        for i, seg in enumerate(json_notes.get("segments", []), start=1):
            lines.append(f"### {i}. {seg['title']}")
            lines.append("")
            lines.append(seg["summary"])
            lines.append("")
            lines.append(f"> **💎 أهم نقطة مستفادة:** {seg['key_insight']}")
            lines.append("")
            lines.append(f"> **لماذا يهم هذا الجزء؟** {seg['why_it_matters']}")
            lines.append("")

        lines.append("---")
        lines.append("")
        lines.append("## 🔖 الخلاصة")
        lines.append("")
        lines.append(f"> {json_notes.get('conclusion', '')}")

        return "\n".join(lines)

    def format_final_notes(
        self,
        notes: str,
        video_title: str,
        video_url: str,
        duration: int,
        detected_language: str = "English",
    ) -> str:
        """Wrap the formatted Markdown body with Source + Duration header."""
        if duration and duration > 0:
            minutes = duration // 60
            secs = duration % 60
            duration_str = f"{minutes:02d}:{secs:02d} دقيقة"
        else:
            duration_str = "غير محدد"

        header = (
            f"# {video_title}\n\n"
            f"---\n\n"
            f"> **المصدر:** {video_url}  \n"
            f"> **المدة:** {duration_str}\n\n"
            f"---\n\n"
        )
        return header + notes

    def chat_with_note(
        self, note_content: str, question: str, history: list[dict] | None = None
    ) -> str:
        """Answer a question about a note using the local Qwen model.

        The model is instructed to ground its answers solely in the provided
        note content to avoid hallucination.
        """
        try:
            messages = [
                {
                    "role": "system",
                    "content": (
                        "You are a helpful study assistant. "
                        "Answer the user's question based ONLY on the note content provided below. "
                        "If the answer is not in the note, say so honestly. "
                        "Reply in the same language the user uses.\n\n"
                        f"--- NOTE CONTENT ---\n{note_content[:4000]}\n--- END NOTE ---"
                    ),
                },
            ]

            # Include conversation history if available
            if history:
                for msg in history[-6:]:  # Keep last 6 turns to fit context
                    role = msg.get("role", "user")
                    content = msg.get("content", "")
                    if role in ("user", "assistant") and content:
                        messages.append({"role": role, "content": content})

            messages.append({"role": "user", "content": question})

            if INFERENCE_MODE == "groq":
                logger.info("🟢 Answering chat question via Groq API (llama-3.3-70b-versatile)...")
                groq_client = get_groq_client()
                chat_completion = groq_client.chat.completions.create(
                    model="llama-3.3-70b-versatile",
                    messages=messages,
                    max_tokens=300,
                    temperature=0.0,
                )
                answer = chat_completion.choices[0].message.content or ""
            else:
                logger.info("🤖 Answering chat question via local Qwen pipeline...")
                answer = _generate_text_local(messages, max_new_tokens=300)

            if not answer or len(answer.strip()) < 3:
                return "عذرًا، لم أتمكن من توليد إجابة. يرجى إعادة صياغة السؤال."

            return answer

        except Exception as e:
            logger.error(f"❌ Chat error: {e}", exc_info=True)
            return "حدث خطأ أثناء معالجة سؤالك. يرجى المحاولة مرة أخرى."