src/analyzer/streaming_summarizer.py

"""
streaming_summarizer.py — Advanced streaming + parallel grant summarization

Implements:
1. True parallelization with configurable concurrency
2. OpenAI streaming (stream=True) for real-time token delivery
3. Per-grant streaming with immediate feedback
4. Graceful fallback for Gradio (not async context)
"""
import asyncio
import logging
from typing import Any, Dict, List, Optional, AsyncGenerator

from .llm_client import LLMClient
from .summarizer_optimized import SummaryCache, extract_minimal_context

logger = logging.getLogger(__name__)


async def summarize_grant_streaming(
    grant: Dict[str, Any],
    client: LLMClient,
    cache: SummaryCache,
    past_winners: Optional[List[Dict[str, Any]]] = None,
) -> AsyncGenerator[str, None]:
    """
    Stream a single grant summary token-by-token.

    Yields tokens as they arrive from OpenAI.
    """
    grant_id = grant.get("id") or grant.get("title") or "unknown"
    title = grant.get("title") or grant.get("name") or "(untitled)"

    # Check cache first
    cached_summary = cache.get(grant)
    if cached_summary:
        logger.info("📦 Cache HIT for %s", grant_id)
        # Stream cached content quickly
        for token in cached_summary.split():
            yield token + " "
        return

    # Build context
    context = extract_minimal_context(grant, past_winners)

    # Stream from OpenAI
    from .prompt_templates import build_prompt

    payload = build_prompt("openai", context)
    full_response = ""

    try:
        # Use stream=True to get token-by-token delivery
        stream = client.chat(payload["messages"], max_tokens=1200, temperature=0.25, stream=True)

        for token in stream:
            full_response += token
            yield token

        # Cache the full response
        cache.set(grant, full_response)
        logger.info("✅ Cached summary for %s", grant_id)

    except Exception as e:
        logger.error("❌ Stream failed for %s: %s", grant_id, e)
        error_msg = f"Error generating summary: {str(e)[:100]}"
        cache.set(grant, error_msg)
        yield error_msg


async def summarize_grants_parallel_streaming(
    current: List[Dict[str, Any]],
    past_winners: Optional[List[Dict[str, Any]]] = None,
    *,
    limit: Optional[int] = None,
    client: Optional[LLMClient] = None,
    cache: Optional[SummaryCache] = None,
    batch_size: int = 5,
    max_concurrent: int = 3,
) -> AsyncGenerator[Dict[str, Any], None]:
    """
    Parallel + Streaming: Process multiple grants concurrently.

    Each grant streams its own summary as it's being generated.
    Results yielded as they complete (not in order).

    Args:
        max_concurrent: Max grants to process simultaneously (default: 3)
    """
    client = client or LLMClient({})
    cache = cache or SummaryCache(ttl_seconds=3600)

    items = current[: limit or len(current)]

    if not items:
        return

    # Create a semaphore to limit concurrent tasks
    semaphore = asyncio.Semaphore(max_concurrent)

    async def bounded_summarize(grant, index):
        """Summarize with concurrency limit."""
        async with semaphore:
            grant_id = grant.get("id") or grant.get("title") or "unknown"
            title = grant.get("title") or grant.get("name") or "(untitled)"

            full_summary = ""

            try:
                logger.info(f"[{index+1}/{len(items)}] Summarizing: {title[:50]}")

                async for token in summarize_grant_streaming(grant, client, cache, past_winners):
                    full_summary += token

                return {
                    "grant_id": grant_id,
                    "title": title,
                    "summary_md": full_summary,
                    "index": index,
                }
            except Exception as e:
                logger.error(f"Failed to summarize {grant_id}: {e}")
                return {
                    "grant_id": grant_id,
                    "title": title,
                    "summary_md": f"Error: {str(e)[:100]}",
                    "index": index,
                }

    # Create all tasks
    tasks = [bounded_summarize(grant, i) for i, grant in enumerate(items)]

    # Yield results as they complete (using as_completed)
    for coro in asyncio.as_completed(tasks):
        result = await coro
        yield result

    logger.info("✅ Completed all %d grants", len(items))


async def summarize_grants_batch_parallel_streaming(
    current: List[Dict[str, Any]],
    past_winners: Optional[List[Dict[str, Any]]] = None,
    *,
    limit: Optional[int] = None,
    client: Optional[LLMClient] = None,
    cache: Optional[SummaryCache] = None,
    batch_size: int = 5,
    max_concurrent_batches: int = 2,
) -> AsyncGenerator[Dict[str, Any], None]:
    """
    Optimized: Process batches in parallel, stream batch results.

    5 grants per batch → reduced API calls
    Multiple batches in parallel → maximum throughput
    Results stream as soon as batch completes

    This is the recommended approach for 30+ grants.
    """
    from .summarizer_optimized import _summarize_batch_async

    client = client or LLMClient({})
    cache = cache or SummaryCache(ttl_seconds=3600)

    items = current[: limit or len(current)]

    if not items:
        return

    # Build contexts
    contexts = [extract_minimal_context(g, past_winners) for g in items]

    # Create batches
    batches = [
        (items[i:i+batch_size], contexts[i:i+batch_size])
        for i in range(0, len(items), batch_size)
    ]

    # Create a semaphore for concurrent batch processing
    semaphore = asyncio.Semaphore(max_concurrent_batches)

    async def process_batch(batch_items, batch_contexts, batch_idx):
        """Process a batch with concurrency limit."""
        async with semaphore:
            try:
                logger.info(f"Processing batch {batch_idx+1}/{len(batches)} ({len(batch_items)} grants)")
                batch_results = await _summarize_batch_async(batch_items, batch_contexts, client, cache)

                for result in batch_results:
                    yield result

            except Exception as e:
                logger.error(f"Batch {batch_idx} failed: {e}")
                for grant in batch_items:
                    yield {
                        "grant_id": grant.get("id") or "unknown",
                        "title": grant.get("title") or "(untitled)",
                        "summary_md": f"Batch error: {str(e)[:100]}",
                    }

    # Process batches concurrently
    tasks = [
        process_batch(batch_items, batch_contexts, i)
        for i, (batch_items, batch_contexts) in enumerate(batches)
    ]

    # Yield from all tasks as they complete
    for task in asyncio.as_completed(tasks):
        async for result in task:
            yield result

    logger.info("✅ All batches processed")


# Synchronous wrapper for use in non-async contexts (e.g., Gradio callbacks)
def summarize_grants_streaming_sync(
    current: List[Dict[str, Any]],
    past_winners: Optional[List[Dict[str, Any]]] = None,
    *,
    limit: Optional[int] = None,
    client: Optional[LLMClient] = None,
    cache: Optional[SummaryCache] = None,
    batch_size: int = 5,
    mode: str = "batch_parallel",  # "grant_parallel" or "batch_parallel"
) -> List[Dict[str, Any]]:
    """
    Synchronous wrapper for streaming summarization.

    Returns all results as a list (blocking until complete).
    Use this when you need results in order.
    """
    async def run():
        results = []
        if mode == "batch_parallel":
            async for result in summarize_grants_batch_parallel_streaming(
                current,
                past_winners=past_winners,
                limit=limit,
                client=client,
                cache=cache,
                batch_size=batch_size,
            ):
                results.append(result)
        else:  # grant_parallel
            async for result in summarize_grants_parallel_streaming(
                current,
                past_winners=past_winners,
                limit=limit,
                client=client,
                cache=cache,
                batch_size=batch_size,
            ):
                results.append(result)

        # Sort by original index if available
        return sorted(results, key=lambda x: x.get("index", float('inf')))

    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # Already in async context, return async generator
            raise RuntimeError("Use async version directly in async context")
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    return loop.run_until_complete(run())