my-chatgpt 0.1.3

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

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Message {
    pub role: String,
    pub content: String,
}

#[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 send_chat<F>(
    instructions: &str, 
    input: &str, 
    api_key: &str, 
    model: &str, 
    stream: bool,
    history: &[Message],
    handler: F
) -> Result<Vec<Message>, Box<dyn std::error::Error>> 
where 
    F: Fn(Option<&UsageInfo>, Option<&ResponseError>, Option<&serde_json::Value>) -> ()
{
    let client: Client = Client::new();

    // Create a new message for the current input
    let user_message = Message {
        role: "user".to_string(),
        content: input.to_string(),
    };

    // Initialize updated history
    let mut updated_history = Vec::new();
    
    // Add system message if not already in history
    let has_system = history.iter().any(|m| m.role == "developer");
    if !has_system {
        updated_history.push(Message {
            role: "developer".to_string(),
            content: instructions.to_string(),
        });
    } else {
        // Copy existing system message(s)
        updated_history.extend(history.iter()
            .filter(|m| m.role == "developer")
            .cloned());
    }
    
    // Add the rest of the history
    updated_history.extend(history.iter()
        .filter(|m| m.role != "developer")
        .cloned());
    
    // Add current user message
    updated_history.push(user_message);

    // Debug print the history
    println!("Current history before API call:");
    for (i, msg) in updated_history.iter().enumerate() {
        println!("  {}: Role: {}, Content: {}", i, msg.role, msg.content);
    }

    // Format the input with conversation history for the API
    let mut formatted_input = String::new();
    
    // Add conversation history to formatted input
    for msg in &updated_history {
        formatted_input.push_str(&format!("{}: {}\n", msg.role, msg.content));
    }

    let body = ResponseRequest {
        model: model.to_string(),
        instructions: instructions.to_string(),
        input: formatted_input,
        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 !res.status().is_success() {
        let error_text = match res.text().await {
            Ok(text) => text,
            Err(e) => format!("Failed to get error response: {}", e),
        };
        println!("Error response: {}", error_text);
        return Err(Box::new(std::io::Error::new(std::io::ErrorKind::Other, format!("API request failed: {}", error_text))));
    }
    
    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())));
            }
        };
        
        println!("Non-streaming response body: {}", response_body);
        
        if let Ok(value) = serde_json::from_str::<serde_json::Value>(&response_body) {
            // Extract and print text output
            let mut assistant_response = String::new();
            
            // Try to extract from "output" field (common format)
            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);
                                        assistant_response.push_str(text);
                                    }
                                }
                            } else if let Some(text) = content.as_str() {
                                print!("{}", text);
                                assistant_response.push_str(text);
                            }
                        } else if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
                            print!("{}", text);
                            assistant_response.push_str(text);
                        }
                    }
                } else if let Some(text) = output.as_str() {
                    print!("{}", text);
                    assistant_response.push_str(text);
                }
            }
            
            // Try other formats
            if assistant_response.is_empty() {
                if let Some(content) = value.get("content").and_then(|c| c.as_str()) {
                    print!("{}", content);
                    assistant_response.push_str(content);
                } else if let Some(text) = value.get("text").and_then(|t| t.as_str()) {
                    print!("{}", text);
                    assistant_response.push_str(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);
                    }
                }
            }
            
            // If direct extraction didn't work, try the helper function
            if assistant_response.is_empty() {
                println!("Initial extraction failed, trying fallback extraction methods");
                debug_response_extraction(&response_body);
                
                if let Some(extracted) = extract_assistant_response(&response_body) {
                    assistant_response = extracted;
                } else {
                    // Last resort: try to extract from raw text
                    if response_body.contains("assistant:") {
                        if let Some(start_idx) = response_body.find("assistant:") {
                            let response = &response_body[start_idx + 10..];
                            if let Some(end_idx) = response.find('\n') {
                                assistant_response = response[..end_idx].trim().to_string();
                            } else {
                                assistant_response = response.trim().to_string();
                            }
                        }
                    }
                }
            }
            
            // Ensure we add the assistant's response if we found one
            if !assistant_response.is_empty() {
                println!("Adding assistant response to history (non-streaming): '{}'", assistant_response);
                updated_history.push(Message {
                    role: "assistant".to_string(),
                    content: assistant_response,
                });
            } else {
                println!("WARNING: No assistant response found in non-streaming response");
            }
            
            // Debug print the final history
            println!("Final history after API response:");
            for (i, msg) in updated_history.iter().enumerate() {
                println!("  {}: Role: {}, Content: {}", i, msg.role, msg.content);
            }
            
            return Ok(updated_history);
        } else {
            let err = ResponseError::ParseError("Failed to parse non-streaming response".to_string());
            handler(None, Some(&err), None);
            
            // Even if parsing failed, return history with user's message
            println!("WARNING: Failed to parse response, returning history without assistant response");
            return Ok(updated_history);
        }
    }
    
    let mut stream = res.bytes_stream();
    let mut last_usage: Option<UsageInfo> = None;
    let mut complete_response = String::new();
    
    while let Some(chunk) = stream.try_next().await? {
        let text = match std::str::from_utf8(&chunk) {
            Ok(t) => {
                println!("chunk: {}", t);
                complete_response.push_str(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()) {
                                    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), Some(&delta));
                                    }
                                }
                            }
                        },
                        Some("response.output_item.done") => {
                            // This is where the complete response is available
                            println!("Received response.output_item.done event: {:?}", value);
                            let assistant_response = value
                                .get("item")
                                .and_then(|item| item.get("content"))
                                .and_then(|content| content.as_array())
                                .and_then(|content_array| content_array.first())
                                .and_then(|first_item| first_item.get("text"))
                                .and_then(|text| text.as_str())
                                .map(|s| s.to_string());
                            
                            if let Some(assistant_response) = assistant_response {
                                updated_history.push(Message {
                                    role: "assistant".to_string(),
                                    content: assistant_response,
                                });
                            }
                        },
                        Some("response.usage.complete") | Some("response.usage") => {
                            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
                            println!("Received response.done event: {:?}", value);
                            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]");
                    
                    // Extract full assistant response from complete_response
                    let mut assistant_response = String::new();
                    
                    // Try to extract assistant message from the complete response
                    if let Some(extracted) = extract_assistant_message_from_stream(&complete_response) {
                        assistant_response = extracted;
                        println!("Extracted assistant response from complete stream: '{}'", assistant_response);
                        
                        updated_history.push(Message {
                            role: "assistant".to_string(),
                            content: assistant_response,
                        });
                    } else {
                        // Fallback: try pattern matching to find assistant's message
                        println!("Falling back to pattern matching extraction");
                        debug_response_extraction(&complete_response);
                    }
                    
                    // Debug print the final history
                    println!("Final history after streaming:");
                    for (i, msg) in updated_history.iter().enumerate() {
                        println!("  {}: Role: {}, Content: {}", i, msg.role, msg.content);
                    }
                    
                    return Ok(updated_history);
                }
            }
        }
    }

    // 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(updated_history)
}

// Helper function to extract assistant's response from the response body
fn extract_assistant_response(response_body: &str) -> Option<String> {
    println!("Extracting assistant response from body");
    
    // If the response body is too large, print a preview
    if response_body.len() > 1000 {
        println!("Response body preview: {}", &response_body[..500]);
    } else {
        println!("Response body: {}", response_body);
    }
    
    if let Ok(value) = serde_json::from_str::<serde_json::Value>(response_body) {
        // Try to extract from multiple possible formats
        
        // Format 1: Output array with content
        if let Some(output) = value.get("output") {
            if let Some(output_array) = output.as_array() {
                let mut response = String::new();
                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()) {
                                    response.push_str(text);
                                }
                            }
                        } else if let Some(text) = content.as_str() {
                            response.push_str(text);
                        }
                    } else if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
                        response.push_str(text);
                    }
                }
                if !response.is_empty() {
                    println!("Found response in output array: {}", response);
                    return Some(response);
                }
            } else if let Some(text) = output.as_str() {
                println!("Found response as direct output string: {}", text);
                return Some(text.to_string());
            }
        }
        
        // Format 2: Direct "content" field
        if let Some(content) = value.get("content") {
            if let Some(text) = content.as_str() {
                println!("Found response in content field: {}", text);
                return Some(text.to_string());
            }
        }
        
        // Format 3: Try to find any text field that might contain the response
        if let Some(text) = value.get("text").and_then(|t| t.as_str()) {
            println!("Found response in text field: {}", text);
            return Some(text.to_string());
        }
        
        println!("Could not find assistant response in structured JSON");
    } else {
        println!("Could not parse response body as JSON");
    }
    
    // If all structured approaches fail, try to extract using heuristics
    if response_body.contains("assistant:") {
        if let Some(start_idx) = response_body.find("assistant:") {
            let response = response_body[start_idx + 10..].trim().to_string();
            println!("Found response using heuristics: {}", response);
            return Some(response);
        }
    }
    
    None
}

// Helper function to debug response extraction failures
fn debug_response_extraction(response_body: &str) {
    println!("\n=== DEBUG RESPONSE EXTRACTION ===");
    println!("Response body length: {}", response_body.len());
    
    // Print a preview of the response
    let preview_length = std::cmp::min(response_body.len(), 500);
    println!("Response preview: {}", &response_body[..preview_length]);
    
    // Check for common patterns
    println!("Contains 'output': {}", response_body.contains("output"));
    println!("Contains 'content': {}", response_body.contains("content"));
    println!("Contains 'text': {}", response_body.contains("text"));
    println!("Contains 'assistant': {}", response_body.contains("assistant"));
    
    // Try to parse as JSON and inspect structure
    if let Ok(value) = serde_json::from_str::<serde_json::Value>(response_body) {
        println!("Successfully parsed as JSON");
        println!("Top-level keys: {:?}", value.as_object().map(|o| o.keys().collect::<Vec<_>>()));
        
        // Check output structure if present
        if let Some(output) = value.get("output") {
            println!("Output type: {}", if output.is_array() { "array" } 
                                       else if output.is_object() { "object" }
                                       else if output.is_string() { "string" }
                                       else { "other" });
            
            if let Some(output_array) = output.as_array() {
                println!("Output array length: {}", output_array.len());
                if !output_array.is_empty() {
                    println!("First output item type: {}", 
                        if output_array[0].is_object() { "object" }
                        else if output_array[0].is_string() { "string" }
                        else { "other" });
                    
                    if let Some(first_item) = output_array.get(0).and_then(|i| i.as_object()) {
                        println!("First output item keys: {:?}", first_item.keys().collect::<Vec<_>>());
                    }
                }
            }
        }
        
        // Check for error information
        if let Some(error) = value.get("error") {
            println!("Error information found: {}", error);
        }
    } else {
        println!("Failed to parse as JSON");
    }
    
    println!("=== END DEBUG RESPONSE EXTRACTION ===\n");
}

// Helper function to extract assistant's message from streaming response
fn extract_assistant_message_from_stream(complete_response: &str) -> Option<String> {
    // Look for the line with "response.output_item.done" type
    for line in complete_response.lines() {
        if line.starts_with("data: ") {
            let payload = &line[6..];
            
            if let Ok(value) = serde_json::from_str::<serde_json::Value>(payload) {
                if value.get("type").and_then(|t| t.as_str()) == Some("response.output_item.done") {
                    println!("Found response.output_item.done event");
                    
                    // Extract text from this line's content
                    if let Some(content) = value.get("content") {
                        if let Some(text) = content.get("text").and_then(|t| t.as_str()) {
                            println!("Extracted complete text: {}", text);
                            return Some(text.to_string());
                        }
                    }
                }
            }
        }
    }
    
    // Fallback: accumulate response from deltas if done event doesn't contain text
    let mut assistant_message = String::new();
    
    for line in complete_response.lines() {
        if line.starts_with("data: ") {
            let payload = &line[6..];
            
            if let Ok(value) = serde_json::from_str::<serde_json::Value>(payload) {
                if value.get("type").and_then(|t| t.as_str()) == 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()) {
                            assistant_message.push_str(text);
                        }
                    }
                }
            }
        }
    }
    
    if !assistant_message.is_empty() {
        println!("Extracted text from deltas: {}", assistant_message);
        Some(assistant_message)
    } else {
        None
    }
}