<?php
/**
 * Costudio_Search_Engine
 *
 * External web search providers (Tavily, Serper, Exa). Returns a markdown
 * blob ready to be prepended to an LLM prompt as research context.
 *
 * Native search (toggle "Use ContentOS Real-time Search") is handled by
 * Costudio_Search_Engine_ContentOS. The Pipeline picks the engine; this
 * class only knows about external providers.
 *
 * @package Contentosapp_Studio
 * @since   3.0.0
 */

defined( 'ABSPATH' ) || exit;

class Costudio_Search_Engine {

	/** Allowed external providers. Single source of truth. */
	const ALLOWED_PROVIDERS = array( 'tavily', 'serper', 'exa', 'gemini' );

	/**
	 * Providers that return STRUCTURED {url,title,snippet} sources (so the pipeline
	 * persists `sources_json` + builds the `## INDEXED SOURCES` block, exactly like
	 * native search). Gemini is excluded: its grounding returns Vertex-redirect
	 * strings the plugin does not yet normalise client-side (a planned follow-up),
	 * so it falls back to raw context.
	 */
	const STRUCTURED_PROVIDERS = array( 'tavily', 'serper', 'exa' );

	private string $provider;
	private string $api_key;
	private string $model;
	private int $max_results;

	public function __construct() {
		$cts = function_exists( 'costudio_settings' ) ? costudio_settings() : null;

		$provider       = $cts ? (string) $cts->get( 'search_api_provider', 'serper' ) : 'serper';
		$this->provider = in_array( $provider, self::ALLOWED_PROVIDERS, true ) ? $provider : 'serper';
		$this->api_key  = $cts ? (string) $cts->get( 'search_api_key', '' ) : '';

		// If Gemini is selected for search but no specific search key is provided, 
		// try to use the main AI API key (since they are usually the same for Google).
		if ( 'gemini' === $this->provider && empty( $this->api_key ) ) {
			$this->api_key = $cts ? (string) $cts->get( 'ai_api_key', '' ) : '';
		}
		$this->model       = self::normalize_provider_model( $this->provider, $cts ? (string) $cts->get( 'search_api_model', '' ) : '' );
		$raw_max           = $cts ? $cts->get( 'search_max_results', 5 ) : 5;
		$this->max_results = class_exists( 'Costudio_Standalone_Settings' )
			? Costudio_Standalone_Settings::clamp_search_results( $raw_max )
			: max( 1, min( 20, (int) $raw_max ) );
	}

	/**
	 * Returns the list of providers currently supported (for UI rendering and
	 * for sanitisation in settings save handlers).
	 *
	 * @return string[]
	 */
	public static function allowed_providers(): array {
		return self::ALLOWED_PROVIDERS;
	}

	/**
	 * Whether the given provider returns structured sources (citable {{ref:N}} pool)
	 * vs raw-context-only (Gemini). Used by the pipeline to skip the legacy prose
	 * URL scraper when structured grounding already populated `sources_json`.
	 */
	public static function provider_is_structured( string $provider ): bool {
		return in_array( $provider, self::STRUCTURED_PROVIDERS, true );
	}

	/**
	 * Execute a search and return a formatted string suitable for LLM context.
	 *
	 * @param  string $query
	 * @return string  Formatted search results, or '' when no key is configured.
	 */
	public function search( string $query ): string {
		$result = $this->search_structured( $query );
		return (string) ( $result['fallback_md'] ?? '' );
	}

	/**
	 * Execute a search and return BOTH structured sources and a fallback string.
	 *
	 * Mirrors the shape of Costudio_Search_Engine_ContentOS::search() so the
	 * pipeline can persist `sources_json` and build the numbered `## INDEXED
	 * SOURCES` block for BYOK providers exactly like the native path. Providers
	 * that don't expose clean structured sources (Gemini grounding → Vertex
	 * redirects) return empty `sources` + a populated `fallback_md` (raw context),
	 * preserving today's behaviour with no regression.
	 *
	 * @param  string $query
	 * @return array{ok:bool,sources:array<int,array{url:string,title:string,snippet:string}>,summary_markdown:string,fallback_md:string,provider?:string,model?:string,model_used?:string}
	 */
	public function search_structured( string $query ): array {
		if ( empty( $this->api_key ) ) {
			return self::empty_result();
		}

		switch ( $this->provider ) {
			case 'tavily':
				return $this->search_tavily( $query );
			case 'exa':
				return $this->search_exa( $query );
			case 'gemini':
				return $this->search_gemini( $query );
			case 'serper':
			default:
				return $this->search_serper( $query );
		}
	}

	private static function empty_result(): array {
		return array(
			'ok'               => false,
			'sources'          => array(),
			'summary_markdown' => '',
			'fallback_md'      => '',
		);
	}

	// ── Providers ────────────────────────────────────────────────────────────

	private function search_tavily( string $query ): array {
		$response = wp_remote_post(
			'https://api.tavily.com/search',
			array(
				'timeout' => 25,
				'headers' => array( 'Content-Type' => 'application/json' ),
				'body'    => wp_json_encode(
					array(
						'api_key'        => $this->api_key,
						'query'          => $query,
						'search_depth'   => $this->model ?: 'advanced',
						'max_results'    => $this->max_results,
						'include_answer' => true,
					)
				),
			)
		);

		if ( is_wp_error( $response ) ) {
			return self::empty_result();
		}
		$code = (int) wp_remote_retrieve_response_code( $response );
		if ( $code < 200 || $code >= 300 ) {
			return self::empty_result();
		}
		$data    = json_decode( wp_remote_retrieve_body( $response ), true );
		$sources = array();
		$parts   = array();
		if ( ! empty( $data['answer'] ) ) {
			$parts[] = "--- DIRECT SUMMARY ---\n" . $data['answer'];
		}
		foreach ( ( $data['results'] ?? array() ) as $item ) {
			$title   = (string) ( $item['title'] ?? '' );
			$url     = (string) ( $item['url'] ?? '' );
			$content = (string) ( $item['content'] ?? '' );
			if ( '' !== trim( $url ) ) {
				$sources[] = array( 'url' => $url, 'title' => $title, 'snippet' => mb_substr( $content, 0, 500 ) );
			}
			$parts[] = "**[{$title}]({$url})**\n{$content}";
		}
		return array(
			'ok'               => true,
			'sources'          => $sources,
			'summary_markdown' => isset( $data['answer'] ) ? (string) $data['answer'] : '',
			'fallback_md'      => implode( "\n\n---\n\n", $parts ),
			'provider'         => 'tavily',
			'model'            => $this->model ?: 'advanced',
			'model_used'       => 'tavily:' . ( $this->model ?: 'advanced' ),
		);
	}

	private function search_exa( string $query ): array {
		$mode = self::normalize_exa_mode( $this->model );
		$response = wp_remote_post(
			'https://api.exa.ai/search',
			array(
				'timeout' => 25,
				'headers' => array(
					'x-api-key'    => $this->api_key,
					'Content-Type' => 'application/json',
				),
				'body'    => wp_json_encode(
					array(
						'query'      => $query,
						'type'       => $mode,
						'numResults' => $this->max_results,
						'contents'   => array(
							'highlights' => true,
							'text'       => array( 'maxCharacters' => 1000 ),
						),
					)
				),
			)
		);

		if ( is_wp_error( $response ) ) {
			return self::empty_result();
		}
		$code = (int) wp_remote_retrieve_response_code( $response );
		if ( $code < 200 || $code >= 300 ) {
			return self::empty_result();
		}
		$data    = json_decode( wp_remote_retrieve_body( $response ), true );
		$sources = array();
		$parts   = array();
		foreach ( ( $data['results'] ?? array() ) as $item ) {
			$title   = (string) ( $item['title'] ?? '' );
			$url     = (string) ( $item['url'] ?? '' );
			$snippet = mb_substr( self::snippet_from_exa_result( is_array( $item ) ? $item : array() ), 0, 500 );
			if ( '' !== trim( $url ) ) {
				$sources[] = array( 'url' => $url, 'title' => $title, 'snippet' => $snippet );
			}
			$parts[] = "**[{$title}]({$url})**\n" . $snippet . '...';
		}
		return array(
			'ok'               => true,
			'sources'          => $sources,
			'summary_markdown' => '',
			'fallback_md'      => implode( "\n\n---\n\n", $parts ),
			'provider'         => 'exa',
			'model'            => $mode,
			'model_used'       => 'exa:' . $mode,
		);
	}

	private function search_serper( string $query ): array {
		list( $gl, $hl ) = self::serper_locale_params();

		$endpoint = 'https://google.serper.dev/search';
		if ( $this->model === 'news' ) {
			$endpoint = 'https://google.serper.dev/news';
		} elseif ( $this->model === 'places' ) {
			$endpoint = 'https://google.serper.dev/places';
		}

		$response = wp_remote_post(
			$endpoint,
			array(
				'timeout' => 20,
				'headers' => array(
					'X-API-KEY'    => $this->api_key,
					'Content-Type' => 'application/json',
				),
				'body'    => wp_json_encode(
					array(
						'q'   => $query,
						'num' => $this->max_results,
						'gl'  => $gl,
						'hl'  => $hl,
					)
				),
			)
		);

		if ( is_wp_error( $response ) ) {
			return self::empty_result();
		}
		$code = (int) wp_remote_retrieve_response_code( $response );
		if ( $code < 200 || $code >= 300 ) {
			return self::empty_result();
		}
		$data    = json_decode( wp_remote_retrieve_body( $response ), true );
		$sources = array();
		$parts   = array();
		$items   = ( $this->model === 'news' ) ? ( $data['news'] ?? array() ) : ( $data['organic'] ?? array() );
		foreach ( $items as $item ) {
			$title   = (string) ( $item['title'] ?? '' );
			$snippet = (string) ( $item['snippet'] ?? '' );
			$link    = (string) ( $item['link'] ?? '' );
			if ( '' !== trim( $link ) ) {
				$sources[] = array( 'url' => $link, 'title' => $title, 'snippet' => $snippet );
			}
			$parts[] = "**{$title}**\n{$snippet}\nURL: {$link}";
		}
		return array(
			'ok'               => true,
			'sources'          => $sources,
			'summary_markdown' => '',
			'fallback_md'      => implode( "\n\n---\n\n", $parts ),
			'provider'         => 'serper',
			'model'            => $this->model ?: 'search',
			'model_used'       => 'serper:' . ( $this->model ?: 'search' ),
		);
	}

	private function search_gemini( string $query ): array {
		if ( ! class_exists( 'Costudio_Search_Engine_Gemini' ) ) {
			return self::empty_result();
		}
		// Pass the resolved Web Search key/model down — the Gemini engine must NOT
		// fall back to the text provider's ai_api_key (wrong key when text != Gemini).
		// It now returns the SAME structured shape: grounding chunks → sources (with
		// Vertex redirects resolved). Empty sources → fallback_md raw context (free
		// tier / no billing), which the pipeline handles.
		$engine = new Costudio_Search_Engine_Gemini();
		$result = $engine->search( $query, $this->model, $this->api_key );
		$result['provider'] = 'gemini';
		if ( empty( $result['model'] ) ) {
			$result['model'] = $this->model ?: 'gemini-2.5-flash';
		}
		if ( empty( $result['model_used'] ) ) {
			$result['model_used'] = (string) $result['model'];
		}
		return $result;
	}

	private static function normalize_provider_model( string $provider, string $model ): string {
		switch ( $provider ) {
			case 'exa':
				return self::normalize_exa_mode( $model );
			case 'tavily':
				$model = strtolower( trim( $model ) );
				return in_array( $model, array( 'basic', 'advanced' ), true ) ? $model : 'advanced';
			case 'serper':
				$model = strtolower( trim( $model ) );
				return in_array( $model, array( 'search', 'news' ), true ) ? $model : 'search';
			default:
				return trim( $model );
		}
	}

	private static function normalize_exa_mode( string $mode ): string {
		$mode = strtolower( trim( $mode ) );
		if ( 0 === strpos( $mode, 'exa:' ) ) {
			$mode = substr( $mode, 4 );
		}
		if ( in_array( $mode, array( 'neural', 'keyword' ), true ) ) {
			$mode = 'auto';
		}
		return in_array( $mode, array( 'auto', 'fast', 'deep' ), true ) ? $mode : 'auto';
	}

	private static function snippet_from_exa_result( array $item ): string {
		$highlights = $item['highlights'] ?? array();
		if ( is_array( $highlights ) && ! empty( $highlights ) ) {
			$parts = array();
			foreach ( $highlights as $highlight ) {
				$highlight = trim( (string) $highlight );
				if ( '' !== $highlight ) {
					$parts[] = $highlight;
				}
			}
			if ( ! empty( $parts ) ) {
				return implode( ' ', $parts );
			}
		}

		$summary = trim( (string) ( $item['summary'] ?? '' ) );
		if ( '' !== $summary ) {
			return $summary;
		}

		return (string) ( $item['text'] ?? '' );
	}

	/**
	 * Map the WordPress locale to Serper's gl (country) / hl (language) params.
	 * Avoids hard-coding pt-BR for users running EN/ES sites.
	 *
	 * @return array{0:string,1:string} [gl, hl]
	 */
	private static function serper_locale_params(): array {
		$locale = function_exists( 'get_locale' ) ? get_locale() : 'en_US';
		$lang   = strtolower( substr( $locale, 0, 2 ) );
		$region = strtolower( substr( $locale, 3, 2 ) );

		$gl = $region ?: ( $lang === 'pt' ? 'br' : ( $lang === 'es' ? 'es' : 'us' ) );
		$hl = $lang ?: 'en';

		return array( $gl, $hl );
	}
}
