package com.ibm.wala.viz.viewer;

import java.awt.BorderLayout;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import javax.swing.JPanel;
import javax.swing.JScrollPane;
import javax.swing.JSplitPane;
import javax.swing.JTextField;
import javax.swing.JTree;
import javax.swing.event.TreeExpansionEvent;
import javax.swing.event.TreeExpansionListener;
import javax.swing.event.TreeSelectionEvent;
import javax.swing.event.TreeSelectionListener;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.TreeNode;
import javax.swing.tree.TreePath;

import com.ibm.wala.analysis.pointers.HeapGraph;
import com.ibm.wala.ipa.callgraph.CGNode;
import com.ibm.wala.ipa.callgraph.CallGraph;
import com.ibm.wala.ipa.callgraph.propagation.InstanceFieldPointerKey;
import com.ibm.wala.ipa.callgraph.propagation.InstanceKey;
import com.ibm.wala.ipa.callgraph.propagation.LocalPointerKey;
import com.ibm.wala.ipa.callgraph.propagation.NormalAllocationInNode;
import com.ibm.wala.ipa.callgraph.propagation.PointerAnalysis;
import com.ibm.wala.ipa.callgraph.propagation.PointerKey;
import com.ibm.wala.ipa.callgraph.propagation.ReturnValueKey;
import com.ibm.wala.ssa.IR;
import com.ibm.wala.ssa.SSAInstruction;
import com.ibm.wala.util.intset.MutableMapping;
import com.ibm.wala.util.intset.OrdinalSetMapping;

/**
 * Panel for showing the Pointer Analysis results. Shows the heap graph on the left, and the ir viewer on the right. Sets the IR of
 * and pc according to the chosen pointer/instance key when possible (e.g., allocation side for NormalAllocationInNode instance
 * keys. Can be customized for clients that use different their own pointer / instance keys (see JsPaPanel)
 * 
 * Note the two steps initialization require (calling init())
 * @author yinnonh
 * 
 */
public class PaPanel extends JSplitPane {

  protected final PointerAnalysis pa;
  protected final CallGraph cg;

  private JTextField fullName;

  private IrAndSourceViewer irViewer;

  private MutableMapping<List<LocalPointerKey>> cgNodeIdToLocalPointers = MutableMapping.<List<LocalPointerKey>> make();
  private MutableMapping<List<ReturnValueKey>> cgNodeIdToReturnValue = MutableMapping.<List<ReturnValueKey>> make();
  private MutableMapping<List<InstanceFieldPointerKey>> instanceKeyIdToInstanceFieldPointers = MutableMapping.<List<InstanceFieldPointerKey>> make();


  public PaPanel(CallGraph cg, PointerAnalysis pa) {
    super(JSplitPane.HORIZONTAL_SPLIT);

    this.pa = pa;
    this.cg = cg;

    initDataStructures(pa);
  }

  /**
   * Two steps initialization is required here is our deriver can choose the roots for the heap tree.
   */
  public void init() {
    this.setDividerLocation(250);

    DefaultMutableTreeNode root = new DefaultMutableTreeNode();
    for (Object rootChildNode : getRootNodes()){
      DefaultMutableTreeNode n = new DefaultMutableTreeNode(rootChildNode);
      root.add(n);
      expandNodeRec(n, 1);
    }
    
    JTree heapTree = new JTree(root);
    
    heapTree.setCellRenderer(new DualTreeCellRenderer());

    this.setLeftComponent(new JScrollPane(heapTree));
    JPanel rightPanel = new JPanel(new BorderLayout());
    this.setRightComponent(rightPanel);
    fullName = new JTextField("");
    rightPanel.add(fullName, BorderLayout.PAGE_START);
    irViewer = new IrAndSourceViewer(cg);
    rightPanel.add(irViewer.getComponent(), BorderLayout.CENTER);

    heapTree.addTreeExpansionListener(new TreeExpansionListener() {

      public void treeExpanded(TreeExpansionEvent event) {
        TreePath path = event.getPath();
        if (path == null) {
          return;
        }
        DefaultMutableTreeNode lastNode = (DefaultMutableTreeNode) path.getLastPathComponent();
        expandNodeRec(lastNode, 2);
      }

      public void treeCollapsed(TreeExpansionEvent event) {
      }
    });

    heapTree.addTreeSelectionListener(new TreeSelectionListener() {

      public void valueChanged(TreeSelectionEvent e) {
        TreePath newLeadSelectionPath = e.getNewLeadSelectionPath();
        if (null == newLeadSelectionPath){
          return;
        }
        DefaultMutableTreeNode treeNode = (DefaultMutableTreeNode) newLeadSelectionPath.getLastPathComponent();
        Object userObject = treeNode.getUserObject();
        fullName.setText(userObject.toString());
        if (userObject instanceof LocalPointerKey){
          LocalPointerKey lpk = (LocalPointerKey) userObject;
          IR ir = lpk.getNode().getIR();
          SSAInstruction def = lpk.getNode().getDU().getDef(lpk.getValueNumber());
          int pc = IrViewer.NA;
          if (def != null){
            SSAInstruction[] instructions = ir.getInstructions();
            for (int i = 0; i < instructions.length; i++) {
              SSAInstruction instruction = instructions[i];
              if (def == instruction){
                pc = i;
              }
            }
          }
          irViewer.setIRAndPc(ir, pc);
        } else if (userObject instanceof InstanceFieldPointerKey){
          InstanceKey ik = ((InstanceFieldPointerKey) userObject).getInstanceKey();
          if (ik instanceof NormalAllocationInNode){
            NormalAllocationInNode normalIk = (NormalAllocationInNode) ik;
            IR ir = normalIk.getNode().getIR();
            int pc = normalIk.getSite().getProgramCounter();
            irViewer.setIRAndPc(ir, pc);
          }
        } else if (userObject instanceof NormalAllocationInNode){
          NormalAllocationInNode normalIk = (NormalAllocationInNode) userObject;
          IR ir = normalIk.getNode().getIR();
          int pc = normalIk.getSite().getProgramCounter();
          irViewer.setIRAndPc(ir, pc);
        } else if (userObject instanceof CGNode){
          irViewer.setIR(((CGNode)userObject).getIR());
        }


      }
    });

  }

  private void initDataStructures(PointerAnalysis pa) {
    HeapGraph heapGraph = pa.getHeapGraph();
    OrdinalSetMapping<InstanceKey> instanceKeyMapping = pa.getInstanceKeyMapping();
    for (Object n : heapGraph){
      if (heapGraph.getPredNodeCount(n) == 0){ // considering only roots of the heap graph.
        if (n instanceof PointerKey){
          if (n instanceof LocalPointerKey){
            LocalPointerKey lpk = (LocalPointerKey) n;
            int nodeId = lpk.getNode().getGraphNodeId();
            mapUsingMutableMapping(cgNodeIdToLocalPointers, nodeId, lpk);
          } else if (n instanceof ReturnValueKey){
            ReturnValueKey rvk = (ReturnValueKey) n;
            int nodeId = rvk.getNode().getGraphNodeId();
            mapUsingMutableMapping(cgNodeIdToReturnValue, nodeId, rvk);
          } else if (n instanceof InstanceFieldPointerKey){
            InstanceFieldPointerKey ifpk = (InstanceFieldPointerKey) n;
            int instanceKeyId = instanceKeyMapping.getMappedIndex(ifpk.getInstanceKey());
            mapUsingMutableMapping(instanceKeyIdToInstanceFieldPointers, instanceKeyId, ifpk);
          }
        } else {
          System.err.println("Non Pointer key root: " + n);
        }
      }
    }
  }


  /**
   * Override if you want different roots for your heap tree.
   * @return
   */
  protected List<Object> getRootNodes(){
    List<Object> ret = new ArrayList<Object>();
    for (CGNode n : cg){
      ret.add(n);
    }
    return ret;
  }

  /**
   * expands the given "treeNode" "rec" levels.
   * @param treeNode
   * @param rec
   */
  private void expandNodeRec(DefaultMutableTreeNode treeNode, int rec) {
    if (rec == 0){
      return;
    }

    if (treeNode.getChildCount() == 0){ // may be expandable.
      List<Object> children = getChildrenFor(treeNode.getUserObject());
      for (Object child : children){
        treeNode.add(new DefaultMutableTreeNode(child));
      }
    }

    for (int i = 0 ; i < treeNode.getChildCount(); i++){
      TreeNode child = treeNode.getChildAt(i);
      expandNodeRec((DefaultMutableTreeNode)child, rec-1);
    }
  }

  /**
   * Used for filling the tree dynamically. Override and handle your own nodes / different links.
   * @param node
   * @return
   */
  protected List<Object> getChildrenFor(Object node) {
    List<Object> ret = new ArrayList<Object>();
    if (node instanceof InstanceKey){
      ret.addAll(getPointerKeysUnderInstanceKey((InstanceKey) node));
    } else if (node instanceof PointerKey){
      for (InstanceKey ik : pa.getPointsToSet((PointerKey) node)){
        ret.add(ik);
      }
    } else if (node instanceof CGNode){
      int nodeId = ((CGNode) node).getGraphNodeId();
      ret.addAll(nonNullList(cgNodeIdToLocalPointers.getMappedObject(nodeId)));
      ret.addAll(nonNullList(cgNodeIdToReturnValue.getMappedObject(nodeId)));
    } else {
      assert false : "Unhandled Node : " + node;
    }
    return ret;
  }

  /**
   * Get the set of pointer keys that should be presented below an instance key in the heap tree. Override if you have special
   * pointer keys (not just for fields)
   * 
   * @param ik
   * @return
   */
  protected List<? extends PointerKey> getPointerKeysUnderInstanceKey(InstanceKey ik) {
    int ikIndex = pa.getInstanceKeyMapping().getMappedIndex(ik);
    List<? extends PointerKey> ret;
    if (ikIndex <= instanceKeyIdToInstanceFieldPointers.getMaximumIndex()){
      ret = nonNullList(instanceKeyIdToInstanceFieldPointers.getMappedObject(ikIndex));
    } else {
      ret = Collections.emptyList();
    }
    return ret;
  }


  /**
   * Utility method for mutable mapping. map[index] U= o
   * @param <T>
   * @param map
   * @param index
   * @param o
   */
  protected static <T> void mapUsingMutableMapping(MutableMapping<List<T>> map, int index, T o){
    List<T> set;
    if (index <= map.getMaximumIndex()){
      set = map.getMappedObject(index);
    } else {
      set = null;
    }
    if (null == set){
      set = new ArrayList<T>();
      map.put(index, set);
    }
    set.add(o);
  }

  protected <T> List<T> nonNullList(List<T> l){
    if (null == l){
      return Collections.emptyList();
    } else {
      return l;
    }
  }
}
