<?php

/**
 * Vvveb
 *
 * Copyright (C) 2022  Ziadin Givan
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 */

namespace Vvveb\Plugins\AiAssistant\Controller;

use function Vvveb\getSetting;
use Vvveb\Controller\Base;

/**
 * AI proxy - relays AI provider requests through the site backend so API keys
 * never reach the browser.
 *
 * Protocol (POST json to /admin/index.php?module=plugins/ai-assistant/proxy):
 *   {provider: 'openai'|'anthropic'|'google', body: <provider native body>, csrf: <token>}
 *
 * The client-side adapter builds the provider-native request body, the proxy
 * injects credentials from plugin settings, re-addresses it to the whitelisted
 * upstream endpoint and streams the raw response (SSE or JSON) back unchanged
 * together with an X-Vvveb-Ai-Format header so the client picks the right
 * parser. No prompt or key material is logged.
 */
class Proxy extends Base {
	/**
	 * Whitelisted providers -> [default endpoint, auth style].
	 * Client supplied urls are ignored; only these endpoints are called.
	 */
	protected array $providers = [
		'openai'     => ['url' => 'https://api.openai.com/v1/chat/completions',       'auth' => 'bearer', 'models' => 'https://api.openai.com/v1/models'],
		'anthropic'  => ['url' => 'https://api.anthropic.com/v1/messages',            'auth' => 'header', 'models' => 'https://api.anthropic.com/v1/models'],
		'google'     => ['url' => 'https://generativelanguage.googleapis.com/v1beta', 'auth' => 'header', 'models' => 'https://generativelanguage.googleapis.com/v1beta/models'],
		// openai chat-completions compatible providers
		'deepseek'   => ['url' => 'https://api.deepseek.com/v1/chat/completions',     'auth' => 'bearer', 'models' => 'https://api.deepseek.com/v1/models'],
		'groq'       => ['url' => 'https://api.groq.com/openai/v1/chat/completions',  'auth' => 'bearer', 'models' => 'https://api.groq.com/openai/v1/models'],
		'mistral'    => ['url' => 'https://api.mistral.ai/v1/chat/completions',       'auth' => 'bearer', 'models' => 'https://api.mistral.ai/v1/models'],
		'xai'        => ['url' => 'https://api.x.ai/v1/chat/completions',             'auth' => 'bearer', 'models' => 'https://api.x.ai/v1/models'],
		'openrouter' => ['url' => 'https://openrouter.ai/api/v1/chat/completions',    'auth' => 'bearer', 'models' => 'https://openrouter.ai/api/v1/models'],
		'opencode'   => ['url' => 'https://opencode.ai/zen/v1/chat/completions',      'auth' => 'bearer', 'models' => 'https://opencode.ai/zen/v1/models'],
		'nim'        => ['url' => 'https://integrate.api.nvidia.com/v1/chat/completions', 'auth' => 'bearer', 'models' => 'https://integrate.api.nvidia.com/v1/models'],
	];

	// fallback model lists when provider has no list endpoint or key missing
	protected array $fallbackModels = [
		'openai' => ['gpt-4o-mini','gpt-4o','gpt-4-turbo','o1-mini','o1-preview','o3-mini'],
		'anthropic' => ['claude-sonnet-4-20250514','claude-3-5-sonnet-20241022','claude-3-opus-20240229','claude-3-haiku-20240307'],
		'google' => ['gemini-2.0-flash','gemini-1.5-pro','gemini-1.5-flash','gemini-2.0-flash-thinking'],
		'deepseek' => ['deepseek-chat','deepseek-coder','deepseek-reasoner'],
		'groq' => ['llama-3.3-70b-versatile','mixtral-8x7b-32768','gemma2-9b-it','llama-3.1-8b-instant'],
		'mistral' => ['mistral-large-latest','mistral-small-latest','open-mistral-nemo','codestral-latest'],
		'xai' => ['grok-3-latest','grok-2-latest','grok-beta'],
		'openrouter' => ['openai/gpt-4o-mini','anthropic/claude-3.5-sonnet','google/gemini-2.0-flash','deepseek/deepseek-chat'],
		'opencode' => ['qwen3-coder','deepseek-r1','llama-3.3-70b','gpt-4o-mini'],
		'nim' => ['meta/llama-3.3-70b-instruct','mistralai/mistral-large','nvidia/llama-3.1-nemotron-70b'],
	];

	protected function jsonResponse($data): void {
		header('Content-Type: application/json');
		echo json_encode($data);
	}

	protected function error(string $message, int $status = 400) {
		http_response_code($status);
		$this->jsonResponse(['error' => ['message' => $message]]);
		die();
	}

	public function index() {
		if (($this->request->getMethod()) !== 'post') {
			$this->error('POST required', 405);
		}

		$input   = json_decode(file_get_contents('php://input'), true);
		$payload = is_array($input) ? $input : [];

		// same check as Base::checkCsrf but against the json payload
		$csrf = $this->session->get('csrf');
		if (! $csrf || ($payload['csrf'] ?? '') !== $csrf) {
			$this->error('Invalid csrf!', 403);
		}

		$provider = strtolower($payload['provider'] ?? 'opencode');
		if (! isset($this->providers[$provider])) {
			$this->error('Unsupported provider: ' . $provider);
		}
		$body = $payload['body'] ?? null;
		if (! is_array($body)) {
			$this->error('Missing request body');
		}

		$config = getSetting('ai-assistant') ?: [];

		// per-provider key falling back to the generic key setting
		$key = $config[$provider . '_key'] ?? ($config['key'] ?? '');
		/*
		if ($key === '') {
			$this->error("No API key configured for \"$provider\"! Add it in the AI Assistant plugin settings page.", 400);
		}*/

		[$url, $headers] = $this->upstream($provider, $body, $config, $key);

		// correlate the upstream request with the chat session: header name is
		// configurable (default x-vvveb-session), empty disables. The id arrives
		// in the proxy envelope, never in the provider body.
		$sessionHeader = preg_replace('/[^A-Za-z0-9-]/', '', (string)($config['session_header'] ?? 'x-vvveb-session'));
		$sessionId = preg_replace('/[\r\n]/', '', substr((string)($payload['session_id'] ?? ''), 0, 128));
		if ($sessionHeader !== '' && $sessionId !== '') {
			$headers[] = $sessionHeader . ': ' . $sessionId;
		}

		$this->relay($url, $headers, $body, $provider);
	}

	public function models() {
		if (($this->request->getMethod()) !== 'post') {
			$this->error('POST required', 405);
		}
		$input = json_decode(file_get_contents('php://input'), true);
		// keys travel in the POST body only, never in the URL (server access logs)
		if (! is_array($input) || ! isset($input['provider'])){
			$this->error('Missing provider');
		}
		$csrf = $this->session->get('csrf');
		if (! $csrf || ($input['csrf'] ?? '') !== $csrf) {
			$this->error('Invalid csrf!', 403);
		}
		$provider = strtolower($input['provider'] ?? 'openai');
		if (! isset($this->providers[$provider])) {
			$this->error('Unsupported provider: ' . $provider);
		}
		$config = getSetting('ai-assistant') ?: [];
		$key = $config[$provider . '_key'] ?? ($config['key'] ?? '');
		// allow key sent from form (unsaved settings) to override stored key
		$inputKey = trim($input['key'] ?? '');
		if ($inputKey !== '') $key = $inputKey;
		// also check provider-specific key field from form
		$inputProviderKey = trim($input[$provider . '_key'] ?? '');
		if ($inputProviderKey !== '') $key = $inputProviderKey;
		// if still no key, try fallback list but also try without auth for providers that allow it (e.g., openrouter, opencode)
		if ($key === '' && !in_array($provider, ['openrouter','opencode'])) {
			$this->jsonResponse(['data' => $this->fallbackModels[$provider] ?? []]);
			return;
		}
		$def = $this->providers[$provider];
		$url = $def['models'] ?? '';
		// handle custom openai url override
		if ($provider === 'openai' && !empty($config['url'])) {
			$url = preg_replace('#/v1/chat/completions/?$#', '/v1/models', $config['url']) ?: $url;
			if (!preg_match('#/v1/models#', $url)) $url = rtrim($config['url'], '/') . '/models';
		}
		$headers = ['Content-Type: application/json'];
		if ($provider === 'anthropic'){
			$headers[] = 'x-api-key: ' . $key;
			$headers[] = 'anthropic-version: 2023-06-01';
		} elseif ($provider === 'google'){
			// google uses key as query param, not header - append to url
			$url .= (str_contains($url, '?') ? '&' : '?') . 'key=' . urlencode($key);
		} else {
			$headers[] = 'Authorization: Bearer ' . $key;
			$headers[] = 'User-Agent: opencode/1.18.23';
			//if ($provider !== 'openai') $headers[] = 'User-Agent: opencode/1.18.23';
		}

		
		$this->session->close();
		$ch = curl_init($url);
		curl_setopt_array($ch, [
			CURLOPT_HTTPHEADER => $headers,
			CURLOPT_RETURNTRANSFER => true,
			CURLOPT_CONNECTTIMEOUT => 8,
			CURLOPT_TIMEOUT => 12,
			CURLOPT_ENCODING => '',
		]);
		$resp = curl_exec($ch);
		$httpCode = curl_getinfo($ch, CURLINFO_HTTP_CODE);
		$err = curl_error($ch);
		//curl_close($ch);
		if ($resp === false || $httpCode >= 400){
			// on failure return fallback list with error hint
			if ($httpCode >= 400){
				$data = json_decode($resp, true);
				$msg = $data['error']['message'] ?? $data['error'] ?? $err ?: 'Upstream error ' . $httpCode;
				$this->jsonResponse(['error' => $msg, 'data' => $this->fallbackModels[$provider] ?? []]);
			} else {
				$this->jsonResponse(['error' => $err ?: 'Failed to fetch models', 'data' => $this->fallbackModels[$provider] ?? []]);
			}
			return;
		}
		$data = json_decode($resp, true);
		// normalize to {data:[...]}
		$models = [];
		if (isset($data['data']) && is_array($data['data'])) $models = $data['data'];
		elseif (isset($data['models']) && is_array($data['models'])) $models = $data['models'];
		elseif (is_array($data)) $models = $data;
		else $models = $this->fallbackModels[$provider] ?? [];
		// google returns {models:[{name:"models/gemini-..."}]}
		if ($provider === 'google' && isset($models[0]['name'])){
			$models = array_map(fn($m) => str_replace('models/', '', $m['name'] ?? $m), $models);
		}
		$this->jsonResponse(['data' => $models]);
	}

	/**
	 * Image generation - relays to OpenAI compatible /v1/images/generations API.
	 *
	 * Protocol (POST json to /admin/index.php?module=plugins/ai-assistant/proxy&action=images):
	 *   {provider, prompt, model?, size?, n?, csrf, key?}
	 * Returns the upstream JSON unchanged: {created, data:[{url|b64_json, revised_prompt}]}.
	 */
	public function images() {
		if (! in_array($this->request->getMethod(), ['post'])) {
			$this->error('POST required', 405);
		}
		$input = json_decode(file_get_contents('php://input'), true);
		if (! is_array($input)) {
			$input = $this->request->post;
		}
		$csrf = $this->session->get('csrf');
		if (! $csrf || ($input['csrf'] ?? '') !== $csrf) {
			$this->error('Invalid csrf!', 403);
		}
		$provider = strtolower($input['provider'] ?? 'openai');
		if (! isset($this->providers[$provider])) {
			$this->error('Unsupported provider: ' . $provider);
		}
		$prompt = trim($input['prompt'] ?? '');
		if ($prompt === '') {
			$this->error('Missing prompt');
		}
		if (strlen($prompt) > 4000) {
			$this->error('Prompt too long (max 4000 chars)');
		}
		$config = getSetting('ai-assistant') ?: [];
		$key = $config[$provider . '_key'] ?? ($config['key'] ?? '');
		$inputKey = trim($input['key'] ?? '');
		if ($inputKey !== '') {
			$key = $inputKey;
		}
		$model = preg_replace('/[^A-Za-z0-9._\-:\/]/', '', $input['model'] ?? ($config['image_model'] ?? 'dall-e-3')) ?: 'dall-e-3';
		$size  = preg_replace('/[^0-9x]/', '', $input['size'] ?? ($config['image_size'] ?? '1024x1024')) ?: '1024x1024';
		$n     = max(1, min(4, (int)($input['n'] ?? 1)));
		$base = rtrim($this->providers[$provider]['url'], '/');
		// providers store chat completions url - swap tail for images endpoint
		if (str_ends_with($base, '/chat/completions')) {
			$url = substr($base, 0, -strlen('/chat/completions')) . '/images/generations';
		} else {
		$this->error('Image generation is only supported by OpenAI compatible providers (openai, openrouter, nim, xai, groq, mistral, deepseek, opencode)');
		}
		if ($provider === 'openai' && ! empty($config['url']) && str_contains($config['url'], '/images/')) {
			$url = $config['url'];
		}
		$headers = ['Content-Type: application/json', 'Authorization: Bearer ' . $key];
		if (in_array($provider, ['openrouter', 'opencode'])) {
			$headers[] = 'User-Agent: opencode/1.18.23';
		}
		$body = ['model' => $model, 'prompt' => $prompt, 'size' => $size, 'n' => $n];
		$this->session->close();
		$ch = curl_init($url);
		curl_setopt_array($ch, [
			CURLOPT_POST           => true,
			CURLOPT_POSTFIELDS     => json_encode($body),
			CURLOPT_HTTPHEADER     => $headers,
			CURLOPT_RETURNTRANSFER => true,
			CURLOPT_CONNECTTIMEOUT => 10,
			CURLOPT_TIMEOUT        => 120,
			CURLOPT_ENCODING       => '',
		]);
		$resp = curl_exec($ch);
		$httpCode = curl_getinfo($ch, CURLINFO_HTTP_CODE);
		$err = curl_error($ch);
		if ($resp === false) {
			$this->error('Proxy request failed: ' . $err, 502);
		}
		if ($httpCode >= 400) {
			http_response_code($httpCode);
			header('Content-Type: application/json');
			echo $resp;
			exit;
		}
		header('Content-Type: application/json');
		header('Cache-Control: no-cache');
		echo $resp;
		exit;
	}

	/**
	 * Builds the upstream url + headers for a provider.
	 */
	protected function upstream(string $provider, array &$body, array $config, string $key): array {
		$def     = $this->providers[$provider];
		$headers = ['Content-Type: application/json'];

		switch ($provider) {
			case 'openai':
				// handle Responses API for custom tools (input + custom tools -> /v1/responses)
				$hasCustom = false;
				if (isset($body['tools']) && is_array($body['tools'])) {
					foreach ($body['tools'] as $t) if (($t['type'] ?? '') === 'custom') { $hasCustom = true; break; }
				}
				if ($hasCustom && isset($body['input'])) {
					// honor a custom OpenAI-compatible base url when configured
					// instead of always hitting api.openai.com
					$url = 'https://api.openai.com/v1/responses';
					$custom = trim($config['url'] ?? '');
					if ($custom !== '') {
						if (preg_match('#/responses/?$#', $custom)) {
							$url = $custom;
						} elseif (preg_match('#/chat/completions/?$#', $custom)) {
							$url = preg_replace('#/chat/completions/?$#', '/responses', $custom);
						} elseif (preg_match('#/completions/?$#', $custom)) {
							$url = preg_replace('#/completions/?$#', '/responses', $custom);
						} else {
							$url = rtrim($custom, '/') . '/responses';
						}
					}
					$headers[] = 'Authorization: Bearer ' . $key;
					break;
				}
				$url = $config['url'] ?? '';
				// rewrite legacy /v1/completions saved by older versions
				if (! $url || preg_match('#/v1/completions/?$#', $url)) {
					$url = preg_replace('#/v1/completions/?$#', '/v1/chat/completions', $url) ?: $def['url'];
				}
				$headers[] = 'Authorization: Bearer ' . $key;
				break;

			case 'anthropic':
				$url = $def['url'];
				$headers[] = 'x-api-key: ' . $key;
				$headers[] = 'anthropic-version: 2023-06-01';
				break;

			case 'google':
				// model goes into the url path - strip everything but id characters
				// to prevent path traversal on the upstream api
				$model   = preg_replace('/[^A-Za-z0-9._-]/', '', $body['model'] ?? 'gemini-2.0-flash');
				// collapse dot segments so traversal sequences can not survive
				while (str_contains($model, '..')) { $model = str_replace('..', '.', $model); }
				$model   = trim($model, '.');
				$method  = ! empty($body['stream']) ? 'streamGenerateContent?alt=sse' : 'generateContent';
				$url     = $def['url'] . '/models/' . ($model !== '' ? $model : 'gemini-2.0-flash') . ':' . $method;
				$headers[] = 'x-goog-api-key: ' . $key;
				break;

			default:
				// openai-compatible providers (deepseek, groq, mistral, xai,
				// openrouter, opencode, nim) - own endpoint + bearer auth
				$url = $def['url'];
				//$key = '';
				$headers[] = 'User-Agent: opencode/1.18.23';
				$headers[] = 'Authorization: Bearer ' . $key;
				break;
		}

		return [$url, $headers];
	}

	/**
	 * Forwards the body and streams the raw upstream response back.
	 */
	protected function relay(string $url, array $headers, array $body, ?string $provider = null): void {
		$this->session->close();
		
		$ch = curl_init($url);

		$started     = false;
		$contentType = '';
		$statusLine  = '';

		$curlOpts = [
			CURLOPT_POST           => true,
			CURLOPT_POSTFIELDS     => json_encode($body),
			CURLOPT_HTTPHEADER     => $headers,
			CURLOPT_ENCODING       => '',             // decode gzip so we forward plain bytes
			CURLOPT_CONNECTTIMEOUT => 10,
			CURLOPT_TIMEOUT        => 600,            // streams can run long, but never forever
		];
		// low-speed abort only for non-streaming calls: SSE streams legitimately
		// go silent (reasoning gaps, big-context time-to-first-token), so the
		// <32 B/s guard would kill healthy runs ("Operation too slow"). Truly
		// dead streams are caught by the client's stall watchdog instead.
		if (empty($body['stream'])) {
			$curlOpts[CURLOPT_LOW_SPEED_LIMIT] = 32;
			$curlOpts[CURLOPT_LOW_SPEED_TIME]  = 60;
		}
		$curlOpts[CURLOPT_HEADERFUNCTION] = function ($ch, $line) use (&$started, &$contentType, &$statusLine, $provider) {
			$trim = trim($line);
			if (stripos($trim, 'HTTP/') === 0) {
				$statusLine = $trim;
			} elseif (preg_match('/^content-type:\s*(.+)$/i', $trim, $m)) {
				$contentType = trim($m[1]);
			} elseif ($trim === '' && ! $started) {
				// end of upstream headers -> emit our own and start streaming
				$this->streamHeaders($statusLine, $contentType, $provider);
				$started = true;
			}
			return strlen($line);
		};
		$curlOpts[CURLOPT_WRITEFUNCTION] = function ($ch, $chunk) use (&$started) {
			if (! $started) {
				$this->streamHeaders('', '', '');
				$started = true;
			}
			while (ob_get_level() > 0) { ob_end_flush(); }
			echo $chunk;
			flush();
			return strlen($chunk);
		};
		curl_setopt_array($ch, $curlOpts);

		$success = curl_exec($ch);

		if ($success === false) {
			if (! $started) {
				$this->error('Proxy request failed: ' . curl_error($ch), 502);
			}
		}

		//curl_close($ch);
		exit;
	}

	/**
	 * Sends streaming response headers mirroring the upstream status/content type.
	 */
	protected function streamHeaders(string $statusLine, string $contentType, ?string $provider = null): void {
		if (preg_match('#HTTP/[\d.]+\s+(\d+)#', $statusLine, $m)) {
			http_response_code((int)$m[1]);
		}
		header('Content-Type: ' . ($contentType ?: 'text/event-stream'));
		header('Cache-Control: no-cache');
		//header('X-Accel-Buffering: no'); // disable nginx buffering
		/*
		if ($provider) {
			header('X-Vvveb-Ai-Format: ' . $provider);
		}*/
		while (ob_get_level() > 0) { ob_end_flush(); }
	}
}
