my-chatgpt 0.1.0

use reqwest::Client;
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
use futures_util::TryStreamExt;

#[derive(Serialize)]
struct ResponseRequest {
    model: String,
    input: String,
    instructions: String,
    stream: bool,
}

#[derive(Deserialize, Debug)]
struct StreamChunk {
    output: Option<Vec<OutputItem>>,
    usage: Option<UsageInfo>,
    #[serde(rename = "type")]
    chunk_type: Option<String>,
}

#[derive(Deserialize, Debug, Clone)]
pub struct UsageInfo {
    pub input_tokens: Option<u32>,
    pub output_tokens: Option<u32>,
    pub total_tokens: Option<u32>,
}

#[derive(Deserialize, Debug)]
struct OutputItem {
    content: Option<Vec<ContentItem>>,
    //#[serde(rename = "type")]
    //type_field: Option<String>,
}

#[derive(Deserialize, Debug)]
struct ContentItem {
    text: Option<String>,
    //#[serde(rename = "type")]
    //content_type: Option<String>,
}

// Helper function to print usage information
fn print_usage(usage: &UsageInfo) {
    println!("\n[Usage Information]");
    println!("Input tokens: {}", usage.input_tokens.unwrap_or(0));
    println!("Output tokens: {}", usage.output_tokens.unwrap_or(0));
    println!("Total tokens: {}", usage.total_tokens.unwrap_or(0));
}

// Helper function to get usage information after streaming
async fn get_final_usage(api_key: &str, model: &str, instructions: &str, input: &str) -> Result<Option<UsageInfo>, Box<dyn std::error::Error>> {
    let client: Client = Client::new();

    let body = ResponseRequest {
        model: model.to_string(),
        instructions: instructions.to_string(),
        input: input.to_string(),
        stream: false, // Non-streaming request to get usage
    };

    let res = client
        .post("https://api.openai.com/v1/responses")
        .header(AUTHORIZATION, format!("Bearer {}", api_key))
        .header(CONTENT_TYPE, "application/json")
        .json(&body)
        .send()
        .await?;
        
    if res.status().is_success() {
        let response_body = res.text().await?;
        if let Ok(value) = serde_json::from_str::<serde_json::Value>(&response_body) {
            if let Some(usage_data) = value.get("usage") {
                if let Ok(usage) = serde_json::from_value::<UsageInfo>(usage_data.clone()) {
                    return Ok(Some(usage));
                }
            }
        }
    }
    
    Ok(None)
}

// Define a general error type that encapsulates all possible error cases
#[derive(Debug, Clone)]
pub enum ResponseError {
    RequestError(String),
    ParseError(String),
    NetworkError(String),
    Unknown(String),
}

pub async fn response<F>(
    instructions: &str, 
    input: &str, 
    api_key: &str, 
    model: &str, 
    stream: bool,
    handler: F
) -> Result<(), Box<dyn std::error::Error>> 
where 
    F: Fn(Option<&UsageInfo>, Option<&ResponseError>, Option<&serde_json::Value>) -> ()
{
    let client: Client = Client::new();

    let body = ResponseRequest {
        model: model.to_string(),
        instructions: instructions.to_string(),
        input: input.to_string(),
        stream,
    };

    let res = match client
        .post("https://api.openai.com/v1/responses")
        .header(AUTHORIZATION, format!("Bearer {}", api_key))
        .header(CONTENT_TYPE, "application/json")
        .json(&body)
        .send()
        .await {
            Ok(res) => res,
            Err(e) => {
                let err = ResponseError::NetworkError(e.to_string());
                handler(None, Some(&err), None);
                return Err(Box::new(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())));
            }
        };

    println!("Response status: {}", res.status());
    
    if !stream {
        // Handle non-streaming response
        let response_body = match res.text().await {
            Ok(body) => body,
            Err(e) => {
                let err = ResponseError::NetworkError(format!("Failed to get response body: {}", e));
                handler(None, Some(&err), None);
                return Err(Box::new(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())));
            }
        };
        
        if let Ok(value) = serde_json::from_str::<serde_json::Value>(&response_body) {
            // Extract and print text output
            if let Some(output) = value.get("output") {
                if let Some(output_array) = output.as_array() {
                    for item in output_array {
                        if let Some(content) = item.get("content") {
                            if let Some(content_array) = content.as_array() {
                                for c in content_array {
                                    if let Some(text) = c.get("text").and_then(|t| t.as_str()) {
                                        print!("{}", text);
                                    }
                                }
                            }
                        }
                    }
                }
            }
            
            // Extract usage information
            let mut usage_found = false;
            if let Some(usage_data) = value.get("usage") {
                if let Ok(usage) = serde_json::from_value::<UsageInfo>(usage_data.clone()) {
                    print_usage(&usage);
                    handler(Some(&usage), None, None);
                    usage_found = true;
                }
            }
            
            // If usage wasn't found in the response, try to get it separately
            if !usage_found {
                println!("\n[Usage information not available in response]");
                println!("Requesting usage information separately...");
                
                match get_final_usage(api_key, model, instructions, input).await {
                    Ok(Some(usage)) => {
                        print_usage(&usage);
                        handler(Some(&usage), None, None);
                    },
                    Ok(None) => {
                        let err = ResponseError::Unknown("Could not retrieve usage information.".to_string());
                        handler(None, Some(&err), None);
                    },
                    Err(e) => {
                        let err = ResponseError::Unknown(format!("Error getting usage: {}", e));
                        handler(None, Some(&err), None);
                    }
                }
            }
            
            return Ok(());
        } else {
            let err = ResponseError::ParseError("Failed to parse non-streaming response".to_string());
            handler(None, Some(&err), None);
            
            // Even if we failed to parse the response, try to get usage information
            match get_final_usage(api_key, model, instructions, input).await {
                Ok(Some(usage)) => {
                    print_usage(&usage);
                    handler(Some(&usage), None, None);
                },
                Ok(None) => {
                    let err = ResponseError::Unknown("Could not retrieve usage information.".to_string());
                    handler(None, Some(&err), None);
                },
                Err(e) => {
                    let err = ResponseError::Unknown(format!("Error getting usage: {}", e));
                    handler(None, Some(&err), None);
                }
            }
            
            return Ok(());
        }
    }
    
    let mut stream = res.bytes_stream();
    let mut last_usage: Option<UsageInfo> = None;

    while let Some(chunk) = stream.try_next().await? {
        let text = match std::str::from_utf8(&chunk) {
            Ok(t) => t,
            Err(e) => {
                let err = ResponseError::ParseError(format!("Failed to parse UTF-8: {}", e));
                handler(None, Some(&err), None);
                continue;
            }
        };

        for line in text.lines() {
            if line.starts_with("data: ") {
                let payload = &line[6..];
                
                // Parse the JSON without type checking first to access any fields
                if let Ok(value) = serde_json::from_str::<serde_json::Value>(payload) {
                    // Pass the raw chunk to the handler
                    handler(None, None, Some(&value));
                    
                    // Check for different event types
                    let event_type = value.get("type").and_then(|t| t.as_str());
                    
                    match event_type {
                        Some("response.output_text.delta") => {
                            if let Some(delta) = value.get("delta") {
                                if let Some(text) = delta.get("text").and_then(|t| t.as_str()) {
                                    print!("{}", text);
                                    if let Err(e) = std::io::Write::flush(&mut std::io::stdout()) {
                                        let err = ResponseError::Unknown(format!("Failed to flush stdout: {}", e));
                                        handler(None, Some(&err), None);
                                    }
                                }
                            }
                        },
                        Some("response.usage.complete") | Some("response.usage") => {
                            // Try to get usage directly
                            if let Some(usage_data) = value.get("usage") {
                                if let Ok(usage) = serde_json::from_value::<UsageInfo>(usage_data.clone()) {
                                    last_usage = Some(usage);
                                }
                            } else {
                                // Try to parse the whole object as usage info
                                if let Ok(usage) = serde_json::from_value::<UsageInfo>(value.clone()) {
                                    last_usage = Some(usage);
                                }
                            }
                        },
                        Some("response.done") | Some("done") | Some("completion") => {
                            // Try to get the final usage data
                            if let Some(usage_data) = value.get("usage") {
                                if let Ok(usage) = serde_json::from_value::<UsageInfo>(usage_data.clone()) {
                                    last_usage = Some(usage.clone());
                                }
                            }
                        },
                        _ => {
                            // Check if there's usage data regardless of event type
                            if let Some(usage_data) = value.get("usage") {
                                if let Ok(usage) = serde_json::from_value::<UsageInfo>(usage_data.clone()) {
                                    last_usage = Some(usage);
                                }
                            }
                        }
                    }
                }

                // Check for [DONE] marker or done-related payloads
                if payload == "[DONE]" || 
                   payload.contains("\"type\":\"done\"") || 
                   payload.contains("\"type\":\"response.done\"") {
                    
                    // If it's not just [DONE] but contains more data, try to parse it
                    if payload != "[DONE]" {
                        if let Ok(value) = serde_json::from_str::<serde_json::Value>(payload) {
                            if let Some(usage_data) = value.get("usage") {
                                if let Ok(usage) = serde_json::from_value::<UsageInfo>(usage_data.clone()) {
                                    last_usage = Some(usage);
                                }
                            }
                        }
                    }
                    
                    // Print usage information if available
                    if let Some(usage) = &last_usage {
                        print_usage(usage);
                        handler(Some(usage), None, None);
                    } else {
                        println!("\n[Usage information not available in stream]");
                        println!("Requesting usage information separately...");
                        
                        // Try to get usage with a separate non-streaming request
                        match get_final_usage(api_key, model, instructions, input).await {
                            Ok(Some(usage)) => {
                                print_usage(&usage);
                                handler(Some(&usage), None, None);
                            },
                            Ok(None) => {
                                let err = ResponseError::Unknown("Could not retrieve usage information.".to_string());
                                handler(None, Some(&err), None);
                            },
                            Err(e) => {
                                let err = ResponseError::Unknown(format!("Error getting usage: {}", e));
                                handler(None, Some(&err), None);
                            }
                        }
                    }
                    println!("\n[Stream finished]");
                    return Ok(());
                }

                match serde_json::from_str::<StreamChunk>(payload) {
                    Ok(parsed) => {
                        // Check if this is a usage-specific chunk
                        if parsed.chunk_type.as_deref() == Some("response.usage.complete") ||
                           parsed.chunk_type.as_deref() == Some("response.usage") {
                            if let Some(usage) = &parsed.usage {
                                last_usage = Some(usage.clone());
                            }
                        }
                        
                        // Process normal output content
                        if let Some(output) = &parsed.output {
                            for item in output {
                                if let Some(content) = &item.content {
                                    for c in content {
                                        if let Some(text) = &c.text {
                                            print!("{}", text);
                                            if let Err(e) = std::io::Write::flush(&mut std::io::stdout()) {
                                                let err = ResponseError::Unknown(format!("Failed to flush stdout: {}", e));
                                                handler(None, Some(&err), None);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    },
                    Err(e) => {
                        if !e.to_string().contains("EOF while parsing") {
                            let err = ResponseError::ParseError(format!("Error parsing JSON: {}", e));
                            handler(None, Some(&err), None);
                        }
                        continue;
                    }
                }
            }
        }
    }

    // Handle case where stream ended without [DONE] marker
    println!("\n[Stream ended unexpectedly]");
    if let Some(usage) = &last_usage {
        print_usage(usage);
        handler(Some(usage), None, None);
    } else {
        println!("Requesting usage information separately...");
        match get_final_usage(api_key, model, instructions, input).await {
            Ok(Some(usage)) => {
                print_usage(&usage);
                handler(Some(&usage), None, None);
            },
            Ok(None) => {
                let err = ResponseError::Unknown("Could not retrieve usage information.".to_string());
                handler(None, Some(&err), None);
            },
            Err(e) => {
                let err = ResponseError::Unknown(format!("Error getting usage: {}", e));
                handler(None, Some(&err), None);
            }
        }
    }

    Ok(())
}