<?php

class MeowPro_MWAI_Google extends Meow_MWAI_Engines_Google {
  private $streamFunctionCall = null;
  private $streamFunctionCalls = [];

  public function __construct( $core, $env ) {
    parent::__construct( $core, $env );
  }

  private function reset_stream() {
    $this->streamContent = '';
    $this->streamBuffer = '';
    $this->streamFunctionCall = null;
    $this->streamFunctionCalls = [];
  }

  public function stream_handler( $handle, $args, $url ) {
    curl_setopt( $handle, CURLOPT_SSL_VERIFYPEER, false );
    curl_setopt( $handle, CURLOPT_SSL_VERIFYHOST, false );

    curl_setopt( $handle, CURLOPT_WRITEFUNCTION, function ( $curl, $data ) {

      $length = strlen( $data );
      $this->streamTemporaryBuffer .= $data;
      $this->streamBuffer .= $data;   // for error checks
      $this->stream_error_check( $this->streamBuffer );

      // ---------------- Find full JSON objects -----------------
      $buf = $this->streamTemporaryBuffer;
      $pos = 0;
      $depth = 0;
      $inStr = false;
      $escape = false;
      $start = null;
      $objects = [];

      while ( $pos < strlen( $buf ) ) {
        $ch = $buf[ $pos ];

        // Handle string state
        if ( $inStr ) {
          if ( $escape ) {
            $escape = false;
          }
          elseif ( $ch === '\\' ) {
            $escape = true;
          }
          elseif ( $ch === '"' ) {
            $inStr = false;
          }
          $pos++;
          continue;
        }

        if ( $ch === '"' ) {
          $inStr = true;
          $pos++;
          continue;
        }

        // Handle brace counting for object detection
        if ( $ch === '{' ) {
          if ( $depth === 0 ) {
            $start = $pos;
          }
          $depth++;
        }
        elseif ( $ch === '}' ) {
          $depth--;
          if ( $depth === 0 && $start !== null ) {
            $jsonChunk = substr( $buf, $start, $pos - $start + 1 );

            // Attempt to decode
            $json = json_decode( $jsonChunk, true );
            if ( json_last_error() === JSON_ERROR_NONE ) {
              $objects[] = $json;

              // Skip trailing spaces / commas / newlines
              $commaPos = $pos + 1;
              while ( $commaPos < strlen( $buf ) && in_array( $buf[ $commaPos ], [ ' ', "\n", "\r", ',' ] ) ) {
                $commaPos++;
              }
              // Trim processed part from buffer and restart scanning
              $buf = substr( $buf, $commaPos );
              $pos = -1;      // will be ++ to 0 at loop bottom
              $start = null;
            }
            else {
              // JSON still incomplete – wait for more bytes
              break;
            }
          }
        }
        $pos++;
      }

      // Keep the unprocessed tail for next callback
      $this->streamTemporaryBuffer = $buf;

      // --------------- Forward each decoded object --------------
      foreach ( $objects as $obj ) {

        // Handle all parts, not just the first one
        if ( isset( $obj['candidates'][0]['content']['parts'] ) ) {
          foreach ( $obj['candidates'][0]['content']['parts'] as $part ) {
            $delta = [ 'role' => 'assistant' ];
            
            if ( isset( $part['functionCall'] ) ) {
              $delta['function_call'] = $part['functionCall'];
            }
            if ( isset( $part['text'] ) ) {
              // Check if this is a thought part (for thinking mode)
              if ( isset( $part['thought'] ) && $part['thought'] === true ) {
                // Emit thinking event
                if ( $this->streamCallback ) {
                  $event = new Meow_MWAI_Event( 'live', MWAI_STREAM_TYPES['THINKING'] );
                  $event->set_content( $part['text'] );
                  call_user_func( $this->streamCallback, $event );
                }
                // Don't add thoughts to the main content
                continue;
              } else {
                $delta['content'] = $part['text'];
              }
            }

            $mapped = [
              'choices' => [ [ 'delta' => $delta ] ],
            ];

            $content = $this->stream_data_handler( $mapped );
            if ( !is_null( $content ) ) {
              if ( $content === "\n" ) {
                $content = "  \n";
              }
              $this->streamContent .= $content;
              call_user_func( $this->streamCallback, $content );
            }
          }
        }
      }

      return $length;
    } );
  }

  protected function stream_data_handler( $json ) {
    if ( !isset( $json['choices'][0]['delta'] ) ) {
      return null;
    }
    $choice = $json['choices'][0];
    $delta = $choice['delta'];
    // Capture a function-call if the model sends one.
    if ( isset( $delta['function_call'] ) ) {
      $this->streamFunctionCall = $delta['function_call'];
      $this->streamFunctionCalls[] = $delta['function_call'];
      
      // Emit function_calling event
      if ( $this->currentDebugMode && $this->streamCallback ) {
        $functionName = $delta['function_call']['name'] ?? 'unknown';
        $functionArgs = isset( $delta['function_call']['args'] ) ? 
          json_encode( $delta['function_call']['args'] ) : '';
        
        $event = Meow_MWAI_Event::function_calling( $functionName, $functionArgs );
        call_user_func( $this->streamCallback, $event );
      }
    }
    if ( isset( $delta['content'] ) && $delta['content'] !== '' ) {
      return $delta['content'];
    }
    return null;
  }

  public function try_decode_error( $data ) {
    $json = json_decode( $data, true );
    if ( isset( $json['error']['message'] ) ) {
      return $json['error']['message'];
    }
    return null;
  }

  public function run_completion_query( $query, $streamCallback = null ): Meow_MWAI_Reply {
    $isStreaming = !is_null( $streamCallback );
    
    // Initialize debug mode
    $this->init_debug_mode( $query );

    if ( $isStreaming ) {
      $this->streamCallback = $streamCallback;
      $this->reset_stream();
    }

    // Build body using the parent's build_body method which handles event emission
    $body = $this->build_body( $query, $streamCallback );

    $base = $this->endpoint . '/models/' . $query->model;
    if ( $isStreaming ) {
      $url = $base . ':streamGenerateContent?key=' . $this->apiKey;
    }
    else {
      $url = $base . ':generateContent?key=' . $this->apiKey;
    }

    if ( $isStreaming ) {
      // Emit "Request sent" event for feedback queries
      if ( $this->currentDebugMode && !empty( $streamCallback ) && 
           ( $query instanceof Meow_MWAI_Query_Feedback || $query instanceof Meow_MWAI_Query_AssistFeedback ) ) {
        $event = Meow_MWAI_Event::request_sent()
          ->set_metadata( 'is_feedback', true )
          ->set_metadata( 'feedback_count', count( $query->blocks ) );
        call_user_func( $streamCallback, $event );
      }
      
      $ch = curl_init();
      curl_setopt_array( $ch, [
        CURLOPT_URL => $url,
        CURLOPT_POST => true,
        CURLOPT_HTTPHEADER => [ 'Content-Type: application/json', 'Accept: text/event-stream' ],
        CURLOPT_POSTFIELDS => json_encode( $body ),
        CURLOPT_HTTP_VERSION => CURL_HTTP_VERSION_2_0
      ] );
      $this->stream_handler( $ch, [], $url );
      curl_exec( $ch );
      curl_close( $ch );

      if ( empty( $this->streamContent ) ) {
        $error = $this->try_decode_error( $this->streamBuffer );
        if ( !is_null( $error ) ) {
          throw new Exception( $error );
        }
      }

      $reply = new Meow_MWAI_Reply( $query );
      
      // If we have multiple function calls, return them in Google's format
      $returned_choices = [];
      
      // Add each function call as a separate choice
      if ( !empty( $this->streamFunctionCalls ) ) {
        foreach ( $this->streamFunctionCalls as $function_call ) {
          $returned_choices[] = [
            'message' => [
              'content' => null,
              'function_call' => $function_call
            ]
          ];
        }
      }
      
      // Add text content if present
      if ( !empty( $this->streamContent ) ) {
        if ( empty( $returned_choices ) ) {
          // No function calls, just text
          $returned_choices[] = [ 'message' => [ 'content' => $this->streamContent ] ];
        } else {
          // Add text as a separate choice
          $returned_choices[] = [ 'role' => 'assistant', 'text' => $this->streamContent ];
        }
      }
      
      // If we still have no choices, add a single function call if available
      if ( empty( $returned_choices ) && $this->streamFunctionCall ) {
        $returned_choices[] = [ 'message' => [ 'content' => null, 'function_call' => $this->streamFunctionCall ] ];
      }
      
      $reply->set_choices( $returned_choices );
      $this->handle_tokens_usage( $reply, $query, $query->model, null, null );
      return $reply;
    }

    // Emit "Request sent" event for feedback queries (non-streaming)
    if ( !$isStreaming && $this->currentDebugMode && !empty( $streamCallback ) && 
         ( $query instanceof Meow_MWAI_Query_Feedback || $query instanceof Meow_MWAI_Query_AssistFeedback ) ) {
      $event = Meow_MWAI_Event::request_sent()
        ->set_metadata( 'is_feedback', true )
        ->set_metadata( 'feedback_count', count( $query->blocks ) );
      call_user_func( $streamCallback, $event );
    }
    
    $headers = $this->build_headers( $query );
    $options = $this->build_options( $headers, $body );
    $res = $this->run_query( $url, $options );
    $reply = new Meow_MWAI_Reply( $query );
    $data = $res['data'];
    if ( empty( $data ) ) {
      throw new Exception( 'No content received (res is null).' );
    }
    $returned_choices = [];
    if ( isset( $data['candidates'] ) ) {
      foreach ( $data['candidates'] as $candidate ) {
        $content = $candidate['content'];
        
        // Check if there are any parts with function calls
        $functionCalls = [];
        $textContent = '';
        
        if ( isset( $content['parts'] ) ) {
          foreach ( $content['parts'] as $part ) {
            if ( isset( $part['functionCall'] ) ) {
              $functionCalls[] = $part['functionCall'];
              
              // Emit function calling event if debug mode is enabled
              if ( $this->currentDebugMode && !empty( $streamCallback ) ) {
                $functionName = $part['functionCall']['name'] ?? 'unknown';
                $functionArgs = isset( $part['functionCall']['args'] ) ? json_encode( $part['functionCall']['args'] ) : '';
                
                $event = Meow_MWAI_Event::function_calling( $functionName, $functionArgs );
                call_user_func( $streamCallback, $event );
              }
            }
            elseif ( isset( $part['text'] ) ) {
              $textContent .= $part['text'];
            }
          }
        }
        
        // If we have function calls, return them in Google's expected format
        if ( !empty( $functionCalls ) ) {
          // Process each function call separately to ensure all are handled
          foreach ( $functionCalls as $function_call ) {
            $returned_choices[] = [
              'message' => [
                'content' => null,
                'function_call' => $function_call
              ]
            ];
          }
        }
        
        // Add text content if present (separate from function calls)
        if ( !empty( $textContent ) ) {
          $returned_choices[] = [ 'role' => 'assistant', 'text' => $textContent ];
        }
      }
    }
    // Create a proper Google-formatted rawMessage
    $googleRawMessage = null;
    if ( isset( $data['candidates'][0]['content'] ) ) {
      $googleRawMessage = $data['candidates'][0]['content'];
    }
    
    $reply->set_choices( $returned_choices, $googleRawMessage );
    $this->handle_tokens_usage( $reply, $query, $query->model, null, null );
    return $reply;
  }

  public function run_embedding_query( Meow_MWAI_Query_Base $query ) {
    // For experimental models, we might need to use a different approach
    // For now, let's use the model as specified
    $modelName = $query->model;
    
    // Build the URL for embeddings
    $url = $this->endpoint . '/models/' . $modelName . ':embedContent';
    if ( strpos( $url, '?' ) === false ) {
      $url .= '?key=' . $this->apiKey;
    }
    else {
      $url .= '&key=' . $this->apiKey;
    }

    // Build the request body
    $body = [
      'content' => [
        'parts' => [
          [ 'text' => $query->get_message() ]
        ]
      ]
    ];

    $headers = $this->build_headers( $query );
    $options = $this->build_options( $headers, $body );
    
    try {
      // Debug logging
      if ( $this->core->get_option( 'queries_debug_mode' ) ) {
        error_log( '[AI Engine] Google Embedding Request URL: ' . $url );
        error_log( '[AI Engine] Google Embedding Request Body: ' . json_encode( $body ) );
      }
      
      $res = $this->run_query( $url, $options );
      $data = $res['data'];
      
      // Debug logging
      if ( $this->core->get_option( 'queries_debug_mode' ) ) {
        // Don't log the full embedding response, just the structure
        if ( isset( $data['embedding']['values'] ) && is_array( $data['embedding']['values'] ) ) {
          error_log( '[AI Engine] Google Embedding Response: Received ' . count( $data['embedding']['values'] ) . ' dimensions' );
        } else {
          error_log( '[AI Engine] Google Embedding Response: ' . json_encode( $data ) );
        }
      }
      
      // Check if we have an error response
      if ( isset( $data['error'] ) ) {
        $errorMsg = isset( $data['error']['message'] ) ? $data['error']['message'] : 'Unknown error';
        $errorCode = isset( $data['error']['code'] ) ? $data['error']['code'] : 'N/A';
        throw new Exception( "Google API Error (Code: {$errorCode}): {$errorMsg}" );
      }
      
      if ( !isset( $data['embedding']['values'] ) ) {
        throw new Exception( 'No embedding values in the response. Response: ' . json_encode( $data ) );
      }
      
      $embedding = $data['embedding']['values'];
      
      // Handle matryoshka truncation if dimensions are specified
      if ( $query->dimensions && $query->dimensions < count( $embedding ) ) {
        // Google Gemini embedding models support matryoshka (dimension truncation)
        // We can safely truncate to the requested dimensions
        $embedding = array_slice( $embedding, 0, $query->dimensions );
        
        if ( $this->core->get_option( 'queries_debug_mode' ) ) {
          error_log( "[AI Engine] Truncated embedding from " . count( $data['embedding']['values'] ) . " to {$query->dimensions} dimensions (matryoshka)" );
        }
      }
      
      $reply = new Meow_MWAI_Reply( $query );
      $reply->type = 'embedding';
      $reply->result = $embedding;
      $reply->results = [ $embedding ];
      
      // Record usage (Google doesn't provide token counts for embeddings)
      $this->handle_tokens_usage( $reply, $query, $query->model, null, null );
      
      return $reply;
    }
    catch ( Exception $e ) {
      Meow_MWAI_Logging::error( '(Google) Embedding error: ' . $e->getMessage() );
      throw new Exception( 'Google Embedding Error: ' . $e->getMessage() );
    }
  }
}
