import java.io.*;
import javax.swing.*;
import java.awt.*;
import java.awt.image.*;
import java.awt.event.*;

class ImagePanel extends JPanel {
    private EasyBufferedImage image;

    public EasyBufferedImage getImage() {
	return image;
    }

    ImagePanel(EasyBufferedImage im) {
	image = im;
	if(image != null) {
	    setPreferredSize(new Dimension(image.getWidth(), image.getHeight()));
	}
    }

    public void paintComponent(Graphics g) {
	Graphics2D gin = (Graphics2D)g;
	int xOffset = 0;
	int yOffset = 0;
	if(image != null) {
	    xOffset = (getWidth() - image.getWidth())/2;
	    yOffset = (getHeight() - image.getHeight())/2;
	}

	gin.setColor(Color.white);
	gin.fillRect(0, 0, getWidth(), getHeight());

	if(image != null) {
	    gin.drawImage(image, xOffset, yOffset, this);
	}
    }

    public void setImage(EasyBufferedImage im) {
	image = im;
	if(image != null) { 
	    setPreferredSize(new Dimension(image.getWidth(), image.getHeight()));
	}
	repaint();
    }
}


class StegoMan extends JPanel implements ActionListener {
    JButton exitButton;
    JFileChooser chooser;
    ImagePanel imagePanel, messagePanel;
    EasyBufferedImage theImage;
    HideMessagePanel hideMessagePanel;
    ExtractMessagePanel extractMessagePanel;

    public JFrame getParentFrame(Container me) {
	Container comp = me.getParent();
	if(comp instanceof JFrame) return (JFrame)comp;
	else return getParentFrame(comp);
    }

    public void load(File f) {
	try {
	    theImage = EasyBufferedImage.createImage(f);
	    hideMessagePanel.setCover(theImage);
	    extractMessagePanel.setStego(theImage);
	} catch(Exception e) {
	    return;
	}

	JFrame top = getParentFrame(this);
	top.pack();
	top.repaint();
    }

    private void hideMessage(EasyBufferedImage cover, EasyBufferedImage message) {
	try {
	Steganographer steg = new LSBSteganographer(3);
	EasyBufferedImage stegoImage = steg.embed(cover, message);

	cover.show("cover");
	message.show("message");

	stegoImage.show("stego");
	
	EasyBufferedImage result = (EasyBufferedImage)steg.extract(stegoImage);
	result.show("extracted message");
	} catch(Exception e){ }
	
    }

    private static class PNGFileFilter extends javax.swing.filechooser.FileFilter {
	public boolean accept(File f) {
	    if(f.isDirectory()) return true;

	    String name = f.getName();
	    int hasDot = name.lastIndexOf('.');
	    if(hasDot >= 0) {
		return name.substring(hasDot+1).equalsIgnoreCase("PNG");
	    } else return false;
	}

	public String getDescription() {
	    return "PNG image files";
	}
    }

    public void actionPerformed(ActionEvent ev) {
	if(ev.getSource() == exitButton) {
	    System.exit(0);
	}
    }

    private class ExtractMessagePanel extends JPanel implements ActionListener {
	ImagePanel extractPanel;
	Steganographer steganographer = new LSBSteganographer(3);
	JButton extractMessageButton;

	public void actionPerformed(ActionEvent ev) {
	    EasyBufferedImage result = null;
	    try {
		result = (EasyBufferedImage)steganographer.extract(extractPanel.getImage());
	    } catch(IOException oops) { return; }
	    if(result != null) {
		JFrame newStego = new JFrame("Steganographer");
		newStego.getContentPane().add(new StegoMan(result));
		newStego.pack();
		newStego.show();
	    } else {
		// notify that no message exists
		JOptionPane.showMessageDialog(this, "No hidden message found", "Extraction Error", JOptionPane.WARNING_MESSAGE);
	    }
	}
	
	public void setStego(EasyBufferedImage im) { 
	    extractPanel.setImage(im);
	}

	ExtractMessagePanel(EasyBufferedImage input) {
	    setLayout(new BorderLayout());
	    extractPanel = new ImagePanel(input);
	    add(new JScrollPane(extractPanel), BorderLayout.CENTER);
	    extractMessageButton = new JButton("Extract Embedded Message");
	    extractMessageButton.addActionListener(this);
	    add(extractMessageButton, BorderLayout.SOUTH);
	}
    }

    private class HideMessagePanel extends JPanel implements ActionListener {
	private ImagePanel coverPanel, messagePanel;
	private JButton hideMessageButton, loadCoverButton, loadMessageButton, switchButton;
	private JSpinner channelSpinner;
	private LSBSteganographer steganographer = new LSBSteganographer(3);

	public void setCover(EasyBufferedImage im) { 
	    coverPanel.setImage(im);
	}

	public void setMessage(EasyBufferedImage im) { 
	    messagePanel.setImage(im);
	}

	public void actionPerformed(ActionEvent ev) {
	    if(ev.getSource() == loadCoverButton ){
		if(chooser == null) {
		    chooser = new JFileChooser();
		    chooser.setFileFilter(new PNGFileFilter());
		}
		
		int returnVal = chooser.showOpenDialog(getParent());
		if(returnVal == JFileChooser.APPROVE_OPTION) {
		    loadCover(chooser.getSelectedFile());
		}
	    } if(ev.getSource() == hideMessageButton) {
		JFrame app = new JFrame("StegoMan");
		steganographer.setNumChannels(((Integer)channelSpinner.getValue()).intValue());
		try {
		    app.getContentPane().add(new StegoMan(steganographer.embed(coverPanel.getImage(), messagePanel.getImage())));
		} catch(IOException oops){ 
		}
		
		app.pack();
		app.show();
	    } else if(ev.getSource() == loadMessageButton) {
		if(chooser == null) {
		    chooser = new JFileChooser();
		    chooser.setFileFilter(new PNGFileFilter());
		}
		
		int returnVal = chooser.showOpenDialog(getParent());
		if(returnVal == JFileChooser.APPROVE_OPTION) {
		    loadMessage(chooser.getSelectedFile());
		}
	    } else if(ev.getSource() == switchButton) {
		EasyBufferedImage tmp = coverPanel.getImage();
		coverPanel.setImage(messagePanel.getImage());
		messagePanel.setImage(tmp);
	    }
	}

	public void loadCover(File f) {
	    try {
		EasyBufferedImage cover = EasyBufferedImage.createImage(f);
		coverPanel.setImage(cover);
		hideMessageButton.setEnabled(true);
	    } catch(IOException oops) {
	    }
	}


	public void loadMessage(File f) {
	    try {
		EasyBufferedImage message = EasyBufferedImage.createImage(f);
		messagePanel.setImage(message);
		hideMessageButton.setEnabled(true);
	    } catch(IOException oops) {
	    }
	}

	HideMessagePanel(EasyBufferedImage input, EasyBufferedImage message) {
	    setLayout(new BorderLayout());
	    coverPanel = new ImagePanel(input);
	    messagePanel = new ImagePanel(message);
	    JSplitPane splitter = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, 
						 new JScrollPane(coverPanel), 
						 new JScrollPane(messagePanel));

	    JPanel controls = new JPanel();
	    controls.add(new JLabel("Number of channels: "));
	    SpinnerNumberModel model = new SpinnerNumberModel(3, 1, 7, 1);
	    channelSpinner = new JSpinner(model);
	    controls.add(channelSpinner);

	    loadCoverButton = new JButton("Load Cover");
	    loadCoverButton.addActionListener(this);
	    controls.add(loadCoverButton);

	    loadMessageButton = new JButton("Load Message");
	    loadMessageButton.addActionListener(this);
	    controls.add(loadMessageButton);


	    hideMessageButton = new JButton("Hide Message in Cover");
	    hideMessageButton.setEnabled(false);
	    hideMessageButton.addActionListener(this);
	    controls.add(hideMessageButton);

	    switchButton = new JButton("Switch cover and message");
	    switchButton.addActionListener(this);
	    controls.add(switchButton);

	    add(splitter, BorderLayout.CENTER);
	    add(controls, BorderLayout.SOUTH);
	    validate();
	    splitter.setDividerLocation(.5);
	}
    }

    StegoMan(EasyBufferedImage image) {
	setLayout(new BorderLayout());
	theImage = image;

	JTabbedPane pane = new JTabbedPane();
	hideMessagePanel = new HideMessagePanel(theImage, null);
	extractMessagePanel = new ExtractMessagePanel(theImage);
	pane.add("Embed Message", hideMessagePanel);
	pane.add("Extract Message", extractMessagePanel);
	
	JPanel controls = new JPanel();
	controls.setLayout(new GridLayout(1, 3));

	exitButton = new JButton("Exit");
	exitButton.addActionListener(this);

	controls.add(exitButton);
	add(controls, BorderLayout.NORTH);
	add(pane, BorderLayout.CENTER);
    }

    public static void main(String[] args) {
	JFrame app = new JFrame("StegoMan");
	app.getContentPane().add(new StegoMan(null));

	app.pack();
	app.show();
    }
}