import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.sql.*;

/** Probe admin_ai_model by provider / identifier patterns (Java 8). */
public class ProbeTargetAiModels {
    private static final String JDBC = "jdbc:mysql://cq-cdb-8fjmemkb.sql.tencentcdb.com:27220/ylrz_saas"
            + "?useUnicode=true&characterEncoding=utf8&serverTimezone=GMT%2B8";
    private static final String PROBE_MSG = "\u4f60\u597d\uff0c\u8bf7\u7528\u4e00\u53e5\u8bdd\u56de\u590dOK";

    public static void main(String[] a) throws Exception {
        Class.forName("com.mysql.cj.jdbc.Driver");
        int ok = 0;
        int fail = 0;
        String sql = "SELECT id, model_name, provider_code, model_identifier, api_endpoint, api_key, status "
                + "FROM admin_ai_model WHERE status=1 AND ("
                + "model_name LIKE '%2.0-lite%' OR model_identifier LIKE '%2-0-lite%' OR model_identifier LIKE '%2.0-lite%' "
                + "OR model_name LIKE '%2.1-pro%' OR model_identifier LIKE '%2-1-pro%' OR model_identifier LIKE '%2.1-pro%' "
                + "OR provider_code='deepseek' OR model_name LIKE 'DeepSeek%' OR model_identifier LIKE 'deepseek%'"
                + ") ORDER BY id";
        try (Connection c = DriverManager.getConnection(JDBC, "root", "Ylrz_1q2w3e4r5t6y");
             Statement st = c.createStatement();
             ResultSet rs = st.executeQuery(sql)) {
            while (rs.next()) {
                long id = rs.getLong("id");
                String modelName = rs.getString("model_name");
                String provider = rs.getString("provider_code");
                String modelId = rs.getString("model_identifier");
                String endpoint = rs.getString("api_endpoint");
                String key = rs.getString("api_key");
                System.out.println("\n=== " + modelName + " (id=" + id + ") ===");
                System.out.println("provider=" + provider + " model=" + modelId);
                System.out.println("endpoint=" + endpoint);
                System.out.println("keyLen=" + (key == null ? 0 : key.length())
                        + (key != null && key.contains("*") ? " [MASKED_IN_DB]" : ""));
                if (key == null || key.isEmpty() || key.contains("*")) {
                    System.out.println("[FAIL] api_key empty or masked");
                    fail++;
                    continue;
                }
                String ep = normalizeChatEndpoint(endpoint, provider);
                if (probeChat(ep, key, modelId)) {
                    ok++;
                    System.out.println("[PASS] HTTP chat probe OK");
                } else {
                    fail++;
                    System.out.println("[FAIL] HTTP chat probe failed");
                }
            }
        }
        System.out.println("\n=== SUMMARY: " + ok + " passed, " + fail + " failed ===");
        if (ok == 0) {
            System.out.println("No matching enabled models found in DB query");
            System.exit(2);
        }
        if (fail > 0) System.exit(1);
    }

    private static String normalizeChatEndpoint(String endpoint, String provider) {
        if (endpoint == null || endpoint.isEmpty()) {
            if ("doubao".equalsIgnoreCase(provider)) {
                return "https://ark.cn-beijing.volces.com/api/v3/chat/completions";
            }
            return endpoint;
        }
        if (endpoint.contains("/responses")) {
            return endpoint.replace("/responses", "/chat/completions");
        }
        return endpoint;
    }

    private static boolean probeChat(String endpoint, String key, String modelId) throws Exception {
        if (endpoint == null || endpoint.isEmpty()) {
            System.out.println("no endpoint");
            return false;
        }
        String body = "{\"model\":\"" + modelId + "\",\"max_tokens\":64,\"temperature\":0.3,"
                + "\"messages\":[{\"role\":\"user\",\"content\":\"" + PROBE_MSG + "\"}]}";
        HttpURLConnection conn = (HttpURLConnection) new URL(endpoint).openConnection();
        conn.setConnectTimeout(20000);
        conn.setReadTimeout(90000);
        conn.setRequestMethod("POST");
        conn.setDoOutput(true);
        conn.setRequestProperty("Authorization", "Bearer " + key);
        conn.setRequestProperty("Content-Type", "application/json; charset=utf-8");
        byte[] bytes = body.getBytes(StandardCharsets.UTF_8);
        conn.setRequestProperty("Content-Length", String.valueOf(bytes.length));
        try (OutputStream os = conn.getOutputStream()) {
            os.write(bytes);
        }
        int code = conn.getResponseCode();
        InputStream is = code >= 400 ? conn.getErrorStream() : conn.getInputStream();
        String resp = readStream(is);
        System.out.println("HTTP " + code);
        if (resp.length() > 400) resp = resp.substring(0, 400) + "...";
        System.out.println(resp);
        return code >= 200 && code < 300 && resp != null && !resp.isEmpty();
    }

    private static String readStream(InputStream is) throws IOException {
        if (is == null) return "";
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        byte[] buf = new byte[4096];
        int n;
        while ((n = is.read(buf)) != -1) bos.write(buf, 0, n);
        return bos.toString("UTF-8");
    }
}