"""CYPHER V12 M34 — Streaming "Thinking" Output.

Stream intermediate reasoning tokens to client (style OpenAI o1 / DeepSeek R1).
Wraps a generate function to emit pseudo-CoT as it produces. UX-revolutionary:
user sees "thinking..." progressively rather than waiting for full response.

Implementation:
  - Generator function yields chunks
  - Optional <think>...</think> tags for explicit thinking phase
  - Final <answer>...</answer> tag for ground-truth response
  - Compatible with SSE (Server-Sent Events) HTTP streaming
"""
from __future__ import annotations

import logging
import time
from typing import Callable, Iterator, Any

logger = logging.getLogger(__name__)


THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
ANSWER_OPEN = "<answer>"
ANSWER_CLOSE = "</answer>"


class StreamingThinker:
    """Wraps a token-generating function to emit streamed thinking + answer."""

    def __init__(
        self,
        step_generator: Callable[[str, int, float], str],
        think_steps: int = 3,
        answer_max_tokens: int = 250,
        thinking_temperature: float = 0.35,
        answer_temperature: float = 0.30,
    ):
        self.step_generator = step_generator
        self.think_steps = think_steps
        self.answer_max_tokens = answer_max_tokens
        self.thinking_temperature = thinking_temperature
        self.answer_temperature = answer_temperature

    def stream(
        self,
        prompt: str,
        category: str | None = None,
    ) -> Iterator[dict]:
        """Yields event dicts: {"type": "think"|"answer"|"done", "delta": "...", "elapsed_ms": N}."""
        t_start = time.perf_counter()
        # 1. Thinking phase (multiple short reasoning steps)
        yield {"type": "think_start", "delta": THINK_OPEN, "elapsed_ms": 0}
        accumulated_thoughts: list[str] = []
        for step in range(self.think_steps):
            think_prompt = (
                f"{prompt}\n\nStep {step+1} of reasoning. "
                f"Brief intermediate thought:"
            )
            try:
                thought = self.step_generator(think_prompt, 60, self.thinking_temperature)
            except Exception as e:
                logger.warning(f"Thinking step {step} fail: {e}")
                thought = ""
            if thought:
                accumulated_thoughts.append(thought.strip())
                yield {
                    "type": "think",
                    "delta": f"\nStep {step+1}: {thought.strip()}",
                    "elapsed_ms": int((time.perf_counter() - t_start) * 1000),
                }
        yield {"type": "think_end", "delta": THINK_CLOSE, "elapsed_ms": int((time.perf_counter() - t_start) * 1000)}
        # 2. Final answer
        yield {"type": "answer_start", "delta": ANSWER_OPEN, "elapsed_ms": int((time.perf_counter() - t_start) * 1000)}
        thoughts_summary = " | ".join(accumulated_thoughts)[:300]
        answer_prompt = (
            f"{prompt}\n\nPrior reasoning: {thoughts_summary}\n\nFinal answer:"
        )
        try:
            final = self.step_generator(answer_prompt, self.answer_max_tokens,
                                         self.answer_temperature)
        except Exception as e:
            logger.warning(f"Final answer fail: {e}")
            final = ""
        # In a real streaming setup, the final would be yielded token-by-token.
        # Here we yield it as a single chunk (encapsulated answer).
        if final:
            yield {
                "type": "answer",
                "delta": final.strip(),
                "elapsed_ms": int((time.perf_counter() - t_start) * 1000),
            }
        yield {"type": "answer_end", "delta": ANSWER_CLOSE, "elapsed_ms": int((time.perf_counter() - t_start) * 1000)}
        yield {
            "type": "done",
            "elapsed_ms": int((time.perf_counter() - t_start) * 1000),
            "thinking_steps": self.think_steps,
            "n_thoughts_captured": len(accumulated_thoughts),
        }

    def collect_all(self, prompt: str, category: str | None = None) -> dict:
        """Drain the stream into a single dict (for non-streaming clients)."""
        thinking_parts: list[str] = []
        answer_parts: list[str] = []
        done_info: dict = {}
        for event in self.stream(prompt, category):
            t = event["type"]
            if t == "think":
                thinking_parts.append(event["delta"])
            elif t == "answer":
                answer_parts.append(event["delta"])
            elif t == "done":
                done_info = event
        return {
            "thinking": "".join(thinking_parts).strip(),
            "answer": "".join(answer_parts).strip(),
            "stats": done_info,
        }


def sse_format(event: dict) -> str:
    """Format event dict as Server-Sent Events string."""
    import json
    return f"data: {json.dumps(event, ensure_ascii=False)}\n\n"


__all__ = ["StreamingThinker", "sse_format",
           "THINK_OPEN", "THINK_CLOSE", "ANSWER_OPEN", "ANSWER_CLOSE"]


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("=== M34 cypher_streaming_thinking SMOKE ===")
    canned = [
        "Identify the CVE: CVE-2021-44228 (Log4Shell).",
        "Assess severity: CVSS 10.0 critical.",
        "Identify mitigation: upgrade Log4j 2.17.1+.",
        "Final answer: Log4Shell is critical RCE in Log4j2, CVSS 10.0, patch to 2.17.1+.",
    ]
    idx = [0]
    def mock_gen(p, mt, t):
        r = canned[idx[0] % len(canned)]
        idx[0] += 1
        return r
    thinker = StreamingThinker(step_generator=mock_gen, think_steps=3)
    print("\n--- Stream events ---")
    for ev in thinker.stream("What is Log4Shell?"):
        print(f"  [{ev['type']}] {ev.get('delta', '')[:80]}")
    print("\n--- Collected ---")
    result = thinker.collect_all("What is Log4Shell?")
    print(f"Thinking: {result['thinking'][:200]}")
    print(f"Answer: {result['answer'][:200]}")
    print(f"Stats: {result['stats']}")
    # SSE format sample
    print(f"\nSSE sample: {sse_format({'type': 'answer', 'delta': 'hello'})[:80]}")
    print("=== SMOKE PASS ===")