"""
summarizer.py
-------------
Two-stage summarization: LED (BigBird) drafts, FLAN-T5 refines.
Keyword extraction via KeyBERT with a frequency-based fallback.

All model access goes through the registry — no imports of transformers/torch here.
Prompt templates are a table, not ad-hoc f-strings buried in conditionals.
"""
from __future__ import annotations

import logging
import re
from collections import Counter
from typing import List

from .model_registry import ModelKey, get_model
from .schemas import SummarizationRequest, SummarizationResult, SummarizationTask

logger = logging.getLogger(__name__)


# ── Prompt tables ─────────────────────────────────────────────────────────────
# Keyed by task. max_input_chars: how much source text to feed LED.
# max_output_tokens: generation ceiling for LED.

_LED_PROMPTS: dict[SummarizationTask, str] = {
    SummarizationTask.ABSTRACT: (
        "Summarize the following content into a 100-150 word SEO-optimized abstract:\n\n{text}"
    ),
    SummarizationTask.TITLE: (
        "Generate a concise, SEO-optimized title (under 70 characters) "
        "for the following content:\n\n{text}"
    ),
    SummarizationTask.CAPTION: (
        "Generate a brief, SEO-optimized video caption (1-2 sentences) "
        "for the following content:\n\n{text}"
    ),
    SummarizationTask.DESCRIPTION: (
        "Generate a detailed, SEO-optimized video description (200-300 words) "
        "covering the main topics, tone, and key takeaways:\n\n{text}"
    ),
}

_LED_MAX_OUTPUT: dict[SummarizationTask, int] = {
    SummarizationTask.ABSTRACT:    300,
    SummarizationTask.TITLE:        80,
    SummarizationTask.CAPTION:     150,
    SummarizationTask.DESCRIPTION: 500,
}

_LED_MAX_INPUT_CHARS: dict[SummarizationTask, int] = {
    SummarizationTask.ABSTRACT:    4000,
    SummarizationTask.TITLE:       1000,
    SummarizationTask.CAPTION:     1000,
    SummarizationTask.DESCRIPTION: 4000,
}

_FLAN_REFINE_PROMPT = (
    "You are an SEO content specialist. Refine the following draft {task} to be "
    "more engaging, keyword-rich, and publication-ready. "
    "Output only the refined text, no preamble.\n\n"
    "Draft:\n{draft}\n\n"
    "Original context (excerpt):\n{context}"
)

# Stopwords for the frequency fallback keyword extractor
_STOPWORDS = frozenset({
    "this", "that", "with", "from", "have", "been", "were", "they",
    "their", "there", "what", "when", "where", "which", "about",
    "would", "could", "should", "will", "just", "also", "then",
    "than", "into", "over", "after", "before", "more", "some",
})


# ── Stage helpers ─────────────────────────────────────────────────────────────

def _run_led(text: str, task: SummarizationTask) -> str:
    """Draft a summary using the LED model. Returns empty string on failure."""
    try:
        import torch
        bundle    = get_model(ModelKey.LED)
        tokenizer = bundle["tokenizer"]
        model     = bundle["model"]
        device    = "cuda" if torch.cuda.is_available() else "cpu"
        model     = model.to(device)

        max_chars = _LED_MAX_INPUT_CHARS[task]
        prompt    = _LED_PROMPTS[task].format(text=text[:max_chars])
        max_out   = _LED_MAX_OUTPUT[task]

        inputs = tokenizer(
            prompt, return_tensors="pt", truncation=True, max_length=1024
        ).to(device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=max_out,
                num_beams=5,
                early_stopping=True,
            )
        return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    except Exception as exc:
        logger.error("[Summarizer] LED stage failed: %s", exc)
        return ""


def _run_flan(
    draft:   str,
    task:    SummarizationTask,
    context: str,
    req:     SummarizationRequest,
) -> str:
    """Refine the LED draft using FLAN-T5. Falls back to the draft on failure."""
    try:
        flan_pipeline = get_model(ModelKey.FLAN)
        prompt = _FLAN_REFINE_PROMPT.format(
            task=task.value,
            draft=draft,
            context=context[:500],
        )
        result = flan_pipeline(
            prompt,
            min_length=req.min_length,
            max_length=req.max_length,
            do_sample=False,
            num_return_sequences=1,
        )
        if isinstance(result, list) and result and "generated_text" in result[0]:
            return result[0]["generated_text"].strip()
        return draft

    except Exception as exc:
        logger.error("[Summarizer] FLAN stage failed: %s", exc)
        return draft


def _extract_keywords(text: str, count: int) -> List[str]:
    """
    Extract keywords via KeyBERT.
    Falls back to naive word-frequency extraction if KeyBERT is unavailable.
    """
    try:
        kb = get_model(ModelKey.KEYBERT)
        pairs = kb.extract_keywords(
            text,
            keyphrase_ngram_range=(1, 2),
            stop_words="english",
            top_n=count,
        )
        return [kw for kw, _ in pairs]

    except Exception as exc:
        logger.warning("[Summarizer] KeyBERT failed (%s); using frequency fallback.", exc)
        words    = re.findall(r"\b[a-zA-Z]{4,}\b", text.lower())
        filtered = [w for w in words if w not in _STOPWORDS]
        return [w for w, _ in Counter(filtered).most_common(count)]


# ── Public API ────────────────────────────────────────────────────────────────

def summarize(req: SummarizationRequest) -> SummarizationResult:
    """
    Run the two-stage summarization pipeline.
    Accepts a SummarizationRequest, returns a SummarizationResult.
    Pure function over schemas — safe to call from a queue consumer.
    """
    text = req.transcription.text
    task = req.task

    # Stage A: LED draft
    led_draft = _run_led(text, task)
    if not led_draft:
        # Degrade gracefully — pass truncated source to FLAN directly
        logger.warning("[Summarizer] LED returned empty; degrading to source truncation.")
        led_draft = text[: _LED_MAX_INPUT_CHARS[task]]

    # Stage B: FLAN refinement
    flan_refined = _run_flan(led_draft, task, text, req)

    # Keyword extraction over full source text (not just the summary)
    keywords = _extract_keywords(text, req.keyword_count)

    return SummarizationResult(
        video_id=req.transcription.video_id,
        raw_summary=led_draft,
        refined_summary=flan_refined,
        keywords=keywords,
        task=task,
    )
