<?php
require_once(__DIR__ . '/litellm_client.php');
require_once(__DIR__ . '/sitesmith-schema.php');
require_once(__DIR__ . '/sitesmith-prompts.php');
require_once(__DIR__ . '/sitesmith-html-sanitizer.php');

class Sitesmith {
    private const SETTINGS_FILE = '/docker/whp/settings.json';
    private PDO $pdo;
    private array $settings;

    public function __construct(PDO $pdo) {
        $this->pdo = $pdo;
        $this->settings = self::loadSettings();
    }

    public function checkEntitlement(string $username): string {
        $st = $this->pdo->prepare('SELECT sitesmith_enabled, sitesmith_monthly_cap, sitesmith_bonus_credits FROM whp.client_allowances WHERE username = ?');
        $st->execute([$username]);
        $a = $st->fetch(PDO::FETCH_ASSOC) ?: ['sitesmith_enabled'=>0,'sitesmith_monthly_cap'=>null,'sitesmith_bonus_credits'=>0];
        $bonus = (int)$a['sitesmith_bonus_credits'];
        $enabled = (int)$a['sitesmith_enabled'] === 1;
        if ($bonus <= 0 && !$enabled) return 'DISABLED';

        $cap = $a['sitesmith_monthly_cap'] !== null ? (int)$a['sitesmith_monthly_cap'] : (int)$this->settings['default_monthly_cap'];
        $usedMonth = (int)$this->scalar("SELECT COUNT(*) FROM whp.sitesmith_usage WHERE username=? AND consumed_from='monthly' AND created_at >= DATE_FORMAT(NOW(),'%Y-%m-01')", [$username]);
        if ($bonus <= 0 && $usedMonth >= $cap) return 'CAP_REACHED';

        $userSpent   = (float)$this->scalar("SELECT COALESCE(SUM(cost_usd),0) FROM whp.sitesmith_usage WHERE username=? AND created_at >= CURDATE()", [$username]);
        $serverSpent = (float)$this->scalar("SELECT COALESCE(SUM(cost_usd),0) FROM whp.sitesmith_usage WHERE created_at >= CURDATE()");
        if ($userSpent   >= (float)$this->settings['max_cost_per_user_per_day_usd'])   return 'USER_KILLSWITCH';
        if ($serverSpent >= (float)$this->settings['max_cost_per_server_per_day_usd']) return 'SERVER_KILLSWITCH';

        $recent = (int)$this->scalar("SELECT COUNT(*) FROM whp.sitesmith_usage WHERE username=? AND created_at > (NOW() - INTERVAL 1 MINUTE)", [$username]);
        if ($recent >= (int)$this->settings['per_user_rate_limit_per_minute']) return 'RATE_LIMITED';

        return $bonus > 0 ? 'OK_BONUS' : 'OK_MONTHLY';
    }

    public function entitlementSummary(string $username): array {
        $st = $this->pdo->prepare('SELECT sitesmith_enabled, sitesmith_monthly_cap, sitesmith_bonus_credits FROM whp.client_allowances WHERE username = ?');
        $st->execute([$username]);
        $a = $st->fetch(PDO::FETCH_ASSOC) ?: ['sitesmith_enabled'=>0,'sitesmith_monthly_cap'=>null,'sitesmith_bonus_credits'=>0];
        $cap = $a['sitesmith_monthly_cap'] !== null ? (int)$a['sitesmith_monthly_cap'] : (int)$this->settings['default_monthly_cap'];
        $usedMonth = (int)$this->scalar("SELECT COUNT(*) FROM whp.sitesmith_usage WHERE username=? AND consumed_from='monthly' AND created_at >= DATE_FORMAT(NOW(),'%Y-%m-01')", [$username]);
        return [
            'enabled'       => (int)$a['sitesmith_enabled'] === 1,
            'monthly_cap'   => $cap,
            'monthly_used'  => $usedMonth,
            'bonus_credits' => (int)$a['sitesmith_bonus_credits'],
            'resets_on'     => date('Y-m-d', strtotime('first day of next month')),
            'status'        => $this->checkEntitlement($username),
        ];
    }

    public function send(string $username, int $siteId, string $userMessage, string $canvasSummary): array {
        $status = $this->checkEntitlement($username);
        if (!in_array($status, ['OK_BONUS','OK_MONTHLY'], true)) {
            return ['ok'=>false,'status'=>$status,'message'=>$this->statusMessage($status)];
        }
        if (!empty($this->settings['prompt_injection_guard'])) {
            require_once(__DIR__ . '/coraza-ai.php');
            if (function_exists('coraza_ai_classify_prompt')) {
                $verdict = coraza_ai_classify_prompt($userMessage);
                if (($verdict['risk'] ?? '') === 'high') {
                    error_log("sitesmith: blocked prompt from {$username}");
                    return ['ok'=>false,'status'=>'BLOCKED','message'=>'Your message was blocked by our content guard. Please rephrase.'];
                }
            }
        }

        $thread = $this->getOrCreateThread($username, $siteId);
        $history = $this->loadHistory($thread['id'], (int)$this->settings['max_history_turns']);
        $this->saveMessage($thread['id'], 'user', $userMessage, null);

        $prefix = sitesmith_system_prefix();
        $assets = $this->getSiteAssets($siteId, $username);
        $suffix = sitesmith_dynamic_suffix($canvasSummary, $history, $assets);
        $messages = [
            ['role'=>'system','content'=>[['type'=>'text','text'=>$prefix,'cache_control'=>['type'=>'ephemeral']]]],
            ['role'=>'system','content'=>$suffix],
            ['role'=>'user','content'=>$userMessage],
        ];
        $opts = [
            'response_format' => ['type'=>'json_object'],
            'max_tokens'      => (int)$this->settings['max_output_tokens'],
            'temperature'     => 0.2,
            'timeout'         => 60,
        ];
        $tier = $this->settings['litellm_tier'] ?? 'advanced';

        $resp = litellm_chat($tier, $messages, $opts);
        if (!$resp['success']) {
            return ['ok'=>false,'status'=>'AI_ERROR','message'=>'AI service error: '.($resp['error']??'unknown')];
        }

        $parsed = $this->parseAndValidate($resp['content']);
        if (!$parsed['ok']) {
            $messages[] = ['role'=>'assistant','content'=>$resp['content']];
            $messages[] = ['role'=>'user','content'=>"Your previous reply was not valid: ".$parsed['error'].". Reply again with a valid JSON object."];
            $resp2 = litellm_chat($tier, $messages, $opts);
            if (!$resp2['success']) {
                $this->saveMessage($thread['id'], 'assistant', 'AI service error', 'error');
                return ['ok'=>false,'status'=>'AI_ERROR','message'=>'AI service error on retry'];
            }
            $parsed = $this->parseAndValidate($resp2['content']);
            if (!$parsed['ok']) {
                $this->saveMessage($thread['id'], 'assistant', 'AI returned invalid output', 'error');
                return ['ok'=>false,'status'=>'AI_INVALID','message'=>'The AI returned a response we could not understand. Please try again.'];
            }
            $usage = $resp2['usage'] ?? ['prompt_tokens'=>0,'completion_tokens'=>0,'cached_tokens'=>0];
            $modelUsed = $resp2['model'] ?? '';
        } else {
            $usage = $resp['usage'] ?? ['prompt_tokens'=>0,'completion_tokens'=>0,'cached_tokens'=>0];
            $modelUsed = $resp['model'] ?? '';
        }

        $data = $parsed['data'];
        $this->sanitizeTreeInPlace($data);

        $assistantMessageId = $this->saveMessage($thread['id'], 'assistant', json_encode($data, JSON_UNESCAPED_SLASHES), $data['type']);
        $this->recordUsage([
            'username'=>$username, 'site_id'=>$siteId, 'thread_id'=>$thread['id'],
            'message_id'=>$assistantMessageId, 'model'=>$modelUsed,
            'prompt_tokens'=>(int)$usage['prompt_tokens'],
            'completion_tokens'=>(int)$usage['completion_tokens'],
            'cached_tokens'=>(int)$usage['cached_tokens'],
            'cost_usd'=>$this->computeCost($modelUsed, $usage),
            'consumed_from'=>$status === 'OK_BONUS' ? 'bonus' : 'monthly',
        ]);
        if ($status === 'OK_BONUS') {
            // Note: small TOCTOU race here — two concurrent requests can both pass checkEntitlement
            // with bonus=1, both decrement, both succeed. GREATEST(..., 0) floor prevents negative.
            // Acceptable for v1; consider SELECT FOR UPDATE or atomic CAS if it becomes a problem.
            $this->pdo->prepare('UPDATE whp.client_allowances SET sitesmith_bonus_credits = GREATEST(sitesmith_bonus_credits - 1, 0) WHERE username = ?')->execute([$username]);
        }
        return ['ok'=>true,'response'=>$data];
    }

    /**
     * Read the site's uploaded image assets so the AI can reference existing
     * media instead of always emitting placeholder URLs. Returns at most $max
     * entries, sorted newest-first. Each entry includes a filename hint
     * (stripped of the upload prefix), the proxy URL, mime type, and — for
     * JPEGs — any EXIF ImageDescription / UserComment as additional context.
     */
    public function getSiteAssets(int $siteId, string $username, int $max = 30): array {
        $st = $this->pdo->prepare('
            SELECT sd.fqdn, d.domain_name
            FROM whp.sites s
            JOIN whp.site_domains sd ON s.id = sd.site_id AND sd.is_primary = 1
            JOIN whp.domains d ON sd.domain_id = d.id
            WHERE s.id = ? AND (s.username = ? OR ? = ?)
        ');
        // The "username = root literal" check lets root see all assets.
        $st->execute([$siteId, $username, $username, 'root']);
        $row = $st->fetch(PDO::FETCH_ASSOC);
        if (!$row) return [];

        $fqdn = $row['fqdn'] ?: $row['domain_name'];
        // Wildcard domains (*.example.com) live under wildcard.example.com on disk.
        $dirName = (strpos($fqdn, '*.') === 0) ? 'wildcard.' . substr($fqdn, 2) : $fqdn;
        $assetsDir = "/docker/users/{$username}/{$dirName}/.site-builder/assets/";
        if (!is_dir($assetsDir)) return [];

        $rows = [];
        foreach (scandir($assetsDir) ?: [] as $entry) {
            if ($entry === '.' || $entry === '..') continue;
            $path = $assetsDir . $entry;
            if (!is_file($path)) continue;
            $mime = function_exists('mime_content_type') ? (mime_content_type($path) ?: '') : '';
            if (strpos($mime, 'image/') !== 0) continue;

            // Strip the upload-prefix (epoch_hex_originalname) for a readable hint.
            $original = preg_replace('/^\d+_[a-f0-9]+_/', '', $entry);

            $asset = [
                'name'     => $entry,
                'original' => $original,
                'url'      => '/api/assets-proxy/?site_id=' . $siteId . '&filename=' . urlencode($entry),
                'mime'     => $mime,
                'mtime'    => filemtime($path),
            ];

            if ($mime === 'image/jpeg' && function_exists('exif_read_data')) {
                $exif = @exif_read_data($path, 'IFD0,COMMENT', false, false);
                if (is_array($exif)) {
                    if (!empty($exif['ImageDescription'])) {
                        $asset['description'] = substr((string)$exif['ImageDescription'], 0, 200);
                    }
                    if (!empty($exif['UserComment'])) {
                        $asset['comment'] = substr((string)$exif['UserComment'], 0, 200);
                    }
                }
            }
            $rows[] = $asset;
        }
        usort($rows, fn($a, $b) => $b['mtime'] - $a['mtime']);
        return array_slice($rows, 0, $max);
    }

    public function getHistory(string $username, int $siteId): array {
        $thread = $this->getOrCreateThread($username, $siteId);
        $st = $this->pdo->prepare('SELECT role, content, response_type, created_at FROM whp.sitesmith_messages WHERE thread_id=? ORDER BY created_at ASC, id ASC');
        $st->execute([$thread['id']]);
        return $st->fetchAll(PDO::FETCH_ASSOC);
    }

    /**
     * Wipe the chat history for one site. Deletes all messages on the site's
     * thread plus the thread row itself; usage rows (cost telemetry) are kept
     * for billing/reporting integrity.
     */
    public function clearHistory(string $username, int $siteId): array {
        $st = $this->pdo->prepare('SELECT id FROM whp.sitesmith_threads WHERE site_id=?');
        $st->execute([$siteId]);
        $threadId = $st->fetchColumn();
        if (!$threadId) return ['ok' => true, 'cleared' => 0];

        $st = $this->pdo->prepare('DELETE FROM whp.sitesmith_messages WHERE thread_id=?');
        $st->execute([$threadId]);
        $cleared = $st->rowCount();

        $this->pdo->prepare('DELETE FROM whp.sitesmith_threads WHERE id=?')->execute([$threadId]);
        return ['ok' => true, 'cleared' => (int)$cleared];
    }

    private function parseAndValidate(string $jsonText): array {
        $clean = preg_replace('/^```(?:json)?\s*|\s*```$/m', '', trim($jsonText));
        $obj = json_decode($clean, true);
        if ($obj === null) return ['ok'=>false,'error'=>'not valid JSON: '.json_last_error_msg()];
        return sitesmith_validate_response($obj);
    }

    private function sanitizeTreeInPlace(array &$resp): void {
        $walk = function(&$node) use (&$walk) {
            if (!is_array($node)) return;
            $type = $node['type']['resolvedName'] ?? null;
            if ($type === 'HtmlBlock' && isset($node['props']['code']) && is_string($node['props']['code'])) {
                $node['props']['code'] = sitesmith_sanitize_html($node['props']['code']);
            }
            if (isset($node['nodes']) && is_array($node['nodes'])) foreach ($node['nodes'] as &$c) $walk($c);
        };
        if ($resp['type'] === 'replace') {
            foreach ($resp['pages'] as &$p) $walk($p['tree']);
            foreach (['header','footer'] as $k) if (isset($resp[$k]['tree'])) $walk($resp[$k]['tree']);
        } elseif ($resp['type'] === 'patch') {
            foreach ($resp['ops'] as &$op) {
                if (isset($op['tree'])) $walk($op['tree']);
                if (($op['op'] ?? '') === 'update_props' && isset($op['props']['code']) && is_string($op['props']['code'])) {
                    $op['props']['code'] = sitesmith_sanitize_html($op['props']['code']);
                }
            }
        }
    }

    private function getOrCreateThread(string $username, int $siteId): array {
        $st = $this->pdo->prepare('SELECT id, created_at FROM whp.sitesmith_threads WHERE site_id = ?');
        $st->execute([$siteId]);
        $t = $st->fetch(PDO::FETCH_ASSOC);
        if ($t) return $t;
        $this->pdo->prepare('INSERT INTO whp.sitesmith_threads (site_id, username) VALUES (?, ?)')->execute([$siteId, $username]);
        return ['id'=>(int)$this->pdo->lastInsertId(),'created_at'=>date('c')];
    }

    private function loadHistory(int $threadId, int $maxTurns): array {
        $st = $this->pdo->prepare('SELECT role, content FROM whp.sitesmith_messages WHERE thread_id=? ORDER BY created_at DESC, id DESC LIMIT ?');
        $st->bindValue(1, $threadId, PDO::PARAM_INT);
        $st->bindValue(2, $maxTurns * 2, PDO::PARAM_INT);
        $st->execute();
        return array_reverse($st->fetchAll(PDO::FETCH_ASSOC));
    }

    private function saveMessage(int $threadId, string $role, string $content, ?string $type): int {
        $st = $this->pdo->prepare('INSERT INTO whp.sitesmith_messages (thread_id, role, content, response_type) VALUES (?, ?, ?, ?)');
        $st->execute([$threadId, $role, $content, $type]);
        return (int)$this->pdo->lastInsertId();
    }

    private function recordUsage(array $row): void {
        $cols = array_keys($row);
        $st = $this->pdo->prepare('INSERT INTO whp.sitesmith_usage (' . implode(',', $cols) . ') VALUES (' . implode(',', array_fill(0, count($cols), '?')) . ')');
        $st->execute(array_values($row));
    }

    private function computeCost(string $model, array $usage): float {
        $pricing = $this->lookupPricing($model);
        if (!$pricing) { error_log("sitesmith: no pricing for model {$model}"); return 0.0; }
        $fresh = max(0, (int)$usage['prompt_tokens'] - (int)$usage['cached_tokens']);
        $cached = (int)$usage['cached_tokens'];
        $out = (int)$usage['completion_tokens'];
        return round(
            ($fresh  / 1_000_000.0) * (float)$pricing['input_per_mtok'] +
            ($cached / 1_000_000.0) * (float)$pricing['cached_per_mtok'] +
            ($out    / 1_000_000.0) * (float)$pricing['output_per_mtok'],
            6
        );
    }

    /**
     * Look up pricing for a model name, tolerating LiteLLM provider prefixes.
     * Order: exact match -> prefix-stripped -> contains-match.
     */
    private function lookupPricing(string $model): ?array {
        $table = $this->settings['model_pricing'] ?? [];
        if (isset($table[$model])) return $table[$model];
        $stripped = preg_replace('#^(global\.|bedrock/|us\.|eu\.)?(anthropic\.)?#', '', $model) ?? $model;
        if ($stripped !== $model && isset($table[$stripped])) return $table[$stripped];
        foreach ($table as $key => $row) {
            if ($key !== '' && stripos($model, $key) !== false) return $row;
        }
        return null;
    }

    private function statusMessage(string $status): string {
        switch ($status) {
            case 'DISABLED':          return 'Sitesmith is a paid addon. Upgrade your hosting plan to enable it.';
            case 'CAP_REACHED':       return "You've used your monthly Sitesmith builds. Resets " . date('F j', strtotime('first day of next month')) . '.';
            case 'USER_KILLSWITCH':   return 'Daily limit reached. Please try again tomorrow.';
            case 'SERVER_KILLSWITCH': return 'Sitesmith is temporarily unavailable. Please try again later.';
            case 'RATE_LIMITED':      return 'Slow down — try again in a moment.';
            default: return 'Sitesmith is unavailable.';
        }
    }

    private function scalar(string $sql, array $params = []): string {
        $st = $this->pdo->prepare($sql);
        $st->execute($params);
        return (string)$st->fetchColumn();
    }

    private static function loadSettings(): array {
        $defaults = [
            'default_monthly_cap'             => 50,
            'max_output_tokens'               => 8000,
            'max_input_tokens'                => 12000,
            'max_history_turns'               => 6,
            'max_cost_per_user_per_day_usd'   => 2.00,
            'max_cost_per_server_per_day_usd' => 50.00,
            'per_user_rate_limit_per_minute'  => 6,
            'litellm_tier'                    => 'advanced',
            'alert_at_pct_of_server_cap'      => 80,
            'prompt_injection_guard'          => false,
            'model_pricing'                   => [],
        ];
        if (!file_exists(self::SETTINGS_FILE)) return $defaults;
        $j = json_decode((string)file_get_contents(self::SETTINGS_FILE), true);
        $s = is_array($j['sitesmith'] ?? null) ? $j['sitesmith'] : [];
        return array_merge($defaults, $s);
    }
}
