package gudusoft.gsqlparser.dlineage.comment;

import gudusoft.gsqlparser.*;
import gudusoft.gsqlparser.nodes.*;
import gudusoft.gsqlparser.dlineage.dataflow.model.*;
import gudusoft.gsqlparser.dlineage.util.Pair;
import gudusoft.gsqlparser.dlineage.comment.CommentBindingBundle.*;
import gudusoft.gsqlparser.dlineage.comment.CommentHostCollector.Draft;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.*;

/** Per-analysis, opt-in draft store. build() releases all AST and token references. */
public final class CommentBindingCollector {
    private static final java.util.regex.Pattern ENGINE_DIRECTIVE = java.util.regex.Pattern.compile("(?m)(?:^|,)\\s*(?:--\\s*)?(?:db-instance|db|schema|process_label|process)\\s*[:=]");
    private static final List<String> KINDS = Arrays.asList("container", "routine", "call", "statement", "column", "table_ref", "object_declaration");
    private static final class Unit {
        SourceUnit source; String sql; int base; CommentTokenIndex index; boolean derived;
        final List<Draft> hosts = new ArrayList<Draft>();
        final List<TSourceToken> comments = new ArrayList<TSourceToken>();
    }
    private static final class Parsed {
        CommentTokenIndex index; CommentHostCollector hosts; List<Unit> units; boolean partial;
        TSourceTokenList tokens;
    }
    private final Option option;
    private final Map<String,List<Unit>> inputs = new LinkedHashMap<String,List<Unit>>();
    private final List<Parsed> parsed = new ArrayList<Parsed>();
    private final List<Unit> units = new ArrayList<Unit>();
    private boolean failed;
    private boolean missingSource;
    private int derivedOrdinal;
    private final Set<TSourceTokenList> captured = Collections.newSetFromMap(new IdentityHashMap<TSourceTokenList, Boolean>());
    public CommentBindingCollector(Option option) { this.option = option; }

    public void register(String key, String sql, int base, int inputOrdinal, String fileName, boolean derived) {
        try { registerUnit(key, sql, base, inputOrdinal, fileName, derived); }
        catch (Exception e) { failed = true; }
    }
    private Unit registerUnit(String key, String sql, int base, int ordinal, String fileName, boolean derived) throws Exception {
        Unit u = new Unit(); u.sql = sql; u.base = base; u.derived = derived;
        String id = "s" + (units.size() + 1);
        byte[] hash = MessageDigest.getInstance("SHA-256").digest(sql.getBytes(StandardCharsets.UTF_8));
        StringBuilder digest = new StringBuilder(); for (byte b : hash) digest.append(String.format(Locale.ROOT, "%02x", b & 255));
        if (fileName != null) fileName = fileName.replace('\\', '/');
        if (fileName != null) fileName = fileName.substring(fileName.lastIndexOf('/') + 1);
        u.source = new SourceUnit(id, ordinal, option.getVendor().toString(), digest.toString(), derived ? "derived" : "original", fileName);
        u.index = new CommentTokenIndex(sql, new TSourceTokenList());
        units.add(u);
        inputs.computeIfAbsent(key, k -> new ArrayList<Unit>()).add(u);
        return u;
    }
    /** Imported lineage without a bundle does not provide original SQL comment facts. */
    public void sourceUnavailable() { missingSource = true; }

    public void captureDerived(TGSqlParser parser, boolean parsePartial) {
        capture("derived:" + (++derivedOrdinal), parser, parsePartial);
    }
    public void capture(String key, TGSqlParser parser, boolean parsePartial) {
        try {
            Parsed primary = captureOne(key, parser, parsePartial);
            Deque<TGSqlParser> foreign = new ArrayDeque<TGSqlParser>(primary.hosts.foreignParsers);
            while (!foreign.isEmpty()) {
                TGSqlParser nested = foreign.removeFirst();
                if (captured.contains(nested.sourcetokenlist)) continue;
                // Routine-body parsers already ran. Capture iteratively without parsing again.
                Parsed child = captureOne("derived:" + (++derivedOrdinal), nested, nested.getErrorCount() != 0);
                foreign.addAll(child.hosts.foreignParsers);
            }
        } catch (Exception e) { failed = true; }
    }
    private Parsed captureOne(String key, TGSqlParser parser, boolean parsePartial) throws Exception {
        captured.add(parser.sourcetokenlist);
        Parsed p = new Parsed();
        p.units = inputs.get(key);
        if (p.units == null) {
            Unit u = registerUnit(key, parser.sqltext, 0, units.size(), null, true);
            p.units = Collections.singletonList(u);
        }
        p.tokens = parser.sourcetokenlist;
        p.index = new CommentTokenIndex(parser.sqltext, parser.sourcetokenlist);
        p.hosts = new CommentHostCollector(p.index, parser);
        p.partial = parsePartial;
        parsed.add(p);
        return p;
    }
    public void clear() { inputs.clear(); parsed.clear(); units.clear(); captured.clear(); }

    public CommentBindingBundle build(ModelBindingManager manager) {
        CommentBindingBundle out = new CommentBindingBundle();
        out.supportedHostKinds.addAll(KINDS);
        out.limits = new Limits(option.getCommentBindingMaxRawLength(), option.getCommentBindingMaxCommentsPerSourceUnit(), option.getCommentBindingMaxTotalRawLength());
        try {
            final IdentityHashMap<Object,List<Target>> models = new IdentityHashMap<Object,List<Target>>();
            manager.visitExactBindings((key, model) -> observe(models, key, model));
            IdentityHashMap<TSourceToken,Parsed> tokenOwners = new IdentityHashMap<TSourceToken,Parsed>();
            for (Parsed p : parsed) for (TSourceToken token : p.index.code) tokenOwners.put(token,p);
            for (Relationship relation : manager.getRelations()) if (relation instanceof CallRelationship) {
                CallRelationship call = (CallRelationship)relation;
                addTarget(models, call.getCallObject(), new Target("relationship", String.valueOf(call.getId()), null));
                Parsed p = call.getCallObject() == null || call.getCallObject().getStartToken() == null ? null
                        : tokenOwners.get(call.getCallObject().getStartToken());
                if (p != null) {
                    Draft h = p.hosts.add(call.getCallObject(), null, "call", "call_occurrence", null, 0);
                    if (h != null) h.kind = "call";
                }
            }
            int total = 0;
            Set<Unit> processed = Collections.newSetFromMap(new IdentityHashMap<Unit, Boolean>());
            for (Unit u : units) out.sourceUnits.add(u.source);
            for (Parsed p : parsed) {
                p.hosts.hosts.sort(Comparator.comparingInt((Draft h) -> h.start).thenComparingInt(h -> -h.end).thenComparing(h -> h.kind));
                for (int i = 0; i < p.hosts.hosts.size(); i++) p.hosts.hosts.get(i).id = "h" + (out.hosts.size() + i + 1);
                NavigableMap<Integer,Unit> boundaries = new TreeMap<Integer,Unit>();
                for (Unit u : p.units) boundaries.put(u.base,u);
                for (Draft h : p.hosts.hosts) {
                    Map.Entry<Integer,Unit> entry = boundaries.floorEntry(h.start);
                    if (entry != null && h.end <= entry.getValue().base + entry.getValue().sql.length()) entry.getValue().hosts.add(h);
                }
                for (TSourceToken token : p.index.comments) {
                    Map.Entry<Integer,Unit> entry = boundaries.floorEntry(CommentTokenIndex.start(token));
                    if (entry != null && CommentTokenIndex.end(token) <= entry.getValue().base + entry.getValue().sql.length()) entry.getValue().comments.add(token);
                }
                CommentBinder binder = new CommentBinder(p.index, p.hosts.hosts, p.hosts.expressionBoundaries);
                for (Unit u : p.units) {
                    processed.add(u);
                    String sourceId = u.source.id;
                    Draft container = new Draft(); container.id = sourceId;
                    container.kind = "container"; container.start = u.base; container.end = u.base + u.sql.length();
                    out.hosts.add(new Host(container.id, sourceId, "container", "script", range(u, u.base, container.end), null, null, null, 0,
                            "source_only", Collections.<Target>emptyList(), null, null));
                    for (Draft h : u.hosts) {
                        List<Target> targets = new ArrayList<Target>();
                        copyTargets(targets, models.get(h.node));
                        if (!"select_item".equals(h.syntax)) copyTargets(targets, models.get(h.anchor));
                        if (h.node instanceof TResultColumn && ((TResultColumn)h.node).getExpr() != null
                                && h.syntax.contains("assignment")) {
                            TExpression lhs = ((TResultColumn)h.node).getExpr().getLeftOperand();
                            copyTargets(targets, models.get(lhs));
                            if (lhs != null) copyTargets(targets, models.get(lhs.getObjectOperand()));
                        }
                        if ("call".equals(h.kind)) targets.removeIf(t -> !"relationship".equals(t.space));
                        String processId = null, procedureId = null;
                        Draft parent = h.parent;
                        if (parent != null && (parent.start < u.base || parent.end > container.end)) parent = null;
                        Draft query = parent;
                        while (query != null && !"sstselect".equals(query.syntax)) query = query.parent;
                        for (Target t : targets) {
                            if ("process".equals(t.space)) processId = t.id;
                            if ("procedure".equals(t.space)) procedureId = t.id;
                        }
                        out.hosts.add(new Host(h.id, sourceId, h.kind, h.syntax, range(u,h.start,h.end), range(u,h.nameStart,h.nameEnd),
                                query == null ? null : query.id, parent == null ? container.id : parent.id, h.ordinal,
                                targets.isEmpty() ? "source_only" : "resolved", targets, processId, procedureId));
                    }
                    boolean limited = false; int count = 0;
                    for (TSourceToken token : u.comments) {
                        int a = CommentTokenIndex.start(token), b = CommentTokenIndex.end(token);
                        if (a < u.base || b > container.end) continue;
                        if (count >= out.limits.maxCommentsPerSourceUnit) { limited = true; break; }
                        count++;
                        String id = sourceId + ":c" + count;
                        if (!p.index.exact(token)) { diagnostic(out,"SOURCE_MAP_UNAVAILABLE",sourceId,id,null); continue; }
                        String raw = token.getAstext(); String lexical = lexicalKind(raw);
                        int length = raw.length(), keep = Math.min(length, Math.min(out.limits.maxRawLength, Math.max(0,out.limits.maxTotalRawLength-total)));
                        if (keep < length && keep > 0 && Character.isHighSurrogate(raw.charAt(keep-1))) keep--;
                        total += keep; String encoding = null;
                        if (invalidSurrogate(raw)) { diagnostic(out,"INVALID_INPUT",sourceId,id,null); raw = ""; keep = 0; }
                        else raw = raw.substring(0,keep);
                        boolean truncated = keep < length;
                        if (truncated) limited = true;
                        if (!xmlSafe(raw)) { raw = Base64.getEncoder().encodeToString(raw.getBytes(StandardCharsets.UTF_8)); encoding = "base64-utf8"; }
                        out.comments.add(new Comment(id, sourceId, range(u,a,b), lexical, raw, encoding, length, truncated));
                        if (!"line".equals(lexical) && !"block".equals(lexical)) continue;
                        CommentBinder.Choice choice = binder.bind(token,u.base,container.end,container);
                        if (choice == null) diagnostic(out,"NO_HOST",sourceId,id,null);
                        else if (choice.host == null) diagnostic(out,choice.rule,sourceId,id,null);
                        else out.bindings.add(new Binding(id,choice.host.id,choice.placement,choice.rule));
                    }
                    String state = (limited || option.isIgnoreCoordinate()) ? "limited" : p.partial ? "parse_partial" : u.derived ? "unsupported" : "complete";
                    for (String kind : KINDS) out.coverage.add(new Coverage(sourceId,kind,state));
                    if (!"complete".equals(state)) out.state = "partial";
                    if (limited) diagnostic(out,"LIMIT_REACHED",sourceId,null,null);
                    if (option.isIgnoreCoordinate()) diagnostic(out,"COORDINATES_DISABLED",sourceId,null,null);
                    if (u.derived) diagnostic(out,"SOURCE_MAP_UNAVAILABLE",sourceId,null,null);
                }
            }
            for (Unit u : units) if (!processed.contains(u)) {
                for (String kind : KINDS) out.coverage.add(new Coverage(u.source.id, kind, "parse_partial"));
                out.state = "partial";
                diagnostic(out, "PARSE_UNAVAILABLE", u.source.id, null, null);
            }
            if (missingSource) diagnostic(out,"SOURCE_MAP_UNAVAILABLE",null,null,null);
            if (failed) { out.state = "failed"; diagnostic(out,"COLLECTION_FAILED",null,null,null); }
        } catch (Exception e) { out.state = "failed"; diagnostic(out,"COLLECTION_FAILED",null,null,null); }
        finally { clear(); }
        return out;
    }
    private Range range(Unit u,int a,int b) { return option.isIgnoreCoordinate() ? null : u.index.range(a-u.base,b-u.base); }
    private static void diagnostic(CommentBindingBundle b,String code,String source,String comment,String host) {
        b.diagnostics.add(new Diagnostic(code,source,comment,host));
        if (code.equals("INVALID_INPUT") || code.equals("SOURCE_MAP_UNAVAILABLE")) b.state = "partial";
    }
    private static String lexicalKind(String raw) {
        if (raw.startsWith("/*+")) return "hint";
        if (raw.startsWith("/*!")) return "executable";
        if (raw.startsWith("/*M!")) return "executable";
        if (raw.startsWith("--+")) return "hint";
        if (raw.startsWith("-- " + TBaseType.sqlflow_stmt_delimiter_str)) return "directive";
        if (raw.startsWith("--") && ENGINE_DIRECTIVE.matcher(raw).find()) return "directive";
        if (raw.startsWith("--#") || raw.startsWith("/*#")) return "directive";
        return raw.startsWith("/*") ? "block" : "line";
    }
    private static boolean invalidSurrogate(String s) {
        for (int i=0;i<s.length();i++) {
            char c=s.charAt(i);
            if (Character.isHighSurrogate(c)) { if (++i>=s.length() || !Character.isLowSurrogate(s.charAt(i))) return true; }
            else if (Character.isLowSurrogate(c)) return true;
        }
        return false;
    }
    private static boolean xmlSafe(String s) {
        for (int i=0;i<s.length();i++) { char c=s.charAt(i); if ((c<32 && c!='\t' && c!='\r' && c!='\n') || c==0xfffe || c==0xffff) return false; }
        return true;
    }
    private static void copyTargets(List<Target> to,List<Target> from) {
        if (from != null) for (Target t : from) {
            boolean found = false;
            for (Target old : to) if (old.space.equals(t.space) && old.id.equals(t.id) && Objects.equals(old.parentId,t.parentId)) { found=true; break; }
            if (!found) to.add(t);
        }
    }
    private static void addTarget(IdentityHashMap<Object,List<Target>> map,Object key,Target target) {
        if (key != null) map.computeIfAbsent(key,k->new ArrayList<Target>()).add(target);
    }
    private static void observe(IdentityHashMap<Object,List<Target>> map,Object key,Object value) {
        if (key instanceof Pair) key = ((Pair<?,?>)key).second;
        if (value instanceof List) { for (Object v:(List<?>)value) observe(map,key,v); return; }
        Target target = null;
        if (value instanceof Table) target = new Target("table",String.valueOf(((Table)value).getId()),null);
        else if (value instanceof ResultSet) target = new Target("resultset",String.valueOf(((ResultSet)value).getId()),null);
        else if (value instanceof TableColumn) { TableColumn c=(TableColumn)value; target=new Target("column",String.valueOf(c.getId()),String.valueOf(c.getTable().getId())); }
        else if (value instanceof ResultColumn) { ResultColumn c=(ResultColumn)value; target=new Target("column",String.valueOf(c.getId()),String.valueOf(c.getResultSet().getId())); }
        else if (value instanceof Procedure) { Procedure p=(Procedure)value; target=new Target("procedure",String.valueOf(p.getId()),null); addTarget(map,p.getProcedureObject(),target); }
        else if (value instanceof OraclePackage) { OraclePackage p=(OraclePackage)value; target=new Target("package",String.valueOf(p.getId()),null); addTarget(map,p.getProcedureObject(),target); }
        else if (value instanceof gudusoft.gsqlparser.dlineage.dataflow.model.Process) target=new Target("process",String.valueOf(((gudusoft.gsqlparser.dlineage.dataflow.model.Process)value).getId()),null);
        if (target != null && key instanceof TParseTreeNode) addTarget(map,key,target);
    }
}