from __future__ import annotations

import argparse
import os
import time
from collections.abc import Iterable


class StreamingAIAgent:
    """Small AI-agent with a streaming interface and an offline fallback."""

    def __init__(self, model: str | None = None) -> None:
        self.model = model or os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free")

    def invoke(self, question: str) -> str:
        return "".join(self.stream(question))

    def stream(self, question: str) -> Iterable[str]:
        client = self._build_client()
        if client is None:
            yield from self._offline_stream(question)
            return

        response = client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": "You are a concise educational assistant.",
                },
                {"role": "user", "content": question},
            ],
            stream=True,
        )
        for chunk in response:
            delta = chunk.choices[0].delta.content
            if delta:
                yield delta

    def _build_client(self):
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            return None
        try:
            from openai import OpenAI
        except ImportError:
            return None
        return OpenAI(
            api_key=api_key,
            base_url=os.getenv("OPENAI_BASE_URL") or None,
        )

    def _offline_stream(self, question: str) -> Iterable[str]:
        answer = (
            "Stream mode returns an answer piece by piece instead of waiting "
            "for the whole response. This makes an AI-agent feel faster and "
            f"lets the UI show progress while processing: {question}"
        )
        for token in answer.split(" "):
            yield token + " "
            time.sleep(0.02)


def main() -> None:
    parser = argparse.ArgumentParser(description="Streaming AI-agent demo")
    parser.add_argument("question", nargs="*", help="User question")
    args = parser.parse_args()

    question = " ".join(args.question).strip() or "Explain streaming mode"
    agent = StreamingAIAgent()
    for token in agent.stream(question):
        print(token, end="", flush=True)
    print()


if __name__ == "__main__":
    main()