<?php
if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

class Arcai_Provider_OpenAI implements Arcai_Provider {

	const API_BASE = 'https://api.openai.com/v1';

	/**
	 * USD per 1M tokens. Approximate list prices; filterable so they can be
	 * corrected without a plugin release.
	 */
	private function prices() {
		return apply_filters(
			'arcai_openai_prices',
			array(
				'text-embedding-3-small' => array( 'in' => 0.02, 'out' => 0 ),
				'text-embedding-3-large' => array( 'in' => 0.13, 'out' => 0 ),
				'gpt-4o-mini'            => array( 'in' => 0.15, 'out' => 0.60 ),
				'gpt-4o'                 => array( 'in' => 2.50, 'out' => 10.00 ),
			)
		);
	}

	private function price_for( $model ) {
		$prices = $this->prices();
		return isset( $prices[ $model ] ) ? $prices[ $model ] : array( 'in' => 1.0, 'out' => 4.0 );
	}

	private function cost( $model, $tokens_in, $tokens_out = 0 ) {
		$price = $this->price_for( $model );
		return ( $tokens_in / 1000000 ) * $price['in'] + ( $tokens_out / 1000000 ) * $price['out'];
	}

	public function estimate_embedding_cost( $tokens ) {
		return $this->cost( Arcai_Settings::get( 'embedding_model' ), $tokens );
	}

	public function embed( array $texts ) {
		$model    = Arcai_Settings::get( 'embedding_model' );
		$response = $this->request(
			'/embeddings',
			array(
				'model' => $model,
				'input' => array_values( $texts ),
			)
		);
		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$embeddings = array();
		foreach ( $response['data'] as $item ) {
			$embeddings[ $item['index'] ] = $item['embedding'];
		}
		ksort( $embeddings );

		$tokens = isset( $response['usage']['total_tokens'] ) ? (int) $response['usage']['total_tokens'] : 0;

		return array(
			'embeddings' => array_values( $embeddings ),
			'tokens'     => $tokens,
			'cost'       => $this->cost( $model, $tokens ),
		);
	}

	public function chat( array $messages, array $opts = array() ) {
		$model = Arcai_Settings::get( 'chat_model' );
		$body  = array(
			'model'      => $model,
			'messages'   => $messages,
			'max_tokens' => isset( $opts['max_tokens'] ) ? (int) $opts['max_tokens'] : 700,
		);
		if ( ! empty( $opts['json'] ) ) {
			$body['response_format'] = array( 'type' => 'json_object' );
		}

		$response = $this->request( '/chat/completions', $body );
		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$tokens_in  = isset( $response['usage']['prompt_tokens'] ) ? (int) $response['usage']['prompt_tokens'] : 0;
		$tokens_out = isset( $response['usage']['completion_tokens'] ) ? (int) $response['usage']['completion_tokens'] : 0;

		return array(
			'content'    => isset( $response['choices'][0]['message']['content'] ) ? $response['choices'][0]['message']['content'] : '',
			'tokens_in'  => $tokens_in,
			'tokens_out' => $tokens_out,
			'cost'       => $this->cost( $model, $tokens_in, $tokens_out ),
		);
	}

	/**
	 * Streamed completion.
	 *
	 * phpcs:disable WordPress.WP.AlternativeFunctions -- the WordPress HTTP API
	 * buffers the entire response body and cannot deliver chunks as they
	 * arrive, which defeats streaming. Raw cURL with a write callback is the
	 * only way to relay tokens to the visitor in real time; when the cURL
	 * extension is missing we fall back to the WP HTTP API (non-streamed).
	 */
	public function chat_stream( array $messages, array $opts, callable $on_delta ) {
		if ( ! function_exists( 'curl_init' ) ) {
			$result = $this->chat( $messages, $opts );
			if ( ! is_wp_error( $result ) && '' !== $result['content'] ) {
				call_user_func( $on_delta, $result['content'] );
			}
			return $result;
		}

		$api_key = Arcai_Settings::get( 'api_key' );
		if ( '' === $api_key ) {
			return new WP_Error( 'arcai_no_key', __( 'No API key configured.', 'archivist-ai' ) );
		}

		$model = Arcai_Settings::get( 'chat_model' );
		$body  = array(
			'model'          => $model,
			'messages'       => $messages,
			'max_tokens'     => isset( $opts['max_tokens'] ) ? (int) $opts['max_tokens'] : 700,
			'stream'         => true,
			'stream_options' => array( 'include_usage' => true ),
		);

		$state = array(
			'buffer'  => '',
			'content' => '',
			'usage'   => null,
			'error'   => '',
		);

		$ch = curl_init();
		curl_setopt_array(
			$ch,
			array(
				CURLOPT_URL            => self::API_BASE . '/chat/completions',
				CURLOPT_POST           => true,
				CURLOPT_HTTPHEADER     => array(
					'Authorization: Bearer ' . $api_key,
					'Content-Type: application/json',
				),
				CURLOPT_POSTFIELDS     => wp_json_encode( $body ),
				CURLOPT_TIMEOUT        => 120,
				CURLOPT_WRITEFUNCTION  => function ( $ch, $chunk ) use ( &$state, $on_delta ) {
					if ( 200 !== (int) curl_getinfo( $ch, CURLINFO_RESPONSE_CODE ) ) {
						$state['error'] .= $chunk;
						return strlen( $chunk );
					}
					$state['buffer'] .= $chunk;
					while ( false !== ( $pos = strpos( $state['buffer'], "\n" ) ) ) {
						$line            = trim( substr( $state['buffer'], 0, $pos ) );
						$state['buffer'] = substr( $state['buffer'], $pos + 1 );
						if ( 0 !== strpos( $line, 'data:' ) ) {
							continue;
						}
						$data = trim( substr( $line, 5 ) );
						if ( '[DONE]' === $data ) {
							continue;
						}
						$json = json_decode( $data, true );
						if ( ! is_array( $json ) ) {
							continue;
						}
						if ( isset( $json['choices'][0]['delta']['content'] ) && '' !== $json['choices'][0]['delta']['content'] ) {
							$state['content'] .= $json['choices'][0]['delta']['content'];
							call_user_func( $on_delta, $json['choices'][0]['delta']['content'] );
						}
						if ( isset( $json['usage']['prompt_tokens'] ) ) {
							$state['usage'] = $json['usage'];
						}
					}
					return strlen( $chunk );
				},
			)
		);

		$ok         = curl_exec( $ch );
		$curl_error = curl_error( $ch );
		$http_code  = (int) curl_getinfo( $ch, CURLINFO_RESPONSE_CODE );
		curl_close( $ch );

		if ( 200 !== $http_code ) {
			$detail  = json_decode( $state['error'], true );
			$message = isset( $detail['error']['message'] ) ? $detail['error']['message'] : sprintf( 'HTTP %d', $http_code );
			return new WP_Error( 'arcai_api_error', $message );
		}
		if ( false === $ok && '' === $state['content'] ) {
			return new WP_Error( 'arcai_api_error', $curl_error ? $curl_error : __( 'Connection failed.', 'archivist-ai' ) );
		}

		$tokens_in  = isset( $state['usage']['prompt_tokens'] ) ? (int) $state['usage']['prompt_tokens'] : $this->rough_tokens( $messages );
		$tokens_out = isset( $state['usage']['completion_tokens'] ) ? (int) $state['usage']['completion_tokens'] : (int) ceil( mb_strlen( $state['content'] ) / 4 );

		return array(
			'content'    => $state['content'],
			'tokens_in'  => $tokens_in,
			'tokens_out' => $tokens_out,
			'cost'       => $this->cost( $model, $tokens_in, $tokens_out ),
		);
	}
	// phpcs:enable WordPress.WP.AlternativeFunctions

	private function rough_tokens( array $messages ) {
		$chars = 0;
		foreach ( $messages as $message ) {
			$chars += mb_strlen( isset( $message['content'] ) ? $message['content'] : '' );
		}
		return (int) ceil( $chars / 4 );
	}

	private function request( $path, array $body ) {
		$api_key = Arcai_Settings::get( 'api_key' );
		if ( '' === $api_key ) {
			return new WP_Error( 'arcai_no_key', __( 'No API key configured.', 'archivist-ai' ) );
		}

		$response = wp_remote_post(
			self::API_BASE . $path,
			array(
				'timeout' => 90,
				'headers' => array(
					'Authorization' => 'Bearer ' . $api_key,
					'Content-Type'  => 'application/json',
				),
				'body'    => wp_json_encode( $body ),
			)
		);

		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$code = wp_remote_retrieve_response_code( $response );
		$data = json_decode( wp_remote_retrieve_body( $response ), true );

		if ( 200 !== $code || ! is_array( $data ) ) {
			$message = isset( $data['error']['message'] ) ? $data['error']['message'] : sprintf( 'HTTP %d', $code );
			return new WP_Error( 'arcai_api_error', $message );
		}

		return $data;
	}
}
