<?php

if ( ! defined( 'ABSPATH' ) ) {
    exit;
}

class Niwa_AI_Embedding_Client {
    const OPENAI_EMBEDDINGS_ENDPOINT = 'https://api.openai.com/v1/embeddings';
    const NIWA_EMBEDDINGS_ENDPOINT = 'https://niwachat.com/wp-json/niwa-calendar-oauth/v1/embeddings';
    const DEFAULT_EMBEDDING_MODEL = 'text-embedding-3-large';

    private $relay;
    private $unavailable = false;

    public function __construct( $relay = null ) {
        $this->relay = $relay;
    }

    public function embed_text( $text ) {
        $embeds = $this->embed_texts( array( $text ) );
        return ! empty( $embeds[0] ) ? $embeds[0] : array();
    }

    public function embed_texts( $texts ) {
        if ( $this->unavailable ) {
            return array();
        }

        $texts = is_array( $texts ) ? array_values( $texts ) : array( $texts );

        // The public return contract is keyed by the caller's original input
        // positions. Blank inputs are not sent to the API, but they remain
        // holes in the returned array; valid falsey strings such as "0" must
        // never be filtered out or cause later vectors to shift positions.
        $request_texts          = array();
        $request_to_input_index = array();
        foreach ( $texts as $input_index => $text ) {
            $text = (string) $text;
            if ( '' === trim( $text ) ) {
                continue;
            }
            $request_to_input_index[] = intval( $input_index );
            $request_texts[]          = $text;
        }
        if ( empty( $request_texts ) ) {
            return array();
        }

        $api_key = $this->get_openai_api_key();
        $endpoint = '' !== $api_key ? self::OPENAI_EMBEDDINGS_ENDPOINT : self::NIWA_EMBEDDINGS_ENDPOINT;
        $model = self::DEFAULT_EMBEDDING_MODEL;
        $audit_id      = class_exists( 'Niwa_AI_LLM_Audit_Logger' ) ? Niwa_AI_LLM_Audit_Logger::new_event_id() : '';
        $audit_started = microtime( true );

        $body = '' !== $api_key ? $this->request_openai_embeddings( $request_texts, $api_key, $model ) : null;
        if ( '' === $api_key || is_wp_error( $body ) ) {
            $endpoint = self::NIWA_EMBEDDINGS_ENDPOINT;
            $body     = $this->relay()->embeddings( $request_texts );
        }
        if ( is_wp_error( $body ) ) {
            $message = 'Embeddings service failed: ' . $body->get_error_message();
            $this->unavailable = true;
            $this->set_last_error( $message );
            $this->record_embedding_audit_event(
                $audit_id,
                $endpoint,
                $model,
                $request_texts,
                $audit_started,
                array(
                    'success' => false,
                    'status' => 0,
                    'error' => $message,
                    'body' => array(),
                )
            );
            return array();
        }

        if ( empty( $body['data'] ) || ! is_array( $body['data'] ) ) {
            $this->set_last_error( 'Embeddings response missing data.' );
            $this->record_embedding_audit_event(
                $audit_id,
                $endpoint,
                $model,
                $request_texts,
                $audit_started,
                array(
                    'success' => false,
                    'status' => 200,
                    'error' => 'Embeddings response missing data.',
                    'body' => is_array( $body ) ? $body : array(),
                )
            );
            return array();
        }

        $out            = array();
        $expected_count = count( $request_texts );
        foreach ( $body['data'] as $item ) {
            $index = isset( $item['index'] ) ? intval( $item['index'] ) : -1;
            if ( $index >= 0 && $index < $expected_count && isset( $item['embedding'] ) && is_array( $item['embedding'] ) ) {
                $out[ $request_to_input_index[ $index ] ] = $item['embedding'];
            }
        }

        if ( empty( $out ) ) {
            $this->set_last_error( 'Embeddings response contained no vectors.' );
            $this->record_embedding_audit_event(
                $audit_id,
                $endpoint,
                $model,
                $request_texts,
                $audit_started,
                array(
                    'success' => false,
                    'status' => 200,
                    'error' => 'Embeddings response contained no vectors.',
                    'body' => is_array( $body ) ? $body : array(),
                )
            );
            return array();
        }

        ksort( $out );
        $missing_indexes = array_values( array_diff( $request_to_input_index, array_map( 'intval', array_keys( $out ) ) ) );
        $partial_error   = '';
        if ( ! empty( $missing_indexes ) ) {
            $partial_error = sprintf(
                'Embeddings response partial: received %d of %d vectors; missing input indexes: %s.',
                count( $out ),
                $expected_count,
                implode( ', ', $missing_indexes )
            );
            $this->set_last_error( $partial_error );
        } else {
            $this->clear_last_error();
        }
        $first_vector = reset( $out );
        $this->record_embedding_audit_event(
            $audit_id,
            $endpoint,
            $model,
            $request_texts,
            $audit_started,
            array(
                'success' => '' === $partial_error,
                'status' => 200,
                'error' => $partial_error,
                'body' => is_array( $body ) ? $body : array(),
                'vector_count' => count( $out ),
                'vector_dimensions' => is_array( $first_vector ) ? count( $first_vector ) : 0,
                'partial' => '' !== $partial_error,
                'missing_input_indexes' => $missing_indexes,
            )
        );
        // Keep API indexes intact. Reindexing a partial response (for example
        // indexes 0 and 2) would attach vector 2 to input 1 and silently corrupt
        // retrieval. Callers can now detect the missing position and retry or
        // skip only that input. Complete responses still have normal 0..N keys.
        return $out;
    }

    private function relay() {
        if ( null === $this->relay ) {
            $this->relay = new Niwa_AI_Calendar_OAuth_Relay_Client();
        }
        return $this->relay;
    }

    private function request_openai_embeddings( array $texts, $api_key, $model ) {
        $response = wp_remote_post(
            self::OPENAI_EMBEDDINGS_ENDPOINT,
            array(
                'timeout' => 45,
                'headers' => array(
                    'Content-Type'  => 'application/json',
                    'Authorization' => 'Bearer ' . (string) $api_key,
                ),
                'body'    => wp_json_encode(
                    array(
                        'model' => (string) $model,
                        'input' => array_values( $texts ),
                    )
                ),
            )
        );
        if ( is_wp_error( $response ) ) {
            return $response;
        }
        $status = (int) wp_remote_retrieve_response_code( $response );
        $body   = json_decode( (string) wp_remote_retrieve_body( $response ), true );
        if ( $status < 200 || $status >= 300 ) {
            $message = is_array( $body ) ? (string) ( $body['error']['message'] ?? '' ) : '';
            return new WP_Error( 'openai_embeddings_http', '' !== $message ? $message : 'OpenAI embeddings request failed.', $status );
        }
        if ( ! is_array( $body ) || empty( $body['data'] ) || ! is_array( $body['data'] ) ) {
            return new WP_Error( 'openai_embeddings_invalid', 'OpenAI embeddings response missing data.' );
        }
        return $body;
    }

    private function get_openai_api_key() {
        if ( defined( 'NIWA_AI_OPENAI_API_KEY' ) ) {
            return trim( (string) NIWA_AI_OPENAI_API_KEY );
        }
        $auth = get_option( Niwa_AI_Constants::OPT_AUTH, array() );
        return isset( $auth['api_key'] ) ? trim( (string) $auth['api_key'] ) : '';
    }

    private function record_embedding_audit_event( $audit_id, $endpoint, $model, $texts, $started_at, $result ) {
        if ( ! class_exists( 'Niwa_AI_LLM_Audit_Logger' ) || ! Niwa_AI_LLM_Audit_Logger::is_enabled() ) {
            return;
        }

        $body = isset( $result['body'] ) && is_array( $result['body'] ) ? $result['body'] : array();
        Niwa_AI_LLM_Audit_Logger::record_event(
            'llm_embeddings_request',
            array(
                'audit_id' => $audit_id,
                'endpoint' => 'openai_embeddings',
                'endpoint_url' => (string) $endpoint,
                'duration_ms' => round( ( microtime( true ) - floatval( $started_at ) ) * 1000, 2 ),
                'request' => array(
                    'model' => (string) $model,
                    'input_count' => is_array( $texts ) ? count( $texts ) : 0,
                    'input' => $texts,
                ),
                'response' => array(
                    'success' => ! empty( $result['success'] ),
                    'status' => isset( $result['status'] ) ? intval( $result['status'] ) : 0,
                    'error' => isset( $result['error'] ) ? (string) $result['error'] : '',
                    'usage' => isset( $body['usage'] ) && is_array( $body['usage'] ) ? $body['usage'] : array(),
                    'vector_count' => isset( $result['vector_count'] ) ? intval( $result['vector_count'] ) : 0,
                    'vector_dimensions' => isset( $result['vector_dimensions'] ) ? intval( $result['vector_dimensions'] ) : 0,
                    'partial' => ! empty( $result['partial'] ),
                    'missing_input_indexes' => isset( $result['missing_input_indexes'] ) ? array_values( array_map( 'intval', (array) $result['missing_input_indexes'] ) ) : array(),
                    'body' => $body,
                ),
            )
        );
    }

    private function set_last_error( $message ) {
        update_option(
            Niwa_AI_Constants::OPT_EMBEDDING_LAST_ERROR,
            array(
                'message' => (string) $message,
                'at'      => time(),
            ),
            false
        );
    }

    private function clear_last_error() {
        delete_option( Niwa_AI_Constants::OPT_EMBEDDING_LAST_ERROR );
    }
}
