<?php

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

if ( ! class_exists( 'AIEO_Batch_Runner' ) ) {
	require_once dirname( __DIR__ ) . '/class-aieo-batch-runner.php';
}

class AIEO_Embedding_Backfill extends AIEO_Batch_Runner {

	const RUNNER_ID        = 'embed_backfill';
	const HOOK_TICK        = 'aieo_embed_backfill_tick';
	const BATCH_SIZE       = 50;
	const PROGRESS_OPTION  = 'aieo_embed_backfill_progress';
	const RUNNING_OPTION   = 'aieo_embed_backfill_running';
	
	const ALLOWED_TYPES    = array( 'post', 'page', 'product' );

	
	public static function allowed_types() {
		$learning = class_exists( 'AIEO_Content_Types' ) ? AIEO_Content_Types::learning_post_types() : array();
		return array_values( array_unique( array_merge( self::ALLOWED_TYPES, $learning ) ) );
	}

	public static function bootstrap() {
		static::bootstrap_runner(); 
		if ( is_admin() ) {
			add_action( 'wp_ajax_aieo_embeddings_estimate', array( __CLASS__, 'ajax_estimate' ) );
			add_action( 'wp_ajax_aieo_embeddings_backfill_start', array( __CLASS__, 'ajax_start' ) );
			add_action( 'wp_ajax_aieo_embeddings_backfill_status', array( __CLASS__, 'ajax_status' ) );
		}
	}

	
	
	

	
	public static function estimate( array $post_types, $engine_type = null, $deduplicated = 0 ) {
		$per_type   = array();
		$items      = 0;
		$est_tokens = 0;
		if ( ! class_exists( 'AI_eShop_Chat_Database_Manager' ) ) {
			return array( 'items' => 0, 'est_tokens' => 0, 'est_usd' => 0.0, 'remaining_usd' => null, 'per_type' => array() );
		}
		$db = AI_eShop_Chat_Database_Manager::get_instance();
		foreach ( self::sanitize_types( $post_types ) as $pt ) {
			$n     = (int) $db->count_posts_without_embeddings( $pt, $engine_type, (int) $deduplicated );
			$tpi   = self::tokens_per_item( $pt );
			$toks  = $n * $tpi;
			$per_type[ $pt ] = array( 'items' => $n, 'tokens_per_item' => $tpi, 'est_tokens' => $toks );
			$items      += $n;
			$est_tokens += $toks;
		}
		$rate_per_m = (float) apply_filters( 'aieo_embedding_usd_per_million_tokens', 0.10, $engine_type );
		$est_usd    = round( ( $est_tokens / 1000000 ) * $rate_per_m, 4 );

		$remaining = null;
		if ( class_exists( 'AIEO_Credit_Notifier' ) && method_exists( 'AIEO_Credit_Notifier', 'remaining_usd' ) ) {
			$remaining = AIEO_Credit_Notifier::remaining_usd();
		}
		return array(
			'items'         => $items,
			'est_tokens'    => $est_tokens,
			'est_usd'       => $est_usd,
			'remaining_usd' => is_numeric( $remaining ) ? round( (float) $remaining, 2 ) : null,
			'per_type'      => $per_type,
		);
	}

	
	private static function tokens_per_item( $post_type ) {
		$defaults = array( 'post' => 1200, 'page' => 1800, 'product' => 800 );
		$base     = isset( $defaults[ $post_type ] ) ? $defaults[ $post_type ] : 1000;
		return (int) apply_filters( 'aieo_embedding_tokens_per_item', $base, $post_type );
	}

	
	
	

	
	public static function start( array $post_types, $engine_type = null, $deduplicated = 0 ) {
		$types = self::sanitize_types( $post_types );
		if ( empty( $types ) ) {
			return array( 'started' => false, 'post_types' => array(), 'reason' => 'no_valid_types' );
		}
		$streams = array();
		foreach ( $types as $pt ) {
			$streams[] = array( $pt, (string) $engine_type, (int) $deduplicated );
		}
		$res = static::start_streams( $streams );
		if ( $res['started'] ) {
			
			$running = get_option( self::RUNNING_OPTION );
			$running = is_array( $running ) ? $running : array();
			$running['post_types'] = $types;
			$running['engine']     = (string) $engine_type;
			$running['dedup']      = (int) $deduplicated;
			update_option( self::RUNNING_OPTION, $running, false );
		}
		return array( 'started' => $res['started'], 'post_types' => $types, 'reason' => $res['reason'] );
	}

	
	public static function tick( $post_type = 'post', $engine_type = '', $deduplicated = 0 ) {
		static::run_tick( array( (string) $post_type, (string) $engine_type, (int) $deduplicated ) );
	}

	
	protected static function do_batch( array $args ) {
		$post_type    = isset( $args[0] ) ? (string) $args[0] : 'post';
		$engine_type  = isset( $args[1] ) ? (string) $args[1] : '';
		$deduplicated = isset( $args[2] ) ? (int) $args[2] : 0;

		if ( ! in_array( $post_type, self::allowed_types(), true ) || ! class_exists( 'AI_eShop_Chat_Embedding_Manager' ) ) {
			return array( 'success' => false, 'remaining' => 0, 'total' => 0 );
		}
		$engine = ( '' === $engine_type ) ? null : $engine_type;
		$mgr    = AI_eShop_Chat_Embedding_Manager::get_instance();
		$res    = $mgr->generate_all_embeddings( $post_type, self::BATCH_SIZE, 0, $deduplicated, $engine );

		$remaining = isset( $res['remaining'] ) ? (int) $res['remaining'] : 0;
		$total     = isset( $res['total_items'] ) ? (int) $res['total_items'] : ( isset( $res['total'] ) ? (int) $res['total'] : 0 );
		return array(
			'success'   => ! empty( $res['success'] ),
			'remaining' => $remaining,
			'total'     => $total,
		);
	}

	
	public static function status() {
		return static::runner_status();
	}

	
	
	

	public static function ajax_estimate() {
		self::guard();
		$types  = self::posted_types();
		$engine = self::posted_engine();
		$dedup  = isset( $_POST['deduplicated'] ) ? (int) $_POST['deduplicated'] : 0;
		wp_send_json_success( self::estimate( $types, $engine, $dedup ) );
	}

	public static function ajax_start() {
		self::guard();
		$types  = self::posted_types();
		$engine = self::posted_engine();
		$dedup  = isset( $_POST['deduplicated'] ) ? (int) $_POST['deduplicated'] : 0;
		wp_send_json_success( self::start( $types, $engine, $dedup ) );
	}

	public static function ajax_status() {
		self::guard();
		wp_send_json_success( self::status() );
	}

	private static function guard() {
		if ( ! current_user_can( 'manage_options' ) ) {
			wp_send_json_error( array( 'message' => 'forbidden' ), 403 );
		}
		check_ajax_referer( 'aieo_embed_backfill', 'nonce' );
	}

	private static function posted_types() {
		$raw = isset( $_POST['post_types'] ) ? wp_unslash( $_POST['post_types'] ) : array();
		if ( is_string( $raw ) ) {
			$raw = array_map( 'trim', explode( ',', $raw ) );
		}
		return self::sanitize_types( (array) $raw );
	}

	private static function posted_engine() {
		$e = isset( $_POST['engine_type'] ) ? sanitize_text_field( wp_unslash( $_POST['engine_type'] ) ) : '';
		return ( '' === $e ) ? null : $e;
	}

	private static function sanitize_types( array $types ) {
		$out = array();
		foreach ( $types as $t ) {
			$t = sanitize_key( (string) $t );
			if ( in_array( $t, self::allowed_types(), true ) && ! in_array( $t, $out, true ) ) {
				$out[] = $t;
			}
		}
		return $out;
	}
}