package jmv;

import java.lang.reflect.*;
import java.util.*;
import java.util.function.*;
import java.util.stream.*;

/** Observes sequential stream traversal once; never enumerates a source to preview it. */
public final class StreamTrace {
    static final List<Chain> chains = new ArrayList<>();
    static final IdentityHashMap<Object, Handler> handles = new IdentityHashMap<>();
    static final Set<String> MAPPERS = Set.of("filter", "map", "mapToInt", "mapToLong", "mapToDouble", "mapToObj", "peek");
    static final Set<String> TERMINALS = Set.of("forEach", "forEachOrdered", "toArray", "toList", "collect", "reduce", "sum", "average", "min", "max", "count", "findFirst", "findAny", "anyMatch", "allMatch", "noneMatch", "summaryStatistics");
    static final Set<String> PASS = Set.of("boxed", "asLongStream", "asDoubleStream", "limit", "skip");
    static final Set<String> FLAT = Set.of("flatMap", "flatMapToInt", "flatMapToLong", "flatMapToDouble");
    // Mapper-created inner pipelines are represented by their actual flattened outputs.
    static int innerMapperDepth;
    static final class Node {
        int id, stage, row; Node parent; Cell cell;
        Node(int id,int stage,int row,Node parent,Cell cell){this.id=id;this.stage=stage;this.row=row;this.parent=parent;this.cell=cell;}
    }
    static final class Expansion {
        int stage; Node parent; String state="opening"; List<Integer> outputs=new ArrayList<>();
        Expansion(int stage,Node parent){this.stage=stage;this.parent=parent;}
    }
    static final class Sort {
        int stage; String state="waiting"; List<Integer> inputs=new ArrayList<>(), outputs=new ArrayList<>();
        Map<Object,ArrayDeque<Node>> pending;
        Sort(int stage,boolean references){this.stage=stage;pending=references?new IdentityHashMap<>():new HashMap<>();}
    }
    static final class Cell {
        Object value; String desc; boolean dropped;
        Cell(Object value, String desc, boolean dropped) { this.value=value; this.desc=desc; this.dropped=dropped; }
    }
    static final class Stage {
        String op, kind; int line;
        Stage(String op,String kind,int line) {this.op=op;this.kind=kind;this.line=line;}
    }
    static final class Chain {
        int id; String state="building", notice="";
        List<Stage> stages=new ArrayList<>();
        List<Map<Integer,Cell>> rows=new ArrayList<>();
        int row=-1, activeStage=0; boolean observing=true;
        Object result; String resultDesc="L"; boolean hasResult;
        boolean advanced;
        Node current;
        List<Node> nodes=new ArrayList<>();
        List<Expansion> expansions=new ArrayList<>();
        List<Sort> sorts=new ArrayList<>();
        Chain(int id){this.id=id;}
    }
    static void reset(){chains.clear();handles.clear();innerMapperDepth=0;}
    static int line(){return Tracer.frames.isEmpty()?0:Tracer.frames.get(Tracer.frames.size()-1).line;}
    static void event(Chain c,int stage){
        if(!Tracer.enabled || !c.observing)return;
        c.activeStage=stage;
        Tracer.step("step_line",null);
    }
    static Class<?> api(Object o){
        if(o instanceof IntStream)return IntStream.class;
        if(o instanceof LongStream)return LongStream.class;
        if(o instanceof DoubleStream)return DoubleStream.class;
        return Stream.class;
    }
    static String desc(Class<?> api){return api==IntStream.class?"I":api==LongStream.class?"J":api==DoubleStream.class?"D":"L";}
    static Object invoke(Method method,Object target,Object[] args)throws Throwable{
        try{return method.invoke(target,args);}catch(InvocationTargetException e){throw e.getCause();}
    }
    static Object innerInvoke(Method method,Object target,Object[] args)throws Throwable{
        innerMapperDepth++;
        try{return invoke(method,target,args);}finally{innerMapperDepth--;}
    }
    static Object peek(Object stream,Consumer<Object> consumer){
        if(stream instanceof IntStream)return ((IntStream)stream).peek(v->consumer.accept(v));
        if(stream instanceof LongStream)return ((LongStream)stream).peek(v->consumer.accept(v));
        if(stream instanceof DoubleStream)return ((DoubleStream)stream).peek(v->consumer.accept(v));
        return ((Stream<?>)stream).peek(consumer);
    }
    static Node record(Chain c,int stage,Object v,String d,boolean dropped){
        if(!Tracer.enabled||!c.observing)return null;
        if(c.nodes.size()>=2000){c.observing=false;c.notice="관측 노드 2,000개 이후의 기록은 생략했습니다.";return null;}
        Cell cell=new Cell(v,d,dropped);
        Node node=new Node(c.nodes.size()+1,stage,c.row,stage==0?null:c.current,cell);
        c.nodes.add(node);c.current=node;
        if(c.row>=0&&c.row<c.rows.size())c.rows.get(c.row).put(stage,cell);
        return node;
    }
    static void cell(Chain c,int stage,Object v,String d,boolean dropped){
        record(c,stage,v,d,dropped);
        event(c,stage);
    }
    public static Object source(Object stream,String label){
        if(!Tracer.enabled || innerMapperDepth>0 || stream==null || handles.containsKey(stream) || chains.size()>=64)return stream;
        Chain c=new Chain(chains.size()+1);chains.add(c);
        c.stages.add(new Stage(label,"source",line()));
        if(((BaseStream<?,?>)stream).isParallel()){
            c.observing=false;c.state="unavailable";c.notice="병렬 스트림의 요소 순서는 추적하지 않습니다.";
            return stream;
        }
        Class<?> type=api(stream);
        Object observed=peek(stream,v->{
            Tracer.tick();
            if(!Tracer.enabled||!c.observing)return;
            if(c.rows.size()>=200){c.observing=false;c.notice="요소 200개 이후의 기록은 생략했습니다.";return;}
            c.rows.add(new LinkedHashMap<>());c.row=c.rows.size()-1;
            cell(c,0,v,desc(type),false);
        });
        Object proxy=wrap(observed,type,c,0); event(c,0); return proxy;
    }
    static Object wrap(Object raw,Class<?> type,Chain c,int stage){
        Handler h=new Handler(raw,type,c,stage);
        Object proxy=Proxy.newProxyInstance(StreamTrace.class.getClassLoader(),new Class<?>[]{type},h);
        handles.put(proxy,h);return proxy;
    }
    static void finish(Chain c){
        for(Sort sort:c.sorts)sort.state="complete";
    }
    static Object flat(Chain c,int stage,Node parent,Object stream,Expansion expansion){
        if(stream==null){expansion.state="closed";event(c,stage);return null;}
        Handler embedded=handles.get(stream);
        if(embedded!=null){
            embedded.c.state="running";
            embedded.c.notice="체인 "+c.id+"의 flatMap에서 이 스트림을 소비합니다.";
            stream=embedded.raw;
        }
        Class<?> type=api(stream);
        Object observed=peek(stream,v->{
            Tracer.tick();
            if(!Tracer.enabled||!c.observing)return;
            c.current=parent;c.row=parent.row;
            Node node=record(c,stage,v,desc(type),false);
            if(node!=null)expansion.outputs.add(node.id);
            event(c,stage);
        });
        expansion.state="open";event(c,stage);
        return ((BaseStream<?,?>)observed).onClose(()->{
            expansion.state="closed";
            if(embedded!=null){embedded.c.state="complete";finish(embedded.c);}
            event(c,stage);
        });
    }
    static final class Handler implements InvocationHandler {
        final Object raw;final Class<?> type;final Chain c;final int stage;
        Handler(Object raw,Class<?> type,Chain c,int stage){this.raw=raw;this.type=type;this.c=c;this.stage=stage;}
        public Object invoke(Object self,Method method,Object[] args)throws Throwable{
            String op=method.getName();
            if(method.getDeclaringClass()==Object.class){
                if(op.equals("equals"))return self==args[0];
                if(op.equals("hashCode"))return System.identityHashCode(self);
                return type.getSimpleName()+" (stream chain "+c.id+")";
            }
            if(!Tracer.enabled)return StreamTrace.invoke(method,raw,args);
            if(op.equals("isParallel"))return StreamTrace.invoke(method,raw,args);
            if(op.equals("close"))return StreamTrace.invoke(method,raw,args);
            if(op.equals("sequential")||op.equals("onClose")){
                Object next=StreamTrace.invoke(method,raw,args);
                return next==raw?self:wrap(next,type,c,stage);
            }
            boolean intermediate=BaseStream.class.isAssignableFrom(method.getReturnType());
            if(c.stages.size()>=32){c.observing=false;c.notice="단계 32개 이후의 기록은 생략했습니다.";}
            int nextStage=c.stages.size();
            if(nextStage<32)c.stages.add(new Stage(op,intermediate?"intermediate":"terminal",line()));
            else nextStage=stage;
            final int s=nextStage;
            if(intermediate){
                Object[] actual=args==null?null:args.clone();
                Object target=raw;
                Sort sorting=null;
                if(op.equals("sorted")){
                    if(actual!=null&&actual.length==1)Objects.requireNonNull(actual[0]);
                    c.advanced=true;
                    sorting=new Sort(s,desc(type).equals("L"));c.sorts.add(sorting);
                    final Sort sort=sorting;
                    target=peek(raw,v->{
                        if(!Tracer.enabled||!c.observing)return;
                        sort.state="buffering";
                        sort.inputs.add(c.current.id);
                        sort.pending.computeIfAbsent(v,k->new ArrayDeque<>()).add(c.current);
                        event(c,s);
                    });
                }else if(FLAT.contains(op)&&actual!=null&&actual[0]!=null){
                    c.advanced=true;
                    Object fn=actual[0];Class<?> sam=method.getParameterTypes()[0];
                    actual[0]=Proxy.newProxyInstance(StreamTrace.class.getClassLoader(),new Class<?>[]{sam},(p,m,a)->{
                        if(!c.observing)return innerInvoke(m,fn,a);
                        if(c.expansions.size()>=200){c.observing=false;c.notice="flatMap 입력 200개 이후의 기록은 생략했습니다.";return innerInvoke(m,fn,a);}
                        Node parent=c.current;
                        Expansion expansion=new Expansion(s,parent);c.expansions.add(expansion);event(c,s);
                        Object out;
                        try{
                            out=innerInvoke(m,fn,a);
                        }catch(Throwable t){expansion.state="failed";event(c,s);throw t;}
                        return flat(c,s,parent,out,expansion);
                    });
                }else if(MAPPERS.contains(op)&&actual!=null&&actual[0]!=null){
                    Object fn=actual[0];Class<?> sam=method.getParameterTypes()[0];
                    actual[0]=Proxy.newProxyInstance(StreamTrace.class.getClassLoader(),new Class<?>[]{sam},(p,m,a)->{
                        Object out=StreamTrace.invoke(m,fn,a);
                        boolean drop=op.equals("filter")&&Boolean.FALSE.equals(out);
                        Object v=op.equals("filter")||op.equals("peek")?a[0]:out;
                        String d=op.equals("filter")||op.equals("peek")?desc(type):desc(method.getReturnType());
                        cell(c,s,v,d,drop);return out;
                    });
                }else if(!PASS.contains(op)){
                    c.observing=false;c.notice=op+"는 요소별 추적을 지원하지 않아 실제 실행 결과만 표시합니다.";
                }
                Object next=StreamTrace.invoke(method,target,actual);
                Class<?> nextType=api(next);
                if(sorting!=null){
                    final Sort sort=sorting;
                    next=peek(next,v->{
                        if(!Tracer.enabled||!c.observing)return;
                        ArrayDeque<Node> matches=sort.pending.get(v);
                        if(matches==null||matches.isEmpty()){
                            c.observing=false;c.notice="sorted의 출력과 입력을 대응하지 못해 요소 연결 기록을 중단했습니다.";return;
                        }
                        Node parent=matches.remove();c.current=parent;c.row=parent.row;
                        sort.state="emitting";
                        Node node=record(c,s,v,desc(nextType),false);
                        if(node!=null)sort.outputs.add(node.id);
                        event(c,s);
                    });
                }
                if(PASS.contains(op))next=peek(next,v->cell(c,s,v,desc(nextType),false));
                Object result=wrap(next,nextType,c,s);event(c,s);return result;
            }
            if(!TERMINALS.contains(op)){
                c.observing=false;c.notice=op+"로 전달된 순회는 추적하지 않습니다.";c.state="unavailable";
                return StreamTrace.invoke(method,raw,args);
            }
            c.state="running";c.hasResult=false;event(c,s);
            try{
                Object terminal=peek(raw,v->{Tracer.tick();cell(c,s,v,desc(type),false);});
                Object result=StreamTrace.invoke(method,terminal,args);
                c.result=result;c.resultDesc=Tracer.descriptor(method.getReturnType());c.hasResult=!c.resultDesc.equals("V");
                c.state="complete";
                finish(c);
                if(c.rows.isEmpty()&&op.equals("count")&&c.notice.isEmpty())c.notice="요소 방문 없이 count 결과가 계산될 수 있습니다. 방문하지 않은 요소는 표시하지 않습니다.";
                event(c,s);return result;
            }catch(Throwable t){c.state="failed";event(c,s);throw t;}
        }
    }
    static String describe(Object o){
        Handler h=handles.get(o);
        return h==null?null:h.type.getSimpleName()+" · 체인 "+h.c.id+" / "+h.c.stages.get(h.stage).op;
    }
    static String json(ArrayDeque<Object> pending){
        StringBuilder b=new StringBuilder("[");
        for(Chain c:chains){
            if(b.length()>1)b.append(',');
            b.append("{\"id\":").append(c.id).append(",\"state\":").append(Tracer.json(c.state));
            b.append(",\"notice\":").append(Tracer.json(c.notice)).append(",\"activeStage\":").append(c.activeStage).append(",\"activeRow\":").append(c.row);
            b.append(",\"stages\":[");
            for(int i=0;i<c.stages.size();i++){
                if(i>0)b.append(',');Stage s=c.stages.get(i);
                b.append("{\"op\":").append(Tracer.json(s.op)).append(",\"kind\":").append(Tracer.json(s.kind)).append(",\"line\":").append(s.line).append('}');
            }
            b.append("],\"rows\":[");
            for(int i=0;i<c.rows.size();i++){
                if(i>0)b.append(',');b.append('{');boolean first=true;
                for(Map.Entry<Integer,Cell> entry:c.rows.get(i).entrySet()){
                    if(!first)b.append(',');first=false;Cell v=entry.getValue();
                    b.append('"').append(entry.getKey()).append("\":{\"value\":").append(Tracer.encode(v.value,v.desc,pending)).append(",\"dropped\":").append(v.dropped).append('}');
                }b.append('}');
            }
            b.append(']');
            if(c.advanced){
                b.append(",\"flow\":{\"activeNode\":").append(c.current==null?"null":c.current.id).append(",\"nodes\":[");
                for(Node node:c.nodes){
                    if(node.id>1)b.append(',');
                    b.append("{\"id\":").append(node.id).append(",\"stage\":").append(node.stage).append(",\"row\":").append(node.row);
                    b.append(",\"parent\":").append(node.parent==null?"null":node.parent.id).append(",\"value\":").append(Tracer.encode(node.cell.value,node.cell.desc,pending));
                    b.append(",\"dropped\":").append(node.cell.dropped).append('}');
                }
                b.append("],\"expansions\":[");
                for(int i=0;i<c.expansions.size();i++){
                    if(i>0)b.append(',');Expansion x=c.expansions.get(i);
                    b.append("{\"stage\":").append(x.stage).append(",\"parent\":").append(x.parent.id).append(",\"state\":").append(Tracer.json(x.state)).append(",\"outputs\":").append(x.outputs).append('}');
                }
                b.append("],\"sorts\":[");
                for(int i=0;i<c.sorts.size();i++){
                    if(i>0)b.append(',');Sort x=c.sorts.get(i);
                    b.append("{\"stage\":").append(x.stage).append(",\"state\":").append(Tracer.json(x.state)).append(",\"inputs\":").append(x.inputs).append(",\"outputs\":").append(x.outputs).append('}');
                }
                b.append("]}");
            }
            if(c.hasResult)b.append(",\"result\":").append(Tracer.encode(c.result,c.resultDesc,pending));
            b.append('}');
        }
        return b.append(']').toString();
    }
}
