package org.batgen.generators;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import org.batgen.Column;
import org.batgen.ForeignNode;
import org.batgen.Table;

public class ForeignKeyGenerator {
    private ArrayList<ForeignNode> foreignNodeList;
    private HashMap<String, Table> tableMap;

    public ForeignKeyGenerator( ArrayList<ForeignNode> foreignNodeList, HashMap<String, Table> tableMap ) {
        this.foreignNodeList = foreignNodeList;
        this.tableMap = tableMap;
    }

    public String createForeignKeys() {
        StringBuilder sb = new StringBuilder();
        boolean fromFieldExist = false;
        boolean toFieldExist = false;
        Table fromTable;
        Table toTable;
        String fromField = "";
        String toField = "";

        for ( ForeignNode node : foreignNodeList ) {
            fromTable = tableMap.get( node.getFromTable() );
            toTable = tableMap.get( node.getToTable() );
            fromFieldExist = false;
            toFieldExist = false;

            if ( fromTable != null && toTable != null ) {
                for ( Column col : fromTable.getColumns() ) {
                    if ( col.getFldName().equals( node.getFromField() ) ) {
                        fromFieldExist = true;
                        fromField = col.getColName();
                        break;
                    }
                }
                for ( Column col : toTable.getColumns() ) {
                    if ( col.getFldName().equals( node.getToField() ) ) {
                        toFieldExist = true;
                        toField = col.getColName();
                        break;
                    }
                }

                if ( fromFieldExist && toFieldExist ) {
                    sb.append( "ALTER TABLE " + fromTable.getTableName() );
                    sb.append( " ADD CONSTRAINT FK_" + fromTable.getTableName() + "_" + fromField );
                    sb.append( " FOREIGN KEY (" + fromField + ") " );
                    sb.append( "REFERENCES " + toTable.getTableName() + "(" + toField + ");\n" );
                }
                else
                    throw new IllegalArgumentException( "In Table" + node.getFromTable() + " and/or "
                            + node.getToTable() + ", either the field names are wrong or don't exist for foreign keys." );
            }
            else
                throw new IllegalArgumentException( "Either the tables names ( " + node.getFromTable() + " and/or "
                        + node.getToTable() + " ) are wrong or don't exist for foreign keys." );
        }

        sb.append( "\n-- PROTECTED CODE -->" );
        List<String> lines = GenUtil.getProtectedLines( "sql/_AlterTables.sql" );
        if ( lines.isEmpty() ) {
            sb.append( "\n" );
        }
        for ( String line : lines ) {
            sb.append( line );
        }
        return sb.toString();
    }

    public String createDropForeignKeys() {
        StringBuilder sb = new StringBuilder();
        boolean fromFieldExist = false;
        Table fromTable;
        String fromField = "";

        for ( ForeignNode node : foreignNodeList ) {
            fromTable = tableMap.get( node.getFromTable() );
            fromFieldExist = false;
            for ( Column col : fromTable.getColumns() ) {
                if ( col.getFldName().equals( node.getFromField() ) ) {
                    fromFieldExist = true;
                    fromField = col.getColName();
                    break;
                }
            }

            if ( fromFieldExist ) {
                sb.append( "ALTER TABLE " + fromTable.getTableName() );
                sb.append( " DROP CONSTRAINT FK_" + fromTable.getTableName() + "_" + fromField + ";\n" );
            }
        }
        return sb.toString();
    }
}

@@ -527,29 +527,388 @@
  public 
-String create
+void createForeignKeys() %7B%0A        StringBuilder sb = new StringBuilder();%0A        sb.append( writeAddForeignKeys() );%0A        GenUtil.writeToFile( %22sql/_AlterTables.sql%22, sb.toString() );%0A        sb = new StringBuilder();%0A        sb.append( writeDropForeignKeys() );%0A        GenUtil.appendToFile( %22sql/_DropTables.sql%22, sb.toString() );%0A    %7D%0A%0A    private String writeAdd
 ForeignK
@@ -3312,21 +3312,22 @@
    p
-ublic
+rivate
  String 
 crea
@@ -3322,20 +3322,19 @@
  String 
-crea
+wri
 teDropFo

Update ForeignKeyGenerator.java