// OpsxDB MongoDB Agent —— 与 OpsxJdbcAgent 同协议的 Mongo 专属桥（默认包，单文件）。
//
// 启动：java -cp opsx-agent-jdbc.jar:gson-2.12.1.jar:mongo-driver-*.jar OpsxMongoAgent --jars <目录>
//
// SQL 框迷你 DSL（文档见 agents/README.md）：
//   db.<coll>.find([filter])        查询（filter 为 JSON，可省略）
//   db.<coll>.count([filter])       计数
//   db.<coll>.insertOne({doc})      插入
//   db.<coll>.updateOne({f},{u})    更新（u 用 $set 等操作符）
//   db.<coll>.deleteOne({f})        删除单条
//   db.<coll>.deleteMany({f})       批量删除（Go 侧强制危险确认）
// metaDatabases=库列表；metaTables(schema=库)=集合；metaColumns=抽样文档字段并集。

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.Updates;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;

import org.bson.BsonDocument;
import org.bson.Document;
import org.bson.conversions.Bson;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public final class OpsxMongoAgent {
    static final String AGENT = "opsx-agent-jdbc";
    static final String VERSION = "0.1.0";
    static final int PROTOCOL = 1;
    private static final Pattern DSL = Pattern.compile(
        "^db\\.(\\w+)\\.(find|count|insertOne|updateOne|deleteOne|deleteMany)\\s*(\\{.*\\})?\\s*\\)?\\s*$",
        Pattern.DOTALL);

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

    public static void main(String[] args) throws Exception {
        new OpsxMongoAgent().loop();
    }

    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 {
                out.println(gson.toJson(resp(id, true, dispatch(op, args), 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": {
                MongoClient cli = MongoClients.create(str(a, "url"));
                String handle = "m" + seq.incrementAndGet();
                handles.put(handle, cli);
                JsonObject o = new JsonObject();
                o.addProperty("handle", handle);
                return o;
            }
            case "ping": {
                Document r = zk(a).getDatabase("admin").runCommand(new Document("ping", 1));
                if (r.getDouble("ok") != 1.0) throw new IllegalStateException("ping 失败: " + r.toJson());
                return ok();
            }
            case "query": return query(a);
            case "exec": return exec(a);
            case "metaDatabases": {
                List<String> names = new ArrayList<>();
                zk(a).listDatabaseNames().into(names);
                JsonObject o = new JsonObject();
                o.add("names", gson.toJsonTree(names));
                return o;
            }
            case "metaTables": {
                List<String> names = new ArrayList<>();
                zk(a).getDatabase(str(a, "schema")).listCollectionNames().into(names);
                JsonObject o = new JsonObject();
                o.add("names", gson.toJsonTree(names));
                return o;
            }
            case "metaColumns": {
                // 抽样前 5 个文档的字段并集（v1 结构探测，非权威 schema）
                List<JsonObject> cols = new ArrayList<>();
                Set<String> seen = new LinkedHashSet<>();
                try (var it = coll(a).find().limit(5).iterator()) {
                    while (it.hasNext()) {
                        for (Map.Entry<String, Object> e : it.next().entrySet()) {
                            if (seen.add(e.getKey())) {
                                JsonObject c = new JsonObject();
                                c.addProperty("name", e.getKey());
                                c.addProperty("type", e.getValue() == null ? "unknown" : e.getValue().getClass().getSimpleName());
                                c.addProperty("nullable", true);
                                cols.add(c);
                            }
                        }
                    }
                }
                JsonObject o = new JsonObject();
                o.add("columns", gson.toJsonTree(cols));
                return o;
            }
            case "close": {
                MongoClient c = handles.remove(str(a, "handle"));
                if (c != null) c.close();
                return ok();
            }
            default: throw new IllegalArgumentException("未知 op: " + op);
        }
    }

    /** find/count 的表格式结果 */
    private JsonObject query(JsonObject a) throws Exception {
        Dsl d = parse(str(a, "sql"));
        int maxRows = a.has("maxRows") ? a.get("maxRows").getAsInt() : 1000;
        if (d.verb.equals("count")) {
            long n = coll(a, d).countDocuments(d.filter == null ? new Document() : d.filter);
            JsonObject o = new JsonObject();
            o.add("columns", gson.toJsonTree(new String[]{"n"}));
            JsonObject row = new JsonObject();
            row.addProperty("n", n);
            o.add("rows", gson.toJsonTree(List.of(row)));
            o.addProperty("truncated", false);
            return o;
        }
        if (!d.verb.equals("find")) {
            throw new IllegalArgumentException("query 仅支持 find/count，收到: " + d.verb);
        }
        List<String> cols = new ArrayList<>();
        cols.add("_id");
        Set<String> seen = new LinkedHashSet<>(cols);
        List<JsonObject> rows = new ArrayList<>();
        boolean truncated = false;
        try (var it = coll(a, d).find(d.filter == null ? new Document() : d.filter).limit(maxRows + 1).iterator()) {
            while (it.hasNext()) {
                Document doc = it.next();
                if (rows.size() >= maxRows) { truncated = true; break; }
                for (String k : doc.keySet()) {
                    if (seen.add(k)) cols.add(k);
                }
                JsonObject row = new JsonObject();
                for (String k : cols) {
                    Object v = doc.get(k);
                    row.addProperty(k, v == null ? null : String.valueOf(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;
    }

    /** insertOne/updateOne/deleteOne/deleteMany */
    private JsonObject exec(JsonObject a) throws Exception {
        Dsl d = parse(str(a, "sql"));
        MongoCollection<Document> c = coll(a, d);
        JsonObject o = new JsonObject();
        switch (d.verb) {
            case "insertOne": {
                Document doc = Document.parse(d.rawJson);
                c.insertOne(doc);
                o.addProperty("affectedRows", 1);
                return o;
            }
            case "updateOne": {
                List<Bson> parts = splitTwo(d.rawJson);
                UpdateResult r = c.updateOne(parts.get(0), parts.get(1));
                o.addProperty("affectedRows", r.getModifiedCount());
                return o;
            }
            case "deleteOne": {
                DeleteResult r = c.deleteOne(d.filter == null ? new Document() : d.filter);
                o.addProperty("affectedRows", r.getDeletedCount());
                return o;
            }
            case "deleteMany": {
                DeleteResult r = c.deleteMany(d.filter == null ? new Document() : d.filter);
                o.addProperty("affectedRows", r.getDeletedCount());
                return o;
            }
            default:
                throw new IllegalArgumentException("exec 仅支持 insertOne/updateOne/deleteOne/deleteMany，收到: " + d.verb);
        }
    }

    /** 把 updateOne 的 {filter},{update} 拆成两个 Bson */
    private List<Bson> splitTwo(String raw) {
        // 顶层逗号切分（括号深度 0 处）
        int depth = 0;
        for (int i = 0; i < raw.length(); i++) {
            char ch = raw.charAt(i);
            if (ch == '{') depth++;
            else if (ch == '}') depth--;
            else if (ch == ',' && depth == 0) {
                return List.of(BsonDocument.parse(raw.substring(0, i).trim()),
                               BsonDocument.parse(raw.substring(i + 1).trim()));
            }
        }
        throw new IllegalArgumentException("updateOne 需要 {filter},{update} 两个 JSON");
    }

    private Dsl parse(String sql) {
        String s = sql == null ? "" : sql.trim();
        if (!s.startsWith("db.")) {
            throw new IllegalArgumentException("无法解析的 Mongo DSL（支持 db.<coll>.find/count/insertOne/updateOne/deleteOne/deleteMany）: " + sql);
        }
        int dot = s.indexOf('.', 3);
        if (dot < 0) {
            throw new IllegalArgumentException("缺少集合名: " + sql);
        }
        Dsl d = new Dsl();
        d.coll = s.substring(3, dot);
        String rest = s.substring(dot + 1).trim();
        int i = 0;
        while (i < rest.length() && Character.isLetter(rest.charAt(i))) i++;
        d.verb = rest.substring(0, i);
        String raw = rest.substring(i).trim();
        // 抽出 JSON 体：第一个 '{' 到最后一个 '}'（容忍 insertOne({...}) 的外层括号与空白）
        int b = raw.indexOf('{');
        if (b >= 0) {
            int e = raw.lastIndexOf('}');
            raw = e > b ? raw.substring(b, e + 1) : raw.substring(b);
        }
        d.rawJson = raw;
        Set<String> verbs = Set.of("find", "count", "insertOne", "updateOne", "deleteOne", "deleteMany");
        if (!verbs.contains(d.verb)) {
            throw new IllegalArgumentException("不支持的操作: " + d.verb);
        }
        if (!d.rawJson.isEmpty() && !d.verb.equals("updateOne")) {
            d.filter = Document.parse(d.rawJson);
        }
        return d;
    }

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

    private MongoCollection<Document> coll(JsonObject a) {
        // metaColumns 等入口：schema=库，table=集合
        return zk(a).getDatabase(str(a, "schema")).getCollection(str(a, "table"));
    }

    private MongoCollection<Document> coll(JsonObject a, Dsl d) {
        String db = a.has("schema") && !str(a, "schema").isEmpty() ? str(a, "schema")
                : ((JsonObject) a).has("database") ? str(a, "database") : "test";
        // ExecSQL 路径未带 schema 时，从会话 Conn 透出：args 里由 Go 侧补 database
        if (a.has("db") && !str(a, "db").isEmpty()) db = str(a, "db");
        return zk(a).getDatabase(db).getCollection(d.coll);
    }

    private static final class Dsl {
        String coll;
        String verb;
        String rawJson;
        Document filter;
    }

    private static String str(JsonObject a, String k) {
        return a.has(k) && !a.get(k).isJsonNull() ? a.get(k).getAsString() : "";
    }

    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;
    }
}