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

class Arcai_Chat_Controller {

	const MAX_HISTORY_TURNS  = 6;
	const MAX_MESSAGE_LENGTH = 1000;
	const RETRIEVE_LIMIT     = 8;

	private $provider;
	private $store;
	private $limiter;

	public function __construct() {
		$this->provider = new Arcai_Provider_OpenAI();
		$this->store    = new Arcai_Vector_Store();
		$this->limiter  = new Arcai_Rate_Limiter();
	}

	public function register_routes() {
		$args = array(
			'message' => array(
				'required' => true,
				'type'     => 'string',
			),
			'history' => array(
				'required' => false,
				'type'     => 'array',
			),
		);

		register_rest_route(
			'archivist-ai/v1',
			'/chat',
			array(
				'methods'             => WP_REST_Server::CREATABLE,
				'callback'            => array( $this, 'handle_chat' ),
				'permission_callback' => '__return_true',
				'args'                => $args,
			)
		);

		register_rest_route(
			'archivist-ai/v1',
			'/chat/stream',
			array(
				'methods'             => WP_REST_Server::CREATABLE,
				'callback'            => array( $this, 'handle_chat_stream' ),
				'permission_callback' => '__return_true',
				'args'                => $args,
			)
		);
	}

	/* ---------------------------------------------------------------------
	 * Endpoints
	 * ------------------------------------------------------------------- */

	public function handle_chat( WP_REST_Request $request ) {
		$input = $this->guards( $request );
		if ( is_wp_error( $input ) ) {
			return $input;
		}

		if ( $input['cache_key'] ) {
			$cached = get_transient( $input['cache_key'] );
			if ( false !== $cached ) {
				return rest_ensure_response( $cached );
			}
		}

		$chunks = $this->retrieve( $input['message'], $input['history'] );
		if ( is_wp_error( $chunks ) ) {
			return $this->api_failure( $chunks );
		}
		if ( empty( $chunks ) ) {
			return rest_ensure_response( $this->no_results_payload() );
		}

		$prompt = $this->build_prompt( $input['message'], $input['history'], $chunks );

		$result = $this->provider->chat( $prompt['messages'], array( 'max_tokens' => 500 ) );
		if ( is_wp_error( $result ) ) {
			return $this->api_failure( $result );
		}
		Arcai_Usage::log( 'chat_answer', $result['tokens_in'], $result['tokens_out'], $result['cost'] );

		$payload = $this->finish( $result['content'], $prompt['sources'] );
		if ( $input['cache_key'] ) {
			set_transient( $input['cache_key'], $payload, 6 * HOUR_IN_SECONDS );
		}

		return rest_ensure_response( $payload );
	}

	/**
	 * SSE variant: same pipeline, but answer tokens are relayed to the
	 * browser as they arrive from the provider. Frames:
	 *   event: delta   data: {"t": "text fragment"}
	 *   event: sources data: [ ...cited sources ]
	 *   event: done    data: {}
	 *   event: error   data: {"message": "..."}
	 * Guard failures return ordinary JSON errors (the client checks the
	 * response content type before attempting to read a stream).
	 */
	public function handle_chat_stream( WP_REST_Request $request ) {
		$input = $this->guards( $request );
		if ( is_wp_error( $input ) ) {
			return $input;
		}

		$this->sse_begin();

		if ( $input['cache_key'] ) {
			$cached = get_transient( $input['cache_key'] );
			if ( false !== $cached ) {
				$this->sse_send( 'delta', array( 't' => $cached['answer'] ) );
				$this->sse_send( 'sources', $cached['sources'] );
				$this->sse_send( 'done', array() );
				exit;
			}
		}

		$chunks = $this->retrieve( $input['message'], $input['history'] );
		if ( is_wp_error( $chunks ) ) {
			$this->sse_send( 'error', array( 'message' => $this->public_error_message() ) );
			exit;
		}
		if ( empty( $chunks ) ) {
			$payload = $this->no_results_payload();
			$this->sse_send( 'delta', array( 't' => $payload['answer'] ) );
			$this->sse_send( 'sources', array() );
			$this->sse_send( 'done', array() );
			exit;
		}

		$prompt     = $this->build_prompt( $input['message'], $input['history'], $chunks );
		$controller = $this;

		$result = $this->provider->chat_stream(
			$prompt['messages'],
			array( 'max_tokens' => 500 ),
			function ( $fragment ) use ( $controller ) {
				$controller->sse_send( 'delta', array( 't' => $fragment ) );
			}
		);

		if ( is_wp_error( $result ) ) {
			$this->sse_send( 'error', array( 'message' => $this->public_error_message() ) );
			exit;
		}
		Arcai_Usage::log( 'chat_answer', $result['tokens_in'], $result['tokens_out'], $result['cost'] );

		$payload = $this->finish( $result['content'], $prompt['sources'] );
		if ( $input['cache_key'] ) {
			set_transient( $input['cache_key'], $payload, 6 * HOUR_IN_SECONDS );
		}

		$this->sse_send( 'sources', $payload['sources'] );
		$this->sse_send( 'done', array() );
		exit;
	}

	/* ---------------------------------------------------------------------
	 * Shared pipeline
	 * ------------------------------------------------------------------- */

	/**
	 * Rate limit, spend cap, and input validation. Returns WP_Error or
	 * [ 'message' => string, 'history' => array, 'cache_key' => string|false ].
	 */
	private function guards( WP_REST_Request $request ) {
		if ( Arcai_Settings::get( 'require_login' ) && ! is_user_logged_in() ) {
			return new WP_Error(
				'arcai_login_required',
				__( 'Please log in to use the archive search.', 'archivist-ai' ),
				array( 'status' => 401 )
			);
		}

		if ( ! $this->limiter->allow( Arcai_Rate_Limiter::client_ip() ) ) {
			return new WP_Error(
				'arcai_rate_limited',
				__( 'You have reached the search limit for now. Please try again later.', 'archivist-ai' ),
				array( 'status' => 429 )
			);
		}

		if ( Arcai_Usage::over_daily_cap() ) {
			return new WP_Error(
				'arcai_cap_reached',
				__( 'AI search is taking a break today. Please use the regular site search.', 'archivist-ai' ),
				array( 'status' => 503 )
			);
		}

		$message = trim( sanitize_textarea_field( (string) $request->get_param( 'message' ) ) );
		if ( '' === $message ) {
			return new WP_Error( 'arcai_empty', __( 'Please enter a question.', 'archivist-ai' ), array( 'status' => 400 ) );
		}

		$message = mb_substr( $message, 0, self::MAX_MESSAGE_LENGTH );
		$history = $this->sanitize_history( $request->get_param( 'history' ) );

		return array(
			'message'   => $message,
			'history'   => $history,
			'cache_key' => $this->cache_key( $message, $history ),
		);
	}

	/**
	 * Parse the question, embed it, and fetch matching chunks (re-checking
	 * post visibility). Returns chunk rows, an empty array, or WP_Error.
	 */
	private function retrieve( $message, array $history ) {
		$parsed = $this->parse_query( $message, $history );
		if ( is_wp_error( $parsed ) ) {
			return $parsed;
		}

		$embedded = $this->provider->embed( array( $parsed['query'] ) );
		if ( is_wp_error( $embedded ) ) {
			return $embedded;
		}
		Arcai_Usage::log( 'chat_embed', $embedded['tokens'], 0, $embedded['cost'] );

		$chunks = $this->store->search(
			$embedded['embeddings'][0],
			self::RETRIEVE_LIMIT,
			$parsed['date_from'],
			$parsed['date_to']
		);

		return array_values(
			array_filter(
				$chunks,
				function ( $chunk ) {
					$post = get_post( $chunk['post_id'] );
					return $post && 'publish' === $post->post_status && '' === $post->post_password;
				}
			)
		);
	}

	private function build_prompt( $message, array $history, array $chunks ) {
		$sources    = array();
		$source_map = array();
		$context    = '';

		foreach ( $chunks as $chunk ) {
			$post_id = $chunk['post_id'];
			if ( ! isset( $source_map[ $post_id ] ) ) {
				$source_map[ $post_id ] = count( $source_map ) + 1;
				$sources[]              = array(
					'n'     => $source_map[ $post_id ],
					'title' => get_the_title( $post_id ),
					'url'   => get_permalink( $post_id ),
					'date'  => get_the_date( '', $post_id ),
				);
			}
			$n        = $source_map[ $post_id ];
			$date     = get_the_date( 'Y-m-d', $post_id );
			$context .= "[{$n}] \"" . get_the_title( $post_id ) . "\" ({$date}):\n" . $chunk['content'] . "\n\n";
		}

		$site   = get_bloginfo( 'name' );
		$system = "You are Archivist AI, the article-archive assistant for the website \"{$site}\". "
			. 'Answer the user using ONLY the article excerpts provided below. '
			. 'Cite sources inline with their bracket numbers, e.g. [1], [2]. '
			. 'If the excerpts do not contain the answer, say so honestly and suggest a different question. '
			. 'Never invent facts, articles, or dates. Keep answers under 200 words. '
			. "Article excerpts:\n\n" . $context;

		$messages   = array( array( 'role' => 'system', 'content' => $system ) );
		$messages   = array_merge( $messages, $history );
		$messages[] = array( 'role' => 'user', 'content' => $message );

		return array(
			'messages' => $messages,
			'sources'  => $sources,
		);
	}

	private function finish( $content, array $sources ) {
		$cited = array();
		foreach ( $sources as $source ) {
			if ( false !== strpos( $content, '[' . $source['n'] . ']' ) ) {
				$cited[] = $source;
			}
		}

		return array(
			'answer'  => $content,
			'sources' => $cited,
		);
	}

	private function no_results_payload() {
		return array(
			'answer'  => __( "I couldn't find any articles matching that. Try rephrasing your question or broadening the time range.", 'archivist-ai' ),
			'sources' => array(),
		);
	}

	private function sanitize_history( $history ) {
		if ( ! is_array( $history ) ) {
			return array();
		}
		$clean = array();
		foreach ( array_slice( $history, -self::MAX_HISTORY_TURNS ) as $turn ) {
			if ( ! is_array( $turn ) || empty( $turn['role'] ) || ! isset( $turn['content'] ) ) {
				continue;
			}
			$role    = 'assistant' === $turn['role'] ? 'assistant' : 'user';
			$clean[] = array(
				'role'    => $role,
				'content' => mb_substr( sanitize_textarea_field( (string) $turn['content'] ), 0, self::MAX_MESSAGE_LENGTH ),
			);
		}
		return $clean;
	}

	/**
	 * Turn the question (plus chat context) into a standalone search query
	 * and optional date range.
	 */
	private function parse_query( $message, array $history ) {
		$today  = current_time( 'Y-m-d' );
		$system = 'You convert a question about a news/article archive into search parameters. '
			. "Today's date is {$today}. Use the conversation context to resolve follow-up questions into a standalone query. "
			. 'Respond with JSON only: {"query": "<standalone search query>", "date_from": "YYYY-MM-DD or null", "date_to": "YYYY-MM-DD or null"}. '
			. 'Only set dates when the user expresses a time constraint (e.g. "last month", "in 2024", "since January").';

		$messages   = array( array( 'role' => 'system', 'content' => $system ) );
		$messages   = array_merge( $messages, $history );
		$messages[] = array( 'role' => 'user', 'content' => $message );

		$result = $this->provider->chat( $messages, array( 'json' => true, 'max_tokens' => 200 ) );
		if ( is_wp_error( $result ) ) {
			return $result;
		}
		Arcai_Usage::log( 'chat_parse', $result['tokens_in'], $result['tokens_out'], $result['cost'] );

		$data  = json_decode( $result['content'], true );
		$query = is_array( $data ) && ! empty( $data['query'] ) ? (string) $data['query'] : $message;

		return array(
			'query'     => $query,
			'date_from' => $this->valid_date( isset( $data['date_from'] ) ? $data['date_from'] : null ),
			'date_to'   => $this->valid_date( isset( $data['date_to'] ) ? $data['date_to'] : null ),
		);
	}

	private function valid_date( $value ) {
		if ( ! is_string( $value ) || ! preg_match( '/^\d{4}-\d{2}-\d{2}$/', $value ) ) {
			return null;
		}
		return $value;
	}

	/**
	 * Cache key for first-turn questions only — follow-ups depend on
	 * conversation context and are not cacheable. Keyed to the index
	 * generation so re-indexing invalidates all cached answers. Time-relative
	 * phrasing ("last month") resolves differently as days pass, hence the
	 * date component.
	 */
	private function cache_key( $message, array $history ) {
		if ( ! empty( $history ) || ! arcai_is_pro() ) {
			return false;
		}
		$normalized = mb_strtolower( trim( preg_replace( '/\s+/', ' ', $message ) ) );
		$generation = get_option( 'arcai_index_generation', '0' );
		return 'arcai_q_' . md5( $normalized . '|' . $generation . '|' . current_time( 'Y-m-d' ) );
	}

	/* ---------------------------------------------------------------------
	 * SSE plumbing
	 * ------------------------------------------------------------------- */

	private function sse_begin() {
		nocache_headers();
		header( 'Content-Type: text/event-stream; charset=utf-8' );
		header( 'X-Accel-Buffering: no' );
		// phpcs:ignore Squiz.PHP.DiscouragedFunctions.Discouraged -- output compression must be off or SSE frames are buffered instead of streamed.
		@ini_set( 'zlib.output_compression', '0' );

		$levels = ob_get_level();
		for ( $i = 0; $i < $levels; $i++ ) {
			@ob_end_clean();
		}
	}

	public function sse_send( $event, $data ) {
		echo 'event: ' . esc_html( $event ) . "\n";
		echo 'data: ' . wp_json_encode( $data ) . "\n\n";
		if ( ob_get_level() > 0 ) {
			@ob_flush();
		}
		flush();
	}

	private function public_error_message() {
		return __( 'The AI search is temporarily unavailable. Please try again in a moment.', 'archivist-ai' );
	}

	private function api_failure( WP_Error $error ) {
		return new WP_Error(
			'arcai_unavailable',
			$this->public_error_message(),
			array(
				'status' => 502,
				'detail' => current_user_can( 'manage_options' ) ? $error->get_error_message() : null,
			)
		);
	}
}
