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

/**
 * Chroma Vector Database Client
 *
 * Implements the generic vector DB contract using Chroma REST API calls.
 *
 * API version follows cromdbapidoc.txt (v2 routes under /api/v2/tenants/{tenant}/databases/{database}).
 *
 * @package    Kitsune_AI_Platform
 * @subpackage Kitsune_AI_Platform/includes/vector-db
 */

class Kitsune_AI_Platform_Vector_DB_Chroma implements Kitsune_AI_Platform_Vector_DB_Interface {

    /**
     * @var string
     */
    private $endpoint;

    /**
     * @var string
     */
    private $api_key;

    /**
     * @var string
     */
    private $tenant_id;

    /**
     * @var string
     */
    private $database;

    /**
     * @var string
     */
    private $collection;

    /**
     * @var string
     */
    private $collection_id;

    /**
     * @var int
     */
    private $timeout;

    /**
     * @var string
     */
    private $v2_base_url;

    /**
     * @param array $config Connection/config options.
     */
    public function __construct($config = array()) {
        $this->endpoint = isset($config['endpoint']) ? rtrim($config['endpoint'], '/') : 'https://api.trychroma.com';
        $this->api_key = isset($config['apiKey']) ? trim($config['apiKey']) : '';
        $this->tenant_id = isset($config['tenantId']) && trim($config['tenantId']) !== '' ? trim($config['tenantId']) : 'default_tenant';
        $this->database = isset($config['database']) && trim($config['database']) !== '' ? trim($config['database']) : 'default_database';
        $this->collection = isset($config['collection']) && trim($config['collection']) !== '' ? trim($config['collection']) : 'ai_platform_kb';
        $this->collection_id = isset($config['collectionId']) ? trim($config['collectionId']) : '';
        $this->timeout = isset($config['timeout']) ? max(5, intval($config['timeout'])) : 20;
        $this->v2_base_url = $this->resolve_v2_base_url();
    }

    /**
     * @inheritDoc
     */
    public function list_records($args = array()) {
        $args = wp_parse_args($args, array(
            'limit' => 50,
            'offset' => 0,
        ));

        $collection_id = $this->resolve_collection_id();
        if (is_wp_error($collection_id)) {
            return array(
                'success' => false,
                'message' => $collection_id->get_error_message(),
            );
        }

        // GET API requires one of ids or where; where:{} fetches all records with pagination.
        $payload = array(
            'include' => array('documents', 'metadatas'),
            'limit' => max(1, intval($args['limit'])),
            'offset' => max(0, intval($args['offset'])),
            'where' => (object) array(),
        );

        $response = $this->post('/collections/' . rawurlencode($collection_id) . '/get', $payload);
        if (empty($response['success'])) {
            return $response;
        }

        $data = $response['data'];
        $ids = isset($data['ids']) && is_array($data['ids']) ? $data['ids'] : array();
        $documents = isset($data['documents']) && is_array($data['documents']) ? $data['documents'] : array();
        $metadatas = isset($data['metadatas']) && is_array($data['metadatas']) ? $data['metadatas'] : array();

        $records = array();
        foreach ($ids as $index => $id) {
            $records[] = array(
                'id' => (string) $id,
                'document' => isset($documents[$index]) ? (string) $documents[$index] : '',
                'metadata' => isset($metadatas[$index]) && is_array($metadatas[$index]) ? $metadatas[$index] : array(),
            );
        }

        return array(
            'success' => true,
            'data' => $records,
            'status' => $response['status'] ?? 200,
        );
    }

    /**
     * Validate connectivity and credentials by resolving the target collection.
     *
     * @return array
     */
    public function probe_connection() {
        $collection_id = $this->resolve_collection_id();
        if (is_wp_error($collection_id)) {
            return array(
                'success' => false,
                'message' => $collection_id->get_error_message(),
            );
        }

        return array(
            'success' => true,
            'data' => array(
                'collection_id' => $collection_id,
            ),
        );
    }

    /**
     * @inheritDoc
     */
    public function vector_search($query, $args = array()) {
        $args = wp_parse_args($args, array(
            'limit'      => 8,
            'query_embedding' => null,
            'source_ids' => array(),
        ));

        $collection_id = $this->resolve_collection_id();
        if (is_wp_error($collection_id)) {
            return array(
                'success' => false,
                'message' => $collection_id->get_error_message(),
            );
        }

        $embedding = isset($args['query_embedding']) && is_array($args['query_embedding'])
            ? $args['query_embedding']
            : null;

        if (!is_array($embedding) || empty($embedding)) {
            return array(
                'success' => false,
                'message' => __('Query embedding is required for vector search.', 'kitsune-ai-platform'),
            );
        }

        $payload = array(
            'query_embeddings' => array($embedding),
            'n_results' => max(1, intval($args['limit'])),
            'include' => array('documents', 'metadatas', 'distances'),
        );

        // Add source_id filter if provided
        $source_ids = isset($args['source_ids']) && is_array($args['source_ids']) ? $args['source_ids'] : array();
        if (!empty($source_ids)) {
            // Chroma filtering: use $in operator to match any source_id in the array
            $payload['where'] = array(
                'source_id' => array(
                    '$in' => array_map('intval', $source_ids),
                ),
            );
        }

        $response = $this->post('/collections/' . rawurlencode($collection_id) . '/query', $payload);
        if (empty($response['success'])) {
            return $response;
        }

        $data = $response['data'];
        $ids = isset($data['ids'][0]) && is_array($data['ids'][0]) ? $data['ids'][0] : array();
        $documents = isset($data['documents'][0]) && is_array($data['documents'][0]) ? $data['documents'][0] : array();
        $metadatas = isset($data['metadatas'][0]) && is_array($data['metadatas'][0]) ? $data['metadatas'][0] : array();
        $distances = isset($data['distances'][0]) && is_array($data['distances'][0]) ? $data['distances'][0] : array();

        $items = array();
        foreach ($ids as $index => $id) {
            $distance = isset($distances[$index]) ? floatval($distances[$index]) : 1.0;
            $score = max(0, min(1, 1 - $distance));

            $items[] = array(
                'id' => (string) $id,
                'content' => isset($documents[$index]) ? (string) $documents[$index] : '',
                'metadata' => isset($metadatas[$index]) && is_array($metadatas[$index]) ? $metadatas[$index] : array(),
                'score' => round($score, 3),
            );
        }

        return array(
            'success' => true,
            'data' => $items,
            'status' => $response['status'] ?? 200,
        );
    }

    /**
     * @inheritDoc
     */
    public function insert_record($record) {
        $id = isset($record['id']) ? (string) $record['id'] : '';
        $document = isset($record['document'])
            ? (string) $record['document']
            : (isset($record['content']) ? (string) $record['content'] : '');
        $metadata = isset($record['metadata']) && is_array($record['metadata']) ? $record['metadata'] : array();

        if ($id === '' || $document === '') {
            return array(
                'success' => false,
                'message' => __('Record ID and document are required.', 'kitsune-ai-platform'),
            );
        }

        // Support richer record shape used by KB sync payloads.
        if (isset($record['source_id']) && !array_key_exists('source_id', $metadata)) {
            $metadata['source_id'] = intval($record['source_id']);
        }
        if (array_key_exists('post_id', $record) && !array_key_exists('post_id', $metadata)) {
            if (!($record['post_id'] === null || $record['post_id'] === '')) {
                $metadata['post_id'] = intval($record['post_id']);
            }
        }
        if (isset($record['model']) && !array_key_exists('model', $metadata)) {
            $metadata['model'] = sanitize_text_field($record['model']);
        }
        $metadata = $this->sanitize_metadata_for_chroma($metadata);

        $collection_id = $this->resolve_collection_id();
        if (is_wp_error($collection_id)) {
            return array(
                'success' => false,
                'message' => $collection_id->get_error_message(),
            );
        }

        // cromdbapidoc requires embeddings in AddCollectionRecordsPayload.
        $embedding = !empty($record['embedding']) && is_array($record['embedding'])
            ? $record['embedding']
            : null;

        if (!is_array($embedding) || empty($embedding)) {
            return array(
                'success' => false,
                'message' => __('Embedding vector is required to insert records.', 'kitsune-ai-platform'),
            );
        }

        $payload = array(
            'ids' => array($id),
            'documents' => array($document),
            'metadatas' => array($metadata),
            'embeddings' => array($embedding),
        );

        return $this->post('/collections/' . rawurlencode($collection_id) . '/add', $payload);
    }

    /**
     * Chroma metadata values must be scalar; null/arrays/objects are rejected.
     *
     * @param array $metadata Raw metadata.
     * @return array
     */
    private function sanitize_metadata_for_chroma($metadata) {
        if (!is_array($metadata)) {
            return array();
        }

        $sanitized = array();
        foreach ($metadata as $key => $value) {
            $meta_key = sanitize_key($key);
            if ($meta_key === '' || $value === null) {
                continue;
            }

            if (is_bool($value) || is_int($value) || is_float($value)) {
                $sanitized[$meta_key] = $value;
                continue;
            }

            if (is_string($value)) {
                $sanitized[$meta_key] = $value;
                continue;
            }

            // Preserve request validity by flattening non-scalar values to strings.
            $sanitized[$meta_key] = wp_json_encode($value);
        }

        return $sanitized;
    }

    /**
     * @inheritDoc
     */
    public function delete_record($record_id, $args = array()) {
        $id = (string) $record_id;
        if ($id === '') {
            return array(
                'success' => false,
                'message' => __('Record ID is required.', 'kitsune-ai-platform'),
            );
        }

        $collection_id = $this->resolve_collection_id();
        if (is_wp_error($collection_id)) {
            return array(
                'success' => false,
                'message' => $collection_id->get_error_message(),
            );
        }

        $payload = array(
            'ids' => array($id),
        );

        return $this->post('/collections/' . rawurlencode($collection_id) . '/delete', $payload);
    }

    /**
     * @inheritDoc
     *
     * Creates a brand-new collection with a unique name so that dimension
     * constraints from a previous embedding model do not carry over.
     * The old collection is left untouched in Chroma.
     */
    public function reset_collection() {
        // Build a unique name based on the configured base name + current timestamp.
        $new_name = $this->collection . '_' . time();

        $response = $this->post('/collections', array(
            'name'           => $new_name,
            'get_or_create'  => false,
        ));

        if (empty($response['success'])) {
            return array(
                'success' => false,
                'message' => $response['message'] ?? __('Failed to create new Chroma collection.', 'kitsune-ai-platform'),
            );
        }

        $new_id = isset($response['data']['id']) ? trim((string) $response['data']['id']) : '';
        if ($new_id === '') {
            return array(
                'success' => false,
                'message' => __('New Chroma collection ID missing in response.', 'kitsune-ai-platform'),
            );
        }

        // Update internal state so subsequent calls on this instance use the new collection.
        $this->collection    = $new_name;
        $this->collection_id = $new_id;

        return array(
            'success'      => true,
            'collection'   => $new_name,
            'collectionId' => $new_id,
        );
    }

    /**
     * Resolve collection UUID for v2 endpoints.
     *
     * @return string|WP_Error
     */
    private function resolve_collection_id() {
        if ($this->collection_id !== '') {
            return $this->collection_id;
        }

        $payload = array(
            'name' => $this->collection,
            'get_or_create' => true,
        );

        $response = $this->post('/collections', $payload);
        if (empty($response['success'])) {
            return new WP_Error(
                'chroma_collection_resolve_failed',
                $response['message'] ?? __('Failed to resolve Chroma collection.', 'kitsune-ai-platform')
            );
        }

        $collection_id = isset($response['data']['id']) ? trim((string) $response['data']['id']) : '';
        if ($collection_id === '') {
            return new WP_Error(
                'chroma_collection_id_missing',
                __('Chroma collection ID is missing in response.', 'kitsune-ai-platform')
            );
        }

        $this->collection_id = $collection_id;
        return $this->collection_id;
    }

    /**
     * POST wrapper that uses the documented v2 route base.
     *
     * @param string $path_v2 V2 path suffix from /collections onward.
     * @param array  $payload JSON payload.
     * @return array
     */
    private function post($path_v2, $payload) {
        $response = wp_remote_post(
            $this->build_url($path_v2),
            array(
                'headers' => $this->build_headers(),
                'body' => wp_json_encode($payload),
                'timeout' => $this->timeout,
            )
        );

        return $this->decode_response($response);
    }

    /**
     * Build request headers.
     *
     * @return array
     */
    private function build_headers() {
        $headers = array(
            'Content-Type' => 'application/json',
            'Accept' => 'application/json',
        );

        if ($this->api_key !== '') {
            $headers['x-chroma-token'] = $this->api_key;
        }

        return $headers;
    }

    /**
     * Build v2 URL.
     *
     * @param string $path_v2 V2 path suffix.
     * @return string
     */
    private function build_url($path_v2) {
        return $this->v2_base_url . $path_v2;
    }

    /**
     * Resolve effective v2 base URL and optionally infer tenant/database from endpoint path.
     *
     * Supported endpoint formats:
     * - https://api.trychroma.com
     * - https://api.trychroma.com/api/v2
     * - https://api.trychroma.com/api/v2/tenants/{tenant}/databases/{database}
     * - https://api.trychroma.com/api/v2/tenants/{tenant}/databases/{database}/collections...
     *
     * @return string
     */
    private function resolve_v2_base_url() {
        $parts = wp_parse_url($this->endpoint);
        if (!is_array($parts) || empty($parts['host'])) {
            return rtrim($this->endpoint, '/') . '/api/v2/tenants/' . rawurlencode($this->tenant_id)
                . '/databases/' . rawurlencode($this->database);
        }

        $scheme = isset($parts['scheme']) ? $parts['scheme'] : 'http';
        $origin = $scheme . '://' . $parts['host'];
        if (!empty($parts['port'])) {
            $origin .= ':' . intval($parts['port']);
        }

        $path = isset($parts['path']) ? rtrim($parts['path'], '/') : '';

        // Endpoint already contains tenant/database path; prefer values from URL.
        if (preg_match('#^(.*/api/v2/tenants/([^/]+)/databases/([^/]+))(?:/.*)?$#', $path, $matches)) {
            $this->tenant_id = rawurldecode($matches[2]);
            $this->database = rawurldecode($matches[3]);
            return $origin . $matches[1];
        }

        // Endpoint points to /api/v2 root.
        if (preg_match('#^(.*/api/v2)$#', $path, $matches)) {
            return $origin . $matches[1] . '/tenants/' . rawurlencode($this->tenant_id)
                . '/databases/' . rawurlencode($this->database);
        }

        // Endpoint is plain origin (or custom path); append canonical v2 route.
        return rtrim($this->endpoint, '/') . '/api/v2/tenants/' . rawurlencode($this->tenant_id)
            . '/databases/' . rawurlencode($this->database);
    }

    /**
     * Parse and normalize remote responses.
     *
     * @param mixed $response wp_remote_* response.
     * @return array
     */
    private function decode_response($response) {
        if (is_wp_error($response)) {
            return array(
                'success' => false,
                'message' => $response->get_error_message(),
            );
        }

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

        if ($code < 200 || $code >= 300) {
            $message = __('Chroma request failed.', 'kitsune-ai-platform');
            if (is_array($data)) {
                $message = $data['message'] ?? $data['error'] ?? $message;
            }
            return array(
                'success' => false,
                'message' => $message,
                'status' => $code,
            );
        }

        // Allow any valid JSON on a 2xx response (Chroma DELETE endpoints may return null or a scalar).
        if ($body !== '' && json_last_error() !== JSON_ERROR_NONE) {
            return array(
                'success' => false,
                'message' => __('Invalid Chroma response payload.', 'kitsune-ai-platform'),
                'status' => $code,
            );
        }

        return array(
            'success' => true,
            'data' => is_array($data) ? $data : array(),
            'status' => $code,
        );
    }
}
