<?php

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

class Niwa_AI_Latency_Tracker {
    const OPTION_KEY = Niwa_AI_Constants::OPT_LATENCY_EVENTS;

	private $active_traces = array();

	public function is_frontend_stage_timing_enabled() {
		$settings = get_option( Niwa_AI_Constants::OPT_SETTINGS, array() );
		$settings = is_array( $settings ) ? $settings : array();
		$enabled  = ! array_key_exists( 'enable_frontend_stage_timing', $settings ) || ! empty( $settings['enable_frontend_stage_timing'] );

		return function_exists( 'apply_filters' )
			? (bool) apply_filters( 'niwa_ai_frontend_stage_timing_enabled', $enabled )
			: $enabled;
	}

	public function start_trace( $request_id, $meta = array() ) {
		$request_id = $this->sanitize_request_id( $request_id );
		if ( '' === $request_id || ! $this->is_frontend_stage_timing_enabled() ) {
			return false;
		}

		$this->active_traces[ $request_id ] = array(
			'started_ns' => $this->now_ns(),
			'stages'     => array(),
			'meta'       => is_array( $meta ) ? $meta : array(),
		);

		return true;
	}

	public function start_stage( $request_id ) {
		$request_id = $this->sanitize_request_id( $request_id );
		return isset( $this->active_traces[ $request_id ] ) ? $this->now_ns() : 0;
	}

	public function has_active_trace( $request_id ) {
		$request_id = $this->sanitize_request_id( $request_id );
		return '' !== $request_id && isset( $this->active_traces[ $request_id ] );
	}

	public function finish_stage( $request_id, $stage, $started_ns, $meta = array() ) {
		$started_ns = (int) $started_ns;
		if ( $started_ns <= 0 ) {
			return 0.0;
		}

		$duration_ms = max( 0, $this->now_ns() - $started_ns ) / 1000000;
		$this->add_stage( $request_id, $stage, $duration_ms, $meta );

		return $duration_ms;
	}

	public function add_stage( $request_id, $stage, $duration_ms, $meta = array() ) {
		$request_id = $this->sanitize_request_id( $request_id );
		$stage      = sanitize_key( (string) $stage );
		if ( '' === $request_id || '' === $stage || ! isset( $this->active_traces[ $request_id ] ) ) {
			return false;
		}

		$duration_ms = round( max( 0, (float) $duration_ms ), 2 );
		if ( ! isset( $this->active_traces[ $request_id ]['stages'][ $stage ] ) ) {
			$this->active_traces[ $request_id ]['stages'][ $stage ] = array(
				'duration_ms' => 0.0,
				'count'       => 0,
				'samples'     => array(),
			);
		}

		$entry                 =& $this->active_traces[ $request_id ]['stages'][ $stage ];
		$entry['duration_ms']   = round( (float) $entry['duration_ms'] + $duration_ms, 2 );
		$entry['count']         = (int) $entry['count'] + 1;
		if ( count( $entry['samples'] ) < 12 ) {
			$sample = array( 'duration_ms' => $duration_ms );
			if ( ! empty( $meta ) && is_array( $meta ) ) {
				$sample['meta'] = $meta;
			}
			$entry['samples'][] = $sample;
		}
		unset( $entry );

		return true;
	}

	public function finish_trace( $request_id, $meta = array() ) {
		$request_id = $this->sanitize_request_id( $request_id );
		if ( '' === $request_id || ! isset( $this->active_traces[ $request_id ] ) ) {
			return false;
		}

		$trace            = $this->active_traces[ $request_id ];
		$processing_ms    = max( 0, $this->now_ns() - (int) $trace['started_ns'] ) / 1000000;
		$this->add_stage( $request_id, 'request_processing', $processing_ms );
		$trace            = $this->active_traces[ $request_id ];
		$queue_wait_ms    = isset( $trace['stages']['queue_wait']['duration_ms'] ) ? (float) $trace['stages']['queue_wait']['duration_ms'] : 0.0;
		$preprocessing_ms = isset( $trace['stages']['request_preprocessing']['duration_ms'] ) ? (float) $trace['stages']['request_preprocessing']['duration_ms'] : 0.0;
		$total_ms         = round( $preprocessing_ms + $queue_wait_ms + $processing_ms, 2 );
		$trace_meta                   = array_merge( $trace['meta'], is_array( $meta ) ? $meta : array() );
		$trace_meta['request_id']     = $request_id;
		$trace_meta['nested_stages']  = true;
		$trace_meta['stages']         = $trace['stages'];
		unset( $this->active_traces[ $request_id ] );

		return $this->record( 'frontend_chat_trace', $total_ms, $trace_meta );
	}

    public function record( $stage, $duration_ms, $meta = array() ) {
		$lock_token = $this->acquire_events_lock();
		if ( '' === $lock_token ) {
			return false;
		}

		try {
			$events   = get_option( self::OPTION_KEY, array() );
			$events   = is_array( $events ) ? $events : array();
			$events[] = array(
				'ts'          => time(),
				'stage'       => sanitize_text_field( $stage ),
				'duration_ms' => round( floatval( $duration_ms ), 2 ),
				'meta'        => is_array( $meta ) ? $meta : array(),
			);

			$events = array_slice( $events, -Niwa_AI_Constants::LATENCY_MAX_EVENTS );
			return (bool) update_option( self::OPTION_KEY, $events, false );
		} finally {
			$this->release_events_lock( $lock_token );
		}
    }

	private function sanitize_request_id( $request_id ) {
		return substr( preg_replace( '/[^a-zA-Z0-9_\-:.]/', '', trim( (string) $request_id ) ), 0, 96 );
	}

	private function now_ns() {
		return function_exists( 'hrtime' ) ? (int) hrtime( true ) : (int) round( microtime( true ) * 1000000000 );
	}

	private function acquire_events_lock() {
		$token = function_exists( 'wp_generate_uuid4' ) ? (string) wp_generate_uuid4() : uniqid( 'latency_', true );
		$key   = Niwa_AI_Constants::OPT_LATENCY_EVENTS_LOCK;

		for ( $attempt = 0; $attempt < 10; $attempt++ ) {
			$record = array(
				'token'      => $token,
				'expires_at' => microtime( true ) + 3,
			);
			if ( add_option( $key, $record, '', false ) ) {
				return $token;
			}

			$current = get_option( $key, array() );
			if ( is_array( $current ) && isset( $current['expires_at'] ) && (float) $current['expires_at'] < microtime( true ) ) {
				delete_option( $key );
				continue;
			}
			usleep( 20000 );
		}

		return '';
	}

	private function release_events_lock( $token ) {
		$key     = Niwa_AI_Constants::OPT_LATENCY_EVENTS_LOCK;
		$current = get_option( $key, array() );
		if ( is_array( $current ) && isset( $current['token'] ) && hash_equals( (string) $current['token'], (string) $token ) ) {
			delete_option( $key );
		}
	}
}
