package com.atlassian.confluence.extra.chart;

import com.atlassian.confluence.servlet.download.ExportDownload;
import com.atlassian.renderer.RenderContext;
import com.atlassian.renderer.v2.RenderMode;
import com.atlassian.renderer.v2.macro.MacroException;
import com.opensymphony.util.TextUtils;
import org.dom4j.Document;
import org.dom4j.DocumentException;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.CategoryPlot;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.chart.renderer.category.LineAndShapeRenderer;
import org.jfree.chart.renderer.category.StackedAreaRenderer;
import org.jfree.chart.renderer.category.StackedBarRenderer;
import org.jfree.chart.renderer.category.StackedBarRenderer3D;
import org.jfree.data.category.DefaultCategoryDataset;
import org.jfree.data.general.Dataset;
import org.jfree.data.general.DefaultPieDataset;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.io.StringReader;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * Generates a chart based on table data contained in its body.
 *
 * @author David Peterson
 */
public class ChartMacro extends
        com.atlassian.confluence.extra.chart.BaseChartMacro {

    private static final RenderMode RENDER_MODE = RenderMode
            .suppress(RenderMode.F_FIRST_PARA);

    private static final String TABLE = "TABLE";

    private static final String THEAD = "THEAD";

    private static final String TBODY = "TBODY";

    private static final String TFOOT = "TFOOT";

    private static final String TR = "TR";

//    private static final String TD = "TD";
//
//    private static final String TH = "TH";

    private static final int DEFAULT_WIDTH = 300;

    private static final int DEFAULT_HEIGHT = 300;

    // parameter values
    private static final String TRUE = "true";
    private static final String FALSE = "false";

    private static final String HORIZONTAL = "horizontal";
    private static final String VERTICAL = "vertical";

    // supported chart types
    private static final String PIE_TYPE = "pie";

    private static final String BAR_TYPE = "bar";

    private static final String LINE_TYPE = "line";

    private static final String AREA_TYPE = "area";

    /**
     * Render a {chart} macro.
     */
    public String execute(Map parameters, String body,
                          RenderContext renderContext) throws MacroException {
        try {
            parameters = toLowerCase(parameters);

            int width = DEFAULT_WIDTH;
            if (TextUtils.stringSet((String) parameters.get("width"))) {
                width = Integer.parseInt((String) parameters.get("width"));
            }

            int height = DEFAULT_HEIGHT;
            if (TextUtils.stringSet((String) parameters.get("height"))) {
                height = Integer.parseInt((String) parameters.get("height"));
            }

            String rendered = subRenderer.render(body, renderContext, RENDER_MODE);

            Document data = parseBody(rendered);

            BufferedImage image = getChart(parameters, data).createBufferedImage(width, height);

            File outputFile = ExportDownload.createTempFile("chart",
                    ".png");
            ImageIO.write(image, "PNG", outputFile);

            StringBuffer out = new StringBuffer();

            out.append("<img src=\"").append(ExportDownload.getUrl(outputFile, "image/png"))
                    .append("\" width=\"").append(width)
                    .append("\" height=\"").append(height).append("\"/>");

            if (TRUE.equals(parameters.get("displaydata")) || TRUE.equals(parameters.get("datadisplay")))
                out.append(rendered);

            return out.toString();
        }
        catch (IOException ioe) {
            throw new MacroException(ioe);
        }
        catch (DocumentException e) {
            throw new MacroException(e);
        }
    }

    /**
     * Get a chart object that represents the data provided using the parameters specified
     * <p/>
     * Supported parameter keys
     * title       - string
     * xLabel      - string (also allow xlabel)
     * yLabel      - string (also allow ylabel)
     * orientation - vertical (default), horizontal
     * legend      - true (default), false
     * stacked     - false (default), true
     * dataOrientation - horizontal false (default), true (also allow dataorientation)
     * type        - pie (default), pie3D, bar, bar3D, line, line3D, area
     */

    private JFreeChart getChart(Map parameters, Document data) throws MacroException {

        String title = (String) parameters.get("title");

        String xLabel = (String) parameters.get("xlabel");
        String yLabel = (String) parameters.get("ylabel");

        boolean tooltips = true;  // not a parameter yet
        boolean urls = false; // not a parameter yet

        String type = PIE_TYPE;   // default to type=pie
        if (TextUtils.stringSet((String) parameters.get("type")))
            type = (String) parameters.get("type");

        // default to show legend unless legend=false
        boolean legend = true;
        if (TextUtils.stringSet((String) parameters.get("legend")))
            legend = !FALSE.equalsIgnoreCase((String) parameters.get("legend"));

        // default to vertical orientation unless orientation=horizontal
        PlotOrientation plotOrientation = PlotOrientation.VERTICAL;
        if (TextUtils.stringSet((String) parameters.get("orientation"))
                && ((String) parameters.get("orientation")).equalsIgnoreCase(HORIZONTAL))
            plotOrientation = PlotOrientation.HORIZONTAL;

        // default to NOT 3D unless 3d=true OR 3D=true
        boolean is3d = false;
        if (TextUtils.stringSet((String) parameters.get("3d")))
            is3d = TRUE.equalsIgnoreCase((String) parameters.get("3d"));

        // default to NOT stacked unless stacked=true
        boolean stacked = false;  // default
        if (TextUtils.stringSet((String) parameters.get("stacked")))
            stacked = TRUE.equalsIgnoreCase((String) parameters.get("stacked"));

        // default to have columns represent the domain or x axis values unless dataOrientation=vertical
        boolean verticalDataOrientation = false;
        if (TextUtils.stringSet((String) parameters.get("dataorientation")))
            verticalDataOrientation = VERTICAL.equalsIgnoreCase((String) parameters.get("dataorientation"));

        type = type.toLowerCase();  // standardize on lowercase

        JFreeChart chart;

        if (PIE_TYPE.equals(type)) {
            DefaultPieDataset dataset = new DefaultPieDataset();
            processData(data, dataset, verticalDataOrientation);

            if (is3d)
                chart = ChartFactory.createPieChart3D(title, dataset, legend, tooltips, urls);
            else
                chart = ChartFactory.createPieChart(title, dataset, legend, tooltips, urls);
        } else if (BAR_TYPE.equals(type)) {
            DefaultCategoryDataset dataset = new DefaultCategoryDataset();
            processData(data, dataset, verticalDataOrientation);

            if (is3d)
                chart = ChartFactory.createBarChart3D(title, xLabel, yLabel, dataset, plotOrientation, legend, tooltips, urls);
            else
                chart = ChartFactory.createBarChart(title, xLabel, yLabel, dataset, plotOrientation, legend, tooltips, urls);

            if (stacked) {
                CategoryPlot plot = chart.getCategoryPlot();

                if (is3d)
                    plot.setRenderer(new StackedBarRenderer3D());
                else
                    plot.setRenderer(new StackedBarRenderer());
            }
        } else if (AREA_TYPE.equals(type)) {
            DefaultCategoryDataset dataset = new DefaultCategoryDataset();
            processData(data, dataset, verticalDataOrientation);

            chart = ChartFactory.createAreaChart(title, xLabel, yLabel, dataset, plotOrientation, legend, tooltips, urls);

            if (stacked) {
                CategoryPlot plot = chart.getCategoryPlot();
                plot.setRenderer(new StackedAreaRenderer());
            }
        } else if (LINE_TYPE.equals(type)) {
            DefaultCategoryDataset dataset = new DefaultCategoryDataset();
            processData(data, dataset, verticalDataOrientation);

            if (is3d)
                chart = ChartFactory.createLineChart3D(title, xLabel, yLabel, dataset, plotOrientation, legend, tooltips, urls);
            else
                chart = ChartFactory.createLineChart(title, xLabel, yLabel, dataset, plotOrientation, legend, tooltips, urls);

            CategoryPlot plot = chart.getCategoryPlot();
            LineAndShapeRenderer renderer = (LineAndShapeRenderer) plot.getRenderer();
            renderer.setShapesVisible(true); // default to show points in line graphs
        } else
            throw new MacroException("Unsupported chart type: " + type);

        return chart;
    }

    /**
     * Parses the body and returns a Dom4J Document.
     *
     * @param rendered The rendered macro body to parse.
     * @return The parsed document.
     * @throws DocumentException
     */
    private Document parseBody(String rendered)
            throws DocumentException {
        rendered = cleanHTML(rendered);
        SAXReader saxReader = new SAXReader();
        return saxReader.read(new StringReader("<data>" + rendered
                + "</data>"));
    }

    /**
     * @param rendered
     * @return The cleaned, XML-friendly HTML
     */
    private String cleanHTML(String rendered) {
        rendered = rendered.replaceAll("\\&nbsp;", "&#160;");
        return rendered;
    }

    /**
     * @param doc
     * @param dataset
     */
    private void processData(Document doc, Dataset dataset, boolean verticalDataOrientation) {
        Element data = doc.getRootElement();

        Iterator i = data.elements().iterator();
        boolean processed = false;

        while (i.hasNext() && !processed) {
            Element e = (Element) i.next();
            if (TABLE.equalsIgnoreCase(e.getName())) {
                processTableContent(e, dataset, null, verticalDataOrientation);
                processed = true;
            }
        }
    }

    private Element processTableContent(Element element, Dataset dataset, Element headerRow, boolean verticalDataOrientation) {
        Iterator i = element.elements().iterator();

        while (i.hasNext()) {

            Element e = (Element) i.next();

            if (THEAD.equalsIgnoreCase(e.getName())
                    || TBODY.equalsIgnoreCase(e.getName())
                    || TFOOT.equalsIgnoreCase(e.getName())) {
                headerRow = processTableContent(e, dataset, headerRow, verticalDataOrientation);
            } else if (TR.equalsIgnoreCase(e.getName())) {
                if (headerRow == null) {
                    headerRow = e;
                } else if (verticalDataOrientation) {
                    processVerticalDataRow(e, dataset, headerRow);
                } else {
                    processHorizontalDataRow(e, dataset, headerRow);
                }
            }
        }
        return headerRow;
    }

    // Columns represent domain or x values
    private void processHorizontalDataRow(Element row, Dataset dataset, Element headerRow) {

        List cells = row.elements();
        List headers = headerRow.elements();

        if (cells.size() > 1 && headers.size() > 1) {
            Iterator iCells = cells.iterator();
            Iterator iHeaders = headers.iterator();
            String domain = ((Element) iHeaders.next()).getTextTrim();
            String range = ((Element) iCells.next()).getTextTrim();

            if (dataset instanceof DefaultPieDataset) {
                DefaultPieDataset pieDataset = (DefaultPieDataset) dataset;
                while (iCells.hasNext() && iHeaders.hasNext()) {
                    Element key = (Element) iHeaders.next();
                    Element value = (Element) iCells.next();

                    pieDataset.setValue(key.getTextTrim(), toDouble(value
                            .getTextTrim()));
                }
            } else if (dataset instanceof DefaultCategoryDataset) {
                DefaultCategoryDataset catDataset = (DefaultCategoryDataset) dataset;

                while (iCells.hasNext() && iHeaders.hasNext()) {
                    Element key = (Element) iHeaders.next();
                    Element value = (Element) iCells.next();

                    catDataset.setValue(toDouble(value.getTextTrim()), range,
                            key.getTextTrim());
                }
            }
        }
    }

    // Rows represent domain or x values
    private void processVerticalDataRow(Element row, Dataset dataset, Element headerRow) {
        List cells = row.elements();
        List headers = headerRow.elements();
        String key;
        String value;

        if (cells.size() > 1 && headers.size() > 1) {

            if (dataset instanceof DefaultPieDataset) {

                DefaultPieDataset pieDataset = (DefaultPieDataset) dataset;

                key = ((Element) cells.get(0)).getTextTrim();
                value = ((Element) cells.get(1)).getTextTrim();
                pieDataset.setValue(key, toDouble(value));

            } else if (dataset instanceof DefaultCategoryDataset) {

                DefaultCategoryDataset catDataset = (DefaultCategoryDataset) dataset;
                Iterator iCells = cells.iterator();
                Iterator iHeaders = headers.iterator();
                String domain = ((Element) iHeaders.next()).getTextTrim();
                String range = ((Element) iCells.next()).getTextTrim();

                while (iCells.hasNext() && iHeaders.hasNext()) {
                    key = ((Element) iHeaders.next()).getTextTrim();
                    value = ((Element) iCells.next()).getTextTrim();

                    catDataset.setValue(toDouble(value), key, range);
                }
            }
        }
    }

    /**
     * @param value
     * @return
     * @throws NumberFormatException
     */
    private Double toDouble(String value) throws NumberFormatException {
        if (value == null)
            return new Double(0);

        value = value.replaceAll("[^0-9\\.]", "");

        if (value.length() == 0)
            return new Double(0);

        return new Double(value);
    }

    private Map toLowerCase(Map params) {
        Map lcParams = new java.util.HashMap();
        Iterator i = params.entrySet().iterator();
        Map.Entry e;
        while (i.hasNext()) {
            e = (Map.Entry) i.next();
            lcParams.put(((String) e.getKey()).toLowerCase(), e.getValue());
        }
        return lcParams;
    }
}
