/*
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */
package org.dbunitng.data;

import java.io.File;
import java.io.FileOutputStream;
import java.io.OutputStream;

import org.apache.commons.lang.StringUtils;
import org.dbunit.IDatabaseTester;
import org.dbunit.JdbcDatabaseTester;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.csv.CsvDataSetWriter;
import org.dbunit.dataset.datatype.IDataTypeFactory;
import org.dbunit.dataset.excel.XlsDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunitng.annotations.FileType;
import org.dbunitng.exception.DbUnitNGRuntimeException;
import org.dbunitng.listeners.internal.DbUnitNGConfig;

/**
 * テストデータファイル作成クラス。 DBからデータを抽出し、テーブルごとにファイルを作成する。
 * 
 * @author jyukutyo
 * 
 */
public class TestDataFileMaker {

	private DbUnitNGConfig config;

	private IDatabaseTester tester;

	/**
	 * コンストラクタ。DB接続に必要な値を引数にとる。
	 * 
	 * @param driver
	 *            JDBCドライバクラス
	 * @param url
	 *            JDBC接続URL
	 * @param userName
	 *            DBユーザー名
	 * @param password
	 *            DBパスワード
	 * @param schema
	 *            DBスキーマ
	 */
	public TestDataFileMaker(String driver, String url, String userName,
			String password, String schema) {
		config =
			DbUnitNGConfig
				.create(driver, url, password, userName, schema, null);
		tester = new JdbcDatabaseTester(driver, url, userName, password);
		if (StringUtils.isNotBlank(schema)) {
			tester.setSchema(schema.toUpperCase());
		}
	}

	/**
	 * DBからすべてのテーブルを抽出し、ファイルに出力する。
	 * 
	 * @param targetDirectory
	 *            ファイルを出力するディレクトリ(存在しなければディレクトリを作成する)
	 * @param type
	 *            ファイルの種類
	 */
	public void extractAllTables(String targetDirectory, FileType type) {
		try {
			IDatabaseConnection connection = tester.getConnection();
			String[] tableNames = connection.createDataSet().getTableNames();
			extractTables(targetDirectory, tableNames, type);
		} catch (Exception e) {
			throw new DbUnitNGRuntimeException(e);
		}
	}

	/**
	 * DBから指定したテーブルを抽出し、ファイルに出力する。
	 * 
	 * @param targetDirectory
	 *            ファイルを出力するディレクトリ(存在しなければディレクトリを作成する)
	 * @param tableNames
	 *            抽出するテーブル
	 * @param type
	 *            ファイルの種類
	 */
	public void extractTables(String targetDirectory, String[] tableNames,
			FileType type) {
		IDatabaseConnection connection;
		try {
			connection = tester.getConnection();

			DatabaseConfig databaseConfig = connection.getConfig();
			IDataTypeFactory factory = config.getFactory();
			if (factory != null) {
				databaseConfig.setProperty(
					DatabaseConfig.PROPERTY_DATATYPE_FACTORY,
					factory);
			}
			File dir = new File(targetDirectory);
			dir.mkdirs();

			if (FileType.CSV == type) {
				CsvDataSetWriter.write(
					connection.createDataSet(tableNames),
					new File(targetDirectory));

				return;
			}

			for (int i = 0; i < tableNames.length; i++) {
				String tableName = tableNames[i];

				IDataSet partialDataSet =
					connection.createDataSet(new String[] { tableName });

				OutputStream out =
					new FileOutputStream(targetDirectory + "/"
						+ tableName.toUpperCase() + "." + type.toString());

				if (FileType.XML == type) {
					FlatXmlDataSet.write(partialDataSet, out);
				} else if (FileType.EXCEL == type) {
					XlsDataSet.write(partialDataSet, out);
				}
			}
		} catch (Exception e) {
			throw new DbUnitNGRuntimeException(e);
		}
	}

}
