#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 10 22:59:42 2025

@author: samd
"""
from typing import Optional, List, Union
import openai
from openai import OpenAI
import anthropic
from google import genai

import pandas as pd
import numpy as np
import bilbystats as bs


def openai_api_stream(content: str, instructions: str = "", model_name: str = "gpt-4o"):
    """
    Generator version of openai_api that yields chunks as they arrive.

    This function yields individual content chunks from the OpenAI API in real-time,
    allowing for true streaming without printing to stdout.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    content : str
        The main user-provided content or prompt to process.
    instructions : str
        The system or developer-level instructions that guide the model's behavior.
    model_name : str, optional (default="gpt-4o")
        The name of the OpenAI model to use for generating the completion.

    ---------------------------------------------------------------------------
    YIELDS:
    chunk : str
        Individual content chunks as they arrive from the API.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if model_name == "deepseek":
        model_name = "deepseek-chat"

    if model_name in ("gpt-4o", "gpt-4o-mini"):
        openai.api_key = bs.read_api_key("openai")
        client = openai
    elif model_name in ("deepseek-chat"):
        client = OpenAI(api_key=bs.read_api_key("deepseek"),
                        base_url="https://api.deepseek.com")
    else:
        raise ValueError(f"Unsupported model_name: {model_name}")

    completion = client.chat.completions.create(
        model=model_name,
        messages=[
            {"role": "system", "content": instructions},
            {"role": "user", "content": content}
        ],
        stream=True
    )

    for chunk in completion:
        if chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content


def llm_api_stream(content: str, instructions: str = "", model_name: str = "gpt-4o"):
    """
    Generator version of llm_api that yields chunks for true streaming.

    This function provides a unified streaming interface across different LLM providers,
    yielding content chunks as they arrive from the API.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    content : str
        User-provided prompt or content for which a response is requested.
    instructions : str
        System-level instructions to guide the model's behavior.
    model_name : str, optional
        The identifier for the target model to be used.

    ---------------------------------------------------------------------------
    YIELDS:
    chunk : str
        Individual content chunks as they arrive from the API.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if model_name in ("deepseek", "gpt-4o", "gpt-4o-mini"):
        for chunk in openai_api_stream(content, instructions, model_name):
            yield chunk
    elif model_name in ("claude", "claude-3-7-sonnet-20250219"):
        # Would need similar streaming implementation for Claude
        # For now, fallback to word-by-word of complete response
        response = bs.claude_api(content, instructions, model_name, 1, False)
        for word in response.split():
            yield word + " "
    elif model_name in ("gemini", "gemini-2.0-flash"):
        # Would need similar streaming implementation for Gemini
        # For now, fallback to word-by-word of complete response
        response = bs.gemini_api(instructions + content, model_name, False)
        for word in response.split():
            yield word + " "
    elif model_name in ("llama", "llama3.2", "deepseek-r1:7b", "ds"):
        # Ollama models already support some streaming via stdout
        # For true streaming, would need to modify ollama function
        response = bs.ollama(content, instructions, model_name, False)
        for word in response.split():
            yield word + " "
    else:
        # Default fallback
        response = bs.ollama(content, instructions, model_name, False)
        for word in response.split():
            yield word + " "
