<?php

require 'vendor/autoload.php';

use Aws\Credentials\Credentials;
use Aws\Signature\SignatureV4;
use GuzzleHttp\Client;
use GuzzleHttp\Psr7\Request;

class BedrockKnowledgeBaseApp
{
    private $client;
    private $credentials;
    private $region;
    private $db;
    private $config;
    private $user;
    private $profile;

    public function __construct($username)
    {
        $this->config = require 'config.php';
        $this->db = new SQLite3('chat_history.sqlite');
        $this->user = $this->getUser($username);
        $this->profile = $this->getProfile($this->user['profile_id']);

        $this->credentials = new Credentials(
            $this->config['credentials']['key'],
            $this->config['credentials']['secret']
        );

        $this->region = $this->profile['region'] ?? $this->config['region'];
        $this->client = new Client();
    }

    private function getUser($username)
    {
        $stmt = $this->db->prepare('SELECT * FROM users WHERE username = :username');
        $stmt->bindValue(':username', $username, SQLITE3_TEXT);
        $result = $stmt->execute();
        return $result->fetchArray(SQLITE3_ASSOC);
    }

    private function getProfile($profileId)
    {
        $stmt = $this->db->prepare('SELECT * FROM profiles WHERE id = :id');
        $stmt->bindValue(':id', $profileId, SQLITE3_INTEGER);
        $result = $stmt->execute();
        return $result->fetchArray(SQLITE3_ASSOC);
    }

    public function ask($question, $sessionId = '')
    {
        if (!$this->profile) {
            return [
                'error' => true,
                'message' => 'Your profile has been removed. Please contact the admin to create a new one.'
            ];
        }

        $endpoint = "https://bedrock-agent-runtime." . $this->region . ".amazonaws.com/retrieveAndGenerate";
        $body = [
            'input' => [
                'text' => $question
            ],
            'retrieveAndGenerateConfiguration' => [
                'knowledgeBaseConfiguration' => [
                    'generationConfiguration' => [
                        'inferenceConfig' => [
                            'textInferenceConfig' => [
                                'maxTokens' => $this->profile['max_tokens'] ?? $this->config['maxTokens'],
                                'stopSequences' => [$this->profile['stop_sequences'] ?? $this->config['stopSequences']],
                                'temperature' => $this->profile['temperature'] ?? $this->config['temperature'],
                                'topP' => $this->profile['top_p'] ?? $this->config['topP']
                            ]
                        ],
                        'promptTemplate' => [
                            'textPromptTemplate' => $this->profile['instructions'] . $this->getPromptTemplate() . '$search_results$ \n $output_format_instructions$'
                        ]
                    ],
                    'knowledgeBaseId' => $this->profile['knowledge_base_id'] ?? $this->config['knowledgeBaseId'],
                    'modelArn' => $this->profile['model_arn'] ?? $this->config['modelArn'],
                    'retrievalConfiguration' => [
                        'vectorSearchConfiguration' => [
                            'numberOfResults' => $this->profile['number_of_results'] ?? $this->config['numberOfResults'],
                            'overrideSearchType' => $this->profile['override_search_type'] ?? $this->config['overrideSearchType']
                        ]
                    ]
                ],
                'type' => 'KNOWLEDGE_BASE'
            ],
        ];

        debug($body);

        if (!empty($sessionId)) {
            $body['sessionId'] = $sessionId;
        }

        $request = new Request(
            'POST',
            $endpoint,
            [
                'Content-Type' => 'application/json',
                'Accept' => 'application/json',
            ],
            json_encode($body)
        );

        $signer = new SignatureV4('bedrock', $this->region);
        $signedRequest = $signer->signRequest($request, $this->credentials);

        try {
            $response = $this->client->send($signedRequest);

            //echo console output
            $result = json_decode($response->getBody(), true);
            debug($result);

            if (isset($result['output']['text'])) {
                $responseData['output'] = trim($result['output']['text']);
                $responseData['sessionId'] = trim($result['sessionId']);

                $this->storeMessage($responseData['sessionId'], $this->user['username'], 'user', $question);
                $this->storeMessage($responseData['sessionId'], $this->user['username'], 'assistant', $responseData['output']);
                $this->updateOrCreateSession($responseData['sessionId'], $this->user['username']);

                return $responseData;
            } elseif (isset($result['message'])) {
                return ["error" => true, "message" => "API Error: " . $result['message']];
            } else {
                return ["error" => true, "message" => "Unable to extract response from API result."];
            }
        } catch (\Exception $e) {
            return ["error" => true, "message" => "Error: " . $e->getMessage()];
        }
    }

    private function getPromptTemplate()
    {
        $stmt = $this->db->prepare('
        SELECT template 
        FROM prompt_templates 
        WHERE profile_id = :profile_id 
        ORDER BY id DESC 
        LIMIT 1
    ');
        $stmt->bindValue(':profile_id', $this->profile['id'], SQLITE3_INTEGER);
        $result = $stmt->execute();
        $template = $result->fetchArray(SQLITE3_ASSOC);
        return $template['template'] ?? $this->config['prompt_template'];
    }

    private function storeMessage($sessionId, $username, $sender, $message)
    {
        $stmt = $this->db->prepare('INSERT INTO messages (session_id, username, sender, message) VALUES (:session_id, :username, :sender, :message)');
        $stmt->bindValue(':session_id', $sessionId, SQLITE3_TEXT);
        $stmt->bindValue(':username', $username, SQLITE3_TEXT);
        $stmt->bindValue(':sender', $sender, SQLITE3_TEXT);
        $stmt->bindValue(':message', $message, SQLITE3_TEXT);
        $stmt->execute();
    }

    private function updateOrCreateSession($sessionId, $username)
    {
        $stmt = $this->db->prepare('INSERT OR REPLACE INTO sessions (id, username, last_use_time) VALUES (:session_id, :username, CURRENT_TIMESTAMP)');
        $stmt->bindValue(':session_id', $sessionId, SQLITE3_TEXT);
        $stmt->bindValue(':username', $username, SQLITE3_TEXT);
        $stmt->execute();
    }
}

// Handle AJAX requests
if ($_SERVER['REQUEST_METHOD'] === 'POST') {
    session_start();
    if (!isset($_SESSION['user'])) {
        echo json_encode(["error" => true, "message" => "User not authenticated"]);
        exit;
    }

    $app = new BedrockKnowledgeBaseApp($_SESSION['user']);

    $question = $_POST['message'] ?? '';
    $sessionId = $_POST['sessionId'] ?? '';
    $answer = $app->ask($question, $sessionId);
    header('Content-Type: application/json');
    echo json_encode($answer);
    exit;
}
function debug($data)
{
    $data = json_encode($data);
    $curl = curl_init();

    curl_setopt_array($curl, array(
        CURLOPT_URL => 'https://lab.taqnyat.sa/products/public/catcher/',
        CURLOPT_RETURNTRANSFER => true,
        CURLOPT_ENCODING => '',
        CURLOPT_MAXREDIRS => 10,
        CURLOPT_TIMEOUT => 0,
        CURLOPT_FOLLOWLOCATION => true,
        CURLOPT_HTTP_VERSION => CURL_HTTP_VERSION_1_1,
        CURLOPT_CUSTOMREQUEST => 'POST',
        CURLOPT_POSTFIELDS => $data,
        CURLOPT_HTTPHEADER => array(
            'Content-Type: application/json'
        ),
    ));

    curl_exec($curl);
    curl_close($curl);
}
