// OpsxDB JDBC Agent —— 纯 JDK + gson 的通用 JDBC 桥（默认包，支持单文件源启动）。
//
// 启动（由 Go 服务端组装）：
//   java -cp opsx-agent-jdbc.jar:gson-2.12.1.jar OpsxJdbcAgent --jars <驱动jar目录>
//
// 协议：stdin/stdout 换行分隔 JSON（JSON Lines）。
//   请求 : {"id":1,"op":"open","args":{...}}
//  响应  : {"id":1,"ok":true,"result":{...}} 或 {"id":1,"ok":false,"error":"..."}
//
// 驱动 jar 由 --jars 目录按 URLClassLoader 加载（商用驱动如达梦不随产品分发，
// 用户把官方 DmJdbcDriver-*.jar 放进该目录即可）。
//
// 安全设计（与 Go 侧约定一致）：
//   - query 必带 maxRows（JDBC Statement.setMaxRows 硬截断，防 OOM/拖库）
//   - 每条 Statement.setQueryTimeout（秒）兜底，防慢查询挂死业务库
//   - open 支持 readOnly（true 时 Connection.setReadOnly，提示库走只读路由）

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.Driver;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicLong;

public final class OpsxJdbcAgent {
    static final String AGENT = "opsx-agent-jdbc";
    static final String VERSION = "0.1.0";
    static final int PROTOCOL = 1;

    private final Gson gson = new Gson();
    private final Map<String, Connection> handles = new HashMap<>();
    private final AtomicLong seq = new AtomicLong(0);

    public static void main(String[] args) throws Exception {
        String jarsDir = null;
        for (int i = 0; i < args.length - 1; i++) {
            if ("--jars".equals(args[i])) jarsDir = args[i + 1];
        }
        if (jarsDir != null) loadVendorDrivers(jarsDir);
        new OpsxJdbcAgent().loop();
    }

    /** 把 --jars 目录下所有 jar 挂进 URLClassLoader 并逐个注册 Driver */
    private static void loadVendorDrivers(String dir) throws Exception {
        java.io.File d = new java.io.File(dir);
        java.io.File[] jars = d.listFiles((f) -> f.getName().endsWith(".jar"));
        if (jars == null || jars.length == 0) {
            throw new IllegalStateException("驱动目录没有 jar: " + dir);
        }
        URL[] urls = new URL[jars.length];
        for (int i = 0; i < jars.length; i++) urls[i] = jars[i].toURI().toURL();
        URLClassLoader cl = new URLClassLoader(urls, OpsxJdbcAgent.class.getClassLoader());
        for (java.io.File jar : jars) {
            try (URLClassLoader probe = new URLClassLoader(new URL[]{jar.toURI().toURL()}, cl)) {
                java.util.jar.JarFile jf = new java.util.jar.JarFile(jar);
                var entries = jf.entries();
                while (entries.hasMoreElements()) {
                    String name = entries.nextElement().getName();
                    if (name.endsWith("META-INF/services/java.sql.Driver")) {
                        try (var in = cl.getResourceAsStream(name)) {
                            if (in == null) continue;
                            java.io.BufferedReader r = new java.io.BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8));
                            String cls;
                            while ((cls = r.readLine()) != null) {
                                cls = cls.trim();
                                if (cls.isEmpty() || cls.startsWith("#")) continue;
                                Driver drv = (Driver) Class.forName(cls, true, cl).getDeclaredConstructor().newInstance();
                                DriverManager.registerDriver(new ShimDriver(drv));
                            }
                        }
                    }
                }
            }
        }
    }

    /** 让 DriverManager 能用 URLClassLoader 里加载的驱动 */
    static final class ShimDriver implements Driver {
        private final Driver delegate;
        ShimDriver(Driver d) { this.delegate = d; }
        public java.sql.Connection connect(String url, Properties info) throws java.sql.SQLException { return delegate.connect(url, info); }
        public boolean acceptsURL(String url) throws java.sql.SQLException { return delegate.acceptsURL(url); }
        public java.sql.DriverPropertyInfo[] getPropertyInfo(String url, Properties info) throws java.sql.SQLException { return delegate.getPropertyInfo(url, info); }
        public int getMajorVersion() { return delegate.getMajorVersion(); }
        public int getMinorVersion() { return delegate.getMinorVersion(); }
        public boolean jdbcCompliant() { return delegate.jdbcCompliant(); }
        public java.util.logging.Logger getParentLogger() { return java.util.logging.Logger.getGlobal(); }
    }

    private void loop() throws Exception {
        BufferedReader in = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
        PrintWriter out = new PrintWriter(new OutputStreamWriter(System.out, StandardCharsets.UTF_8), true);
        out.println(gson.toJson(resp(0, true, helloResult(), null))); // 首行即握手公告
        String line;
        while ((line = in.readLine()) != null) {
            line = line.trim();
            if (line.isEmpty()) continue;
            JsonObject req;
            try {
                req = JsonParser.parseString(line).getAsJsonObject();
            } catch (Exception e) {
                out.println(gson.toJson(resp(0, false, null, "bad json: " + e.getMessage())));
                continue;
            }
            long id = req.has("id") ? req.get("id").getAsLong() : 0;
            String op = req.has("op") ? req.get("op").getAsString() : "";
            JsonObject args = req.has("args") && req.get("args").isJsonObject() ? req.getAsJsonObject("args") : new JsonObject();
            try {
                JsonObject result = dispatch(op, args);
                out.println(gson.toJson(resp(id, true, result, null)));
            } catch (Exception e) {
                out.println(gson.toJson(resp(id, false, null, e.getMessage() == null ? String.valueOf(e) : e.getMessage())));
            }
            if ("shutdown".equals(op)) {
                out.flush();
                System.exit(0);
            }
        }
        System.exit(0);
    }

    private JsonObject helloResult() {
        JsonObject o = new JsonObject();
        o.addProperty("agent", AGENT);
        o.addProperty("version", VERSION);
        o.addProperty("protocol", PROTOCOL);
        return o;
    }

    private JsonObject dispatch(String op, JsonObject a) throws Exception {
        switch (op) {
            case "hello": return helloResult();
            case "open": return open(a);
            case "ping": { stmt(getHandle(a), "SELECT 1", 5, 1); return ok(); }
            case "query": return query(a);
            case "exec": return exec(a);
            case "metaDatabases": return metaDatabases(getHandle(a));
            case "metaTables": return metaTables(getHandle(a), str(a, "schema"));
            case "metaColumns": return metaColumns(getHandle(a), str(a, "schema"), str(a, "table"));
            case "close": {
                Connection c = handles.remove(getHandle(a));
                if (c != null) c.close();
                return ok();
            }
            default: throw new IllegalArgumentException("未知 op: " + op);
        }
    }

    private JsonObject open(JsonObject a) throws Exception {
        String url = str(a, "url");
        Properties props = new Properties();
        if (a.has("user") && !a.get("user").isJsonNull()) props.setProperty("user", a.get("user").getAsString());
        if (a.has("password") && !a.get("password").isJsonNull()) props.setProperty("password", a.get("password").getAsString());
        // 其余属性原样透传（SSL 等）
        if (a.has("props") && a.get("props").isJsonObject()) {
            for (var e : a.getAsJsonObject("props").entrySet()) {
                if (!props.containsKey(e.getKey())) props.setProperty(e.getKey(), e.getValue().getAsString());
            }
        }
        Connection conn = DriverManager.getConnection(url, props);
        if (a.has("readOnly") && a.get("readOnly").getAsBoolean()) {
            try { conn.setReadOnly(true); } catch (Exception ignored) { /* 部分驱动不支持 */ }
        }
        String handle = "c" + seq.incrementAndGet();
        handles.put(handle, conn);
        JsonObject o = new JsonObject();
        o.addProperty("handle", handle);
        return o;
    }

    private JsonObject query(JsonObject a) throws Exception {
        Connection conn = getHandle(a);
        String sql = str(a, "sql");
        int maxRows = a.has("maxRows") ? a.get("maxRows").getAsInt() : 1000;
        if (maxRows <= 0 || maxRows > 20000) maxRows = 1000;
        int timeoutSec = a.has("timeoutSec") ? a.get("timeoutSec").getAsInt() : 30;
        try (Statement st = conn.createStatement()) {
            st.setMaxRows(maxRows);
            st.setQueryTimeout(timeoutSec);
            try (ResultSet rs = st.executeQuery(sql)) {
                return readRows(rs, maxRows);
            }
        }
    }

    private JsonObject exec(JsonObject a) throws Exception {
        Connection conn = getHandle(a);
        int timeoutSec = a.has("timeoutSec") ? a.get("timeoutSec").getAsInt() : 30;
        try (Statement st = conn.createStatement()) {
            st.setQueryTimeout(timeoutSec);
            boolean hasRs = st.execute(str(a, "sql"));
            JsonObject o = new JsonObject();
            if (hasRs) {
                try (ResultSet rs = st.getResultSet()) {
                    o = readRows(rs, 1000);
                    o.addProperty("rowCount", o.getAsJsonArray("rows").size());
                    return o;
                }
            }
            o.addProperty("affectedRows", st.getUpdateCount());
            return o;
        }
    }

    private JsonObject readRows(ResultSet rs, int maxRows) throws Exception {
        ResultSetMetaData md = rs.getMetaData();
        int n = md.getColumnCount();
        List<String> cols = new ArrayList<>(n);
        for (int i = 1; i <= n; i++) cols.add(md.getColumnLabel(i));
        List<JsonObject> rows = new ArrayList<>();
        int count = 0;
        boolean truncated = false;
        while (rs.next()) {
            if (++count > maxRows) { truncated = true; break; }
            JsonObject row = new JsonObject();
            for (int i = 1; i <= n; i++) {
                String v = rs.getString(i);
                row.addProperty(cols.get(i - 1), v == null ? null : v);
            }
            rows.add(row);
        }
        JsonObject o = new JsonObject();
        o.add("columns", gson.toJsonTree(cols));
        o.add("rows", gson.toJsonTree(rows));
        o.addProperty("truncated", truncated);
        return o;
    }

    private JsonObject metaDatabases(Connection conn) throws Exception {
        List<String> out = new ArrayList<>();
        try (ResultSet rs = conn.getMetaData().getCatalogs()) {
            while (rs.next()) out.add(rs.getString(1));
        }
        if (out.isEmpty()) { // 无 catalog 概念的库（如 Oracle/DM 常用 schema）退化为 schemas
            try (ResultSet rs = conn.getMetaData().getSchemas()) {
                while (rs.next()) out.add(rs.getString(1));
            }
        }
        JsonObject o = new JsonObject();
        o.add("names", gson.toJsonTree(out));
        return o;
    }

    private JsonObject metaTables(Connection conn, String schema) throws Exception {
        List<String> out = new ArrayList<>();
        try (ResultSet rs = conn.getMetaData().getTables(conn.getCatalog(), emptyToNull(schema), "%", new String[]{"TABLE", "VIEW"})) {
            while (rs.next()) out.add(rs.getString("TABLE_NAME"));
        }
        JsonObject o = new JsonObject();
        o.add("names", gson.toJsonTree(out));
        return o;
    }

    private JsonObject metaColumns(Connection conn, String schema, String table) throws Exception {
        List<JsonObject> out = new ArrayList<>();
        try (ResultSet rs = conn.getMetaData().getColumns(conn.getCatalog(), emptyToNull(schema), table, "%")) {
            while (rs.next()) {
                JsonObject col = new JsonObject();
                col.addProperty("name", rs.getString("COLUMN_NAME"));
                col.addProperty("type", rs.getString("TYPE_NAME"));
                col.addProperty("nullable", rs.getInt("NULLABLE") == ResultSetMetaData.columnNullable);
                String key = rs.getString("IS_AUTOINCREMENT");
                col.addProperty("key", "YES".equalsIgnoreCase(key) ? "PRI" : "");
                out.add(col);
            }
        }
        JsonObject o = new JsonObject();
        o.add("columns", gson.toJsonTree(out));
        return o;
    }

    /** 探活：驱动差异兜底（Oracle/DM 老版本无 dual-less SELECT 1 时退 DUAL） */
    private void stmt(Connection conn, String sql, int timeoutSec, int maxRows) throws Exception {
        try (Statement st = conn.createStatement()) {
            st.setQueryTimeout(timeoutSec);
            st.setMaxRows(maxRows);
            st.execute(sql);
        } catch (java.sql.SQLException e) {
            if (sql.equals("SELECT 1")) {
                try (Statement st = conn.createStatement()) {
                    st.setQueryTimeout(timeoutSec);
                    st.setMaxRows(maxRows);
                    st.execute("SELECT 1 FROM DUAL");
                    return;
                }
            }
            throw e;
        }
    }

    private Connection getHandle(JsonObject a) {
        String h = str(a, "handle");
        Connection c = handles.get(h);
        if (c == null) throw new IllegalStateException("句柄不存在或已关闭: " + h);
        return c;
    }

    private static String str(JsonObject a, String k) {
        return a.has(k) && !a.get(k).isJsonNull() ? a.get(k).getAsString() : "";
    }
    private static String emptyToNull(String s) { return s == null || s.isEmpty() ? null : s; }
    private static JsonObject ok() { return new JsonObject(); }
    private static JsonObject resp(long id, boolean ok, JsonObject result, String error) {
        JsonObject o = new JsonObject();
        o.addProperty("id", id);
        o.addProperty("ok", ok);
        if (result != null) o.add("result", result);
        if (error != null) o.addProperty("error", error);
        return o;
    }
}