<?php
// Expects to be in /plugin_root/includes folder

class TWXChatOpenAIAssistant
{

    // Class properties (module variables)
    private $_apiKey;
    private $_assistantId;
    private $_assistant;
    private $_apiKeyFiles;
	
    // Constructor method
    public function __construct($apiKey, $assistantId)
    {
        $this->_apiKey = $apiKey;
        $this->_assistantId = $assistantId;
		$this->_assistant = $this->getAssistantPrivate($assistantId);
		$this->_apiKeyFiles = [];
    }

    // Destructor method
    public function __destruct()
    {
        // Log or perform actions needed at object destruction
    }

	public function queryAssistant($query, $threadId = null) {
		
		if(empty($threadId)) {
			$thread = $this->createThread();
			$threadId = $thread['id'];
		}

		$ret = [
			'query' => $query,
			'threadId' => $threadId
		];
		
		//TWXChatUtil::chatLog("$threadId query $query");

		$message = $this->sendMessage($query, $threadId);
		if(isset($message['error'])){
			$ret['answer'] = 'Error: '.$message['error']['message'];
			$ret['error'] = true;
			return $ret;
		}
		
		$response = $this->getMessageResponse($message['id'], $threadId);
		if(isset($response['error'])){
			$ret['answer'] = 'Error: '.$response['error']['message'];
			$ret['error'] = true;
			return $ret;
		}

		$ret['answer'] = TWXChatUtil::parseAnswer($response['answer']);
		$ret['usage'] = $this->getUseageCosts($response['usage']);
		$ret['model'] = $this->_assistant['model'];
		
		return $ret;
	}
	
	public function uploadFileFromPath($newFileName, $sourceFilePath, $contentType, $andAttach = false){

		// create a new File from contents at $sourceFilePath and optionally attaches to Assistant
		// returns Assistant File object
		
		if(class_exists('TWXChatSpider')) TWXChatSpider::log("uploadFileFromPath $newFileName $sourceFilePath $contentType");
		
		$contents = @file_get_contents($sourceFilePath); // Suppress errors for failed fetches
		
		if(empty($contents)){
			$ret['error']['message'] = 'could not read contents from file: '.$sourceFilePath;
			return $ret;
		}

		$response = $this->uploadFile($contents, $newFileName, $contentType, 'assistants');
		if(isset($response['error'])){
			$ret['error']['message'] = $response['error']['message'];
			return $ret;
		}

		if(class_exists('TWXChatSpider')) TWXChatSpider::log("uploadFileFromPath upload successful");

		$file = $response;
		
		if($andAttach === true){
			$assistantFile = $this->attachFile($file['id']);
			if(isset($assistantFile['error'])){
				$ret['error']['message'] = $assistantFile['error']['message'];
				return $ret;
			}
		}

		return $file;  // Assistant file object
	}
	
	public function loadFiles($purpose = null){

		$this->_apiKeyFiles = [];
		
		$files = $this->getFiles($purpose);
		
		foreach($files as $file){
			$this->_apiKeyFiles[$file['id']] = $file;
		}
		
	}
	
	public function getAssistant(){
		
		return $this->_assistant;
	}
	
	public function getAssistants() {
		return TWXChatUtil::getAssistants($this->_apiKey);
	}
	
	public function getVectorStoreFiles($vectorStoreId){
		
		$url = "https://api.openai.com/v1/vector_stores/$vectorStoreId/files";
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'GET');
		if(isset($response['error'])){
			$ret['error']['message'] = $response['error']['message'];
			return $ret;
		}

		return $response['data'];  // Vector Store Files
	}
	
	public function getAssistantFiles() {
		$ret = [];

		// Make sure the assistant is loaded
		if (empty($this->_assistant)) {
			$this->_assistant = $this->getAssistantPrivate($this->_assistantId);
			if (empty($this->_assistant)) {
				return [
					'error' => [
						'message' => 'Failed to load assistant details.'
					]
				];
			}
		}

		// Get vector store IDs from assistant tools
		$vectorStoreIds = [];
		if (isset($this->_assistant['tools'])) {
			foreach ($this->_assistant['tools'] as $tool) {
				if ($tool['type'] === 'file_search' && isset($tool['vector_store_ids'])) {
					$vectorStoreIds = array_merge($vectorStoreIds, $tool['vector_store_ids']);
				}
			}
		}

		$vectorStoreIds = array_unique($vectorStoreIds);

		// If no vector stores are found, return empty array
		if (empty($vectorStoreIds)) {
			return [];
		}

		// Load all files attached to the assistant via vector stores
		$attachedFileIds = [];

		foreach ($vectorStoreIds as $vsId) {
			$url = "https://api.openai.com/v1/vector_stores/{$vsId}/files";

			$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'GET');

			if (isset($response['error'])) {
				// Optionally log error or include it in results
				continue;
			}

			foreach ($response['data'] as $file) {
				$attachedFileIds[] = $file['id'];
			}
		}

		$attachedFileIds = array_unique($attachedFileIds);

		// Load all available files for the API key, if not already loaded
		if (empty($this->_apiKeyFiles)) {
			$this->loadFiles();
		}

		// Match and return detailed file info
		foreach ($attachedFileIds as $fileId) {
			if (isset($this->_apiKeyFiles[$fileId])) {
				$ret[] = $this->_apiKeyFiles[$fileId];
			}
		}

		return $ret;
	}

	public function getVectorStoreIds() {
		if(!isset($this->_assistant['tool_resources']['file_search']['vector_store_ids'])) return [];
		return $this->_assistant['tool_resources']['file_search']['vector_store_ids'];
	}

	public function getAssistantFiles_OLD(){

		$assistantId = $this->_assistantId;
		
		$url = "https://api.openai.com/v1/assistants/$assistantId/files";
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'GET');

		if(isset($response['error'])){
			$ret['error']['message'] = $response['error']['message'];
			return $ret;
		}

		$assistFileIds = array_column($response['data'],'id'); 
		
		// For some reason, sometimes this result can be empty even if there are files attached to the Assistant
		// Lets also search  vector stores attached to the Assistant
		
		$vectorStoreFileIds = [];
		if(isset($this->_assistant['tool_resources']['file_search']['vector_store_ids'])){
			$vectorStoreIds = $this->_assistant['tool_resources']['file_search']['vector_store_ids'];
			foreach($vectorStoreIds as $vectorStoreId){
				$vectoreStoreFiles = $this->getVectorStoreFiles($vectorStoreId);
				if(!isset($vectoreStoreFiles['error'])){
					foreach($vectoreStoreFiles as $vectorStoreFile){
						$vectorStoreFileIds[] = $vectorStoreFile['id'];
					}
				}
			}
		}

		$assistFileIds = array_unique(array_merge($assistFileIds, $vectorStoreFileIds));

		$ret = [];

		if(empty($this->_apiKeyFiles)) $this->loadFiles();

		foreach($assistFileIds as $fileId){
			if(isset($this->_apiKeyFiles[$fileId])){
				$ret[] = $this->_apiKeyFiles[$fileId];
			}
		}
		
		return $ret;
	}
	
	public function detachFile($fileId){

		$assistantId = $this->_assistantId;
		$url = "https://api.openai.com/v2/assistants/$assistantId/files/$fileId";
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'DELETE');

		return $response;  // Detach Status
	}
	
	
	public function replaceAssistantFile($fileId, $newFileName, $sourceFilePath, $contentType){
		
		// First detach fileId from Assistant
		$response = $this->detachFile($fileId);
		if(isset($response['error'])){
			$ret['error']['message'] = $response['error']['message'];
			return $ret;
		}

		// Now delete the file
		$response = $this->deleteFile($fileId);
		if(isset($response['error'])){
			$ret['error']['message'] = $response['error']['message'];
			return $ret;
		}
		
		// Now upload new
		$file = $this->uploadFileFromPath($newFileName, $sourceFilePath, $contentType, true);

		return $file;
	}
	
	public function submitFeedback($message, $feedback, $feedbackFilePath){
		
		// Creates/replaces the twxfeedback.json OpenAI file and attaches to the Assistant

		$ret = [];
		if(!file_exists($feedbackFilePath)){
			try {
				TWXChatUtil::touch($feedbackFilePath);
				if(!file_exists($feedbackFilePath)){
					$ret['error']['message'] = 'unable to create feedback file';
					return $ret;
				}
			}catch(Exception $ex){
				$ret['error']['message'] = 'error creating feedback file: '.$ex->getMessage();
				return $ret;
			}
		}

		$allFeedback = [];
		try {
			$contents = file_get_contents($feedbackFilePath);
			if(TWXChatUtil::isValidJson($contents)) $allFeedback = json_decode($contents,true);
		}catch(Exception $ex){
			// Do Nothing, will overwrite feedback
		}
		
		TWXChatUtil::log("previous feedback: ".json_encode($allFeedback));
		
		// Append new feedback
		$allFeedback[] = [
			'question' => $message,
			'answer' => $feedback
		];

		file_put_contents($feedbackFilePath, json_encode($allFeedback));

		$this->loadFiles(); // Loads all files for the API KEY
		TWXChatUtil::log("all apikey files: ".count($this->_apiKeyFiles));

		$assistantFiles = $this->getAssistantFiles(); // Get any files attached to the Assistant
		TWXChatUtil::log("assistant files: ".count($assistantFiles));

		$assistantFeedbackFileName = 'twxfeedback.json';

		// Detach any previous feedback files from Assistant
		foreach($assistantFiles as $af){

			$pos = strpos($af['filename'], $assistantFeedbackFileName);
			if($pos === 0){

				// Detach the file from the Assistant
				TWXChatUtil::log("detach file {$af['filename']} from Assistant");
				$ret = $this->detachFile($af['id']);
				if(isset($ret['error']['message'])){
					TWXChatUtil::log($ret['error']['message']);
				}else{
					//TWXChatUtil::deleteFile($this->_apiKey, $af['id']);
				}
			}
		}

		// Delete any detached feedback files
		foreach($this->_apiKeyFiles as $af){
			$pos = strpos($af['filename'], $assistantFeedbackFileName);
			if($pos === 0){
				TWXChatUtil::log("delete file {$af['filename']} from Storage");
				TWXChatUtil::deleteFile($this->_apiKey, $af['id']);
			}

		}
		
		TWXChatUtil::log("upload new feedback file with contents: ".json_encode($allFeedback));
		
		//$fileName = 'twxfeedback-'.date('YmdHis').'.json';
		
		// upload feedback contents as a file to OpenAI
		$response = $this->uploadFile(json_encode($allFeedback), $assistantFeedbackFileName, 'json', 'assistants');
		if(isset($response['error'])){
			$ret['error']['message'] = $response['error']['message'];
			return $ret;
		}

		if(!isset($response['id'])){
			$ret['error']['message'] = 'Error creating feedback file in OpenAI';
			return $ret;
		}

		// Attach the file to the Assistant in OpenAI
		$fileId = $response['id'];
		TWXChatUtil::log("attach new feedback file $fileId to Assistant.");

		$assistantFile = $this->attachFile($fileId);
		if(isset($assistantFile['error'])){
			$ret['error']['message'] = $assistantFile['error']['message'];
			return $ret;
		}
		
		$ret = [
			'file_id' => $fileId,
			'file_name' => $assistantFeedbackFileName
		];

		TWXChatUtil::log("return success");

		return $ret;
	}
	
	
	///////////////////////// PRIVATES ///////////////////////////

	private function getAssistantPrivate($assistantId){

		$url = "https://api.openai.com/v1/assistants/$assistantId";
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'POST');
		if(isset($response['error'])){
			return null;
		}

		return $response;
	}
	
	private function createThread()
    {
		
		$url = 'https://api.openai.com/v1/threads';

		$data = [];
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');

		if(isset($response['error'])){
			return $response['error']['message'];
		}

		return $response;	
    }
	
	private function sendMessage($query, $threadId, $role = 'user') {
		
		// asyncronously send message to Assistant. Returns Message object
		
		$url = "https://api.openai.com/v1/threads/$threadId/messages";

		$data = [
			'role' => $role,
			'content' => $query
		];
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');

		return $response;	
	}
	
	private function getMessageResponse($messageId , $threadId){
		
		$assistantId = $this->_assistantId;
		
		// create a run
		$url = "https://api.openai.com/v1/threads/$threadId/runs";
		
		$data = [
			'assistant_id' => $assistantId,
			// "additional_instructions" => "Please address the user as Jane Doe. The user has a premium account."
		];

		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');
		if(isset($response['error'])) return $response;
		
		$runId = $response['id'];
		
		$usage = [];
		
		// Poll Run
		$i = 0; $max_i = 40;
		while($i<$max_i){
			
			$url = "https://api.openai.com/v1/threads/$threadId/runs/$runId";
			$data = [];
			
			$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'GET');
			if(isset($response['error'])){
				return $response['error']['message'];
			}

			if($response['status'] == 'completed') {
				$usage = $response['usage'];
				break;
			}
			sleep(2); // 2 seconds
			$i++;
		}
		
		if($i >= $max_i){
			$ret = [];
			$ret['error']['message'] = 'Timeout getting response. Try again later.';
			return $ret;
		}

		// Get Messages
		$url = "https://api.openai.com/v1/threads/$threadId/messages";
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'GET');
		if(isset($response['error'])){
			return $response['error']['message'];
		}

		foreach($response['data'] as $message){
			if($message['run_id'] == $runId) {
				$ret = [
					'answer' => $message['content'][0]['text']['value'],
					'usage' => $usage
				];
				return $ret; // $message['content'][0]['text']['value'];
			}
		}
		
		return null;
	}
	
	private function uploadFile($contents, $fileName, $contentType = 'application/json', $purpose = 'assistants') {

		// Uploads file contents to the Organisation belonging to the apiKey
		// Returns newly created file object
		
		try{
			$tmpFilePath = TWXChatUtil::getUploadsDir('tmp') . uniqid('feedback_', true).'.json';

			if(file_exists($tmpFilePath)) wp_delete_file($tmpFilePath);
			
			file_put_contents($tmpFilePath, $contents);
			
			$c_file = new CURLFile($tmpFilePath, $contentType, $fileName);

		}catch(Exception $ex){
			
			if(class_exists('TWXChatSpider')) TWXChatSpider::log('uploadFile error: '.$ex->getMessage(),'a',true);
		}
		
		
		$url = "https://api.openai.com/v1/files";

		$data = [
			'purpose' => $purpose,
			'file' => $c_file
		];

		$response = null;
		if(class_exists('TWXChatPlus')){
			$response = TWXChatPlus::goCURLFormData($this->_apiKey, $url, $data);
		}else{
			$response = $this->goWPFormData($url, $data);
		}
		
		wp_delete_file($tmpFilePath);
		
		return $response;
	}
	
	private function attachFile($fileId){


		$vectorStoreId = null;
		$vectorStoreIds = $this->getVectorStoreIds();
		if(empty($vectorStoreIds)){
			// create VS
			if(class_exists('TWXChatSpider')) TWXChatSpider::log("attachFile() no vector store found. create new");

			$response = $this->createVectorStore(true);
			if(isset($response['error'])){
				$ret['error']['message'] = $response['error']['message'];
				return $ret;
			}

			if(!isset($response['id'])){
				$ret['error']['message'] = 'Error creating Vector Store in OpenAI';
				return $ret;
			}

			$vectorStoreId = $response['id'];

		}else{
			$vectorStoreId = $vectorStoreIds[0];
		}
		
		// Attach the File to the VS
		$url = "https://api.openai.com/v1/vector_stores/{$vectorStoreId}/files";

		$data = [
			'file_id' => $fileId
		];

		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');
		
		return $response;  // File Object

	}
	
	public function createVectorStore($andAttachToAssistant = true) {

		$url = "https://api.openai.com/v1/vector_stores";

		$data = [
  			"name"=> "Vector Store for ".$this->_assistant['name'],
  			"description"=> "Vector Store for ".$this->_assistant['name'],
		];

		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');

		if(isset($response['error'])){
			return $response;
		}

		if(!isset($response['id'])){
			$ret['error']['message'] = 'Error creating Vector Store in OpenAI';
			return $ret;
		}

		if($andAttachToAssistant){
			$response2 = $this->attachVectorStoreToAssistant($response['id']);
			if(isset($response2['error'])){
				return $response2;
			}
		}

		return $response;
	}

	public function attachVectorStoreToAssistant($vectorStoreId) {

		$assistantId = $this->_assistant['id'] ?? '';
		if(empty($assistantId)) {
			$ret['error']['message'] = 'Missing Assistant Id';
			return $ret;
		}

		$vectorStoreIds = $this->getVectorStoreIds();
		if(in_array($vectorStoreId, $vectorStoreIds)) {
			$ret['error']['message'] = 'Vector Store Id already attached';
			return $ret;
		}

		$assistant = $this->_assistant;

		// Find existing file_search tool or create a new one
		$tools_resources = isset($assistant['tool_resources']) ? $assistant['tool_resources'] : [];

		$tools_resources['file_search']['vector_store_ids'][] = $vectorStoreId;

		$data['tool_resources'] = $tools_resources;

		$url = "https://api.openai.com/v1/assistants/{$assistantId}";

		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');

		return $response;
	}


	public function deleteFile($fileId){
		
		return TWXChatUtil::deleteFile($this->_apiKey, $fileId);
	}

	public function getFile($fileId){
		
		$url = "https://api.openai.com/v1/files/$fileId";

		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'GET');

		return $response;
	}
	
	public function getFiles($purpose = 'assistants') {

		return TWXChatUtil::getFiles($this->_apiKey, $purpose);
	}
	
	private function getFileContents($fileId){
		
		$url = "https://api.openai.com/v1/files/$fileId/content";
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, null, 'GET');
		
		return $response;
	}
	
	private function getUseageCosts($tokenUsage){
		
		$model = $this->_assistant['model'];

		if(!isset($tokenUsage['prompt_tokens']) || !isset($tokenUsage['completion_tokens']) || empty(TWXChatUtil::getModelCosts($model))) return false;
		
		$ret = [];
		
		$prompt_tokens = $tokenUsage['prompt_tokens'];
		$completion_tokens = $tokenUsage['completion_tokens'];
		
		$modelCosts = TWXChatUtil::getModelCosts($model);
		$totalCost = $prompt_tokens * $modelCosts['in'];
		$totalCost += $completion_tokens * $modelCosts['out'];
		
		$ret['prompt_tokens'] = $prompt_tokens;
		$ret['completion_tokens'] = $completion_tokens;
		$ret['total_tokens'] = $prompt_tokens + $completion_tokens;
		$ret['cost_total'] = number_format($totalCost, 4, '.', '');
		$ret['currency'] = 'USD';
		$ret['is_total_usage'] = true; // is thread total
		
		return $ret;
	}
	
	private function goWPFormData($url, $data) {
		// Submits Form Data POST using wp_remote_post()

		$apiKey = $this->_apiKey;

		// Prepare headers
		$headers = [
			'Authorization' => 'Bearer ' . $apiKey,
			'OpenAI-Beta'   => 'assistants=v2',
			// 'Content-Type' => 'multipart/form-data', // Let WP handle this automatically
		];

		// Prepare args
		$args = [
			'headers' => $headers,
			'body'    => $data,
			'timeout' => 30,
		];

		$response = wp_remote_post($url, $args);

		if (is_wp_error($response)) {
			// Optional: log error
			TWXChatUtil::log('goWPFormData error: ' . $response->get_error_message());
			return false;
		}

		$body = wp_remote_retrieve_body($response);

		return json_decode($body, true);
	}

	
}

class TWXChatOpenAICompletions
{
    // Class properties (module variables)
    private $_apiKey;
    private $_model;
    private $_temperature;
	private $_tools;
	private $_toolChoice;

    // Constructor method
    public function __construct($apiKey, $model = null, $temperature = 0.5)
    {
        $this->_apiKey = $apiKey;
        $this->_model = ($model ? $model : 'gpt-4o-mini');
		$this->_temperature = $temperature;  // creativity 0-2
		$this->_tools = null;
		$this->_toolChoice = 'auto';
    }

    // Destructor method
    public function __destruct()
    {
        // Log or perform actions needed at object destruction
    }
	
	public function getModel(){
		return $this->_model;
	}

	public function setUserFieldsTools($userFields){

		// Single function approach
		if(empty($userFields)) {
			$this->_tools = null;
			$this->_toolChoice = 'auto';
			return;
		}

		// Initialise with full response (answer) property
		$functionName = 'user_fields_main_function';
		$properties = [];
		$properties['full_response'] = [
			"type" => "string",
			"description" => "The full response to the user",
		];
		$required = ["full_response"];

		foreach($userFields as $field){
			$fieldName = $field['name'] ?? '';
			if(empty($fieldName)) continue;

			$properties[$fieldName] = [
				"type" => "string",
				"description" => $field['description'] ?? '',
			];
			//$required[] = $fieldName;
		}

		$this->_tools = [
			[
				"type" => "function",
				"function" => [
					"name" => $functionName,
					"description" => "Always call this function",
					"parameters" => [
						"type" => "object",
						"properties" => $properties,
					],
					"required" => $required,
				]
			]
		];

		$this->_toolChoice = [
			'type' => 'function',
			'function' => [
				'name' => $functionName
			]
		];

	}

	public static function getUserFieldFunction($userField){
		// field must have name and description and optionally a value
		$function = [];
		$fieldName = $userField['name'];
		$fieldDesc = $userField['description'];
		$function['name'] = "{$fieldName}";
		$function['description'] = "Called when {$fieldName} ($fieldDesc) is in the user message";
		$function['parameters']['type'] = "object";
		//$function['parameters']['properties'] = [$fieldName => $userField];
		$function['parameters']['properties'][$fieldName] = ['type'=>'string','description'=>$fieldDesc];
		$function['parameters']['required'] = [$fieldName];
		return $function;
	}
	
	public function query($query) {
		
		$messages = [];
		if(is_array($query)){
			$messages = $query;
		}else{
			$messages = [['role'=>'user', 'content'=>$query]];
		}

		$ret = [];
		$url = "https://api.openai.com/v1/chat/completions";

		$data = [
			'model' => $this->_model,
			'messages' => $messages,
			'temperature' => $this->_temperature,
		];

		if(!empty($this->_tools)){
			$data['tools'] = $this->_tools;
			$data['tool_choice'] = $this->_toolChoice;
		}
		
		$response = TWXChatUtil::getHttpOpenAI_WP($this->_apiKey, $url, $data, 'POST');

		if (isset($response['error']['message'])) {
			$ret['error']['message'] = $response['error']['message'];
		} else {

			$answer = '';

			if(isset($response['choices'][0]['message']['content'])){
				$responseText = $response['choices'][0]['message']['content'];
				$source_ids = TWXChatUtil::extractSourceIdsFromText($responseText);
				$parsedAnswer = TWXChatUtil::parseAnswer($responseText);
				//list($parsed, $source_ids) = TWXChatUtil::parseInlineSources($parsedAnswer);
				$answer = $parsedAnswer;
				if(!empty($source_ids)) $ret['relevant_source_ids'] = $source_ids;
			}

			// Check if OpenAI has tool/function output
			if (isset($response['choices'][0]['message']['tool_calls'])) {
				$tool_calls = $response['choices'][0]['message']['tool_calls'];

				foreach($tool_calls as $tool_call){
					if($tool_call['type'] == 'function'){
						$function = $tool_call['function'];
						$funcName = $function['name'];
						if($funcName === 'user_fields_main_function'){
							$args = json_decode($function['arguments'],true);
							foreach($args as $arg_name => $arg_val){
								if(!empty($arg_val)) {
									if($arg_name === 'full_response') {
										$answer = $arg_val;
									}else{
										$ret['user_fields'][$arg_name] = $arg_val;
									}
								}
							}
						}
					}
				}
			}

			if(empty($answer)){
				$ret['error']['message'] = 'no response';
			}else{
				$ret['answer'] = $answer;
			}

			if(isset($response['usage'])) $ret['usage'] = $this->getUseageCosts($response['usage']);
			$ret['model'] = $this->_model;
		}
		
		return $ret;
	}

	/*
	public function queryStreamOLDWAY($query, callable $callback) {

		$messages = [];
		if(is_array($query)){
			$messages = $query;
		}else{
			$messages = [['role'=>'user', 'content'=>$query]];
		}
		
		$ret = [];
		$url = "https://api.openai.com/v1/chat/completions";

		$data = [
			'model' => $this->_model,
			'messages' => $messages,
			'temperature' => $this->_temperature,
			'stream' => true,
		];

		$fullStreamOutput = ''; // Buffer the full output
		
        $ch = curl_init($url);

        curl_setopt($ch, CURLOPT_RETURNTRANSFER, true);
        curl_setopt($ch, CURLOPT_HTTPHEADER, [
            'Authorization: Bearer ' . $this->_apiKey,
            'Content-Type: application/json'
        ]);
        curl_setopt($ch, CURLOPT_POST, true);
        curl_setopt($ch, CURLOPT_POSTFIELDS, json_encode($data));

        // Set up the response handler to stream data
        curl_setopt($ch, CURLOPT_WRITEFUNCTION, function($ch, $data) use ($callback, &$fullStreamOutput) {
            // For each chunk of data, call the user-provided callback function
            // Typically, OpenAI will send small chunks of the response, so we'll pass them to the callback
            if (!empty($data)) {
				$fullStreamOutput .= $data;
                $callback($data);
            }
            return strlen($data);
        });

        curl_exec($ch); // Cannot use Wordpress functions as streaming is not suported by them

        // Handle any cURL errors
        if ($error = curl_error($ch)) {
			$ret['error']['message'] = $error;
        }else{
			$ret['answer'] = 'streamed_to_callback';
			$ret['fullStreamOutput'] = $fullStreamOutput;
		}

        // Close cURL session
        curl_close($ch);

		return $ret;
	}
	*/

	public function queryStream($query, callable $callback) {

		// Executes OpenAI query and returns text stream back via $callback function - uses Wordpress functions/classes

		$messages = is_array($query) ? $query : [['role' => 'user', 'content' => $query]];

		$data = [
			'model'       => $this->_model,
			'messages'    => $messages,
			'temperature' => $this->_temperature,
			'stream'      => true,
		];

		$fullStreamOutput = '';

		// Add a temporary curl handler to stream data
		add_action('http_api_curl', function($handle) use (&$fullStreamOutput, $callback) {
			curl_setopt($handle, CURLOPT_WRITEFUNCTION, function($ch, $data) use (&$fullStreamOutput, $callback) {
				if (!empty($data)) {
					$fullStreamOutput .= $data;
					$callback($data);
				}
				return strlen($data);
			});
			curl_setopt($handle, CURLOPT_TIMEOUT, 0); // keep stream alive
		});

		$http = new WP_Http(); // Tried using wp_remote_post() but kept closing mid-stream
		$response = $http->request('https://api.openai.com/v1/chat/completions', [
			'method'  => 'POST',
			'headers' => [
				'Authorization' => 'Bearer ' . $this->_apiKey,
				'Content-Type'  => 'application/json',
			],
			'body'    => wp_json_encode($data),
			'timeout' => 0, // streaming can take long
		]);

		// remove curl hook after request
		remove_all_actions('http_api_curl');

		$ret = [];
		if (is_wp_error($response)) {
			$ret['error']['message'] = $response->get_error_message();
		} else {
			$ret['answer'] = 'streamed_to_callback';
			$ret['fullStreamOutput'] = $fullStreamOutput;
		}

		return $ret;
	}

	//// PRIVATES ////

	private function getUseageCosts($tokenUsage){
		
		$model = $this->_model;

		if(!isset($tokenUsage['prompt_tokens']) || !isset($tokenUsage['completion_tokens']) || empty(TWXChatUtil::getModelCosts($model))) return false;
		
		$ret = [];
		
		$prompt_tokens = $tokenUsage['prompt_tokens'];
		$completion_tokens = $tokenUsage['completion_tokens'];
		
		$modelCosts = TWXChatUtil::getModelCosts($model);
		$totalCost = $prompt_tokens * $modelCosts['in'];
		$totalCost += $completion_tokens * $modelCosts['out'];
		
		$ret['prompt_tokens'] = $prompt_tokens;
		$ret['completion_tokens'] = $completion_tokens;
		$ret['total_tokens'] = $prompt_tokens + $completion_tokens;
		$ret['cost_total'] = number_format($totalCost, 4, '.', '');
		$ret['currency'] = 'USD';
		$ret['is_total_usage'] = false;
		
		return $ret;
	}
	
}

