/*
Copyright (c) 2007 Joseph Gleason

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Current versions of this and other code can be downloaded at:
http://gleason.cc/

*/
package cc.glsn.v15.argus;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.LinkedList;
import java.util.Random;
import java.util.Scanner;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

import cc.glsn.ConfigFile;
import cc.glsn.Util;
import cc.glsn.v15.DBConnMgr;
import cc.glsn.v15.SQLConnection;
import cc.glsn.v15.ParityFileTool;
import cc.glsn.v15.StringID;

public class ArgusParityMaint
{

    /**
     * @param args
     * @throws SQLException 
     * @throws IOException 
     */
    public static void main(String[] args) throws SQLException, IOException
    {
        new ArgusParityMaint(args);
    }
    
    
    private TreeMap<ParitySetID,ParitySet> paritySets;
    private TreeMap<ImageID,ImageData> images;
    private TreeMap<DriveID,File> drives;
    private TreeMap<DriveID,Double> drivePriority;
    private TreeSet<DriveID> noplace_drives;
    private DBConnMgr dbcm;
    private ConfigFile cf;
    
    public ArgusParityMaint(String[] args) throws SQLException, IOException
    {
        cf=new ConfigFile("/argus/config/argus.config");
        
        images=new TreeMap<ImageID,ImageData>();
        paritySets=new TreeMap<ParitySetID,ParitySet>();
        drives=new TreeMap<DriveID,File>();
        drivePriority=new TreeMap<DriveID,Double>();
        noplace_drives=new TreeSet<DriveID>();
        
        dbcm=new DBConnMgr(10,
                cf.getString("SQLDriver"),
                cf.getString("SQLUrl"),
                cf.getString("SQLUsername"),
                cf.getString("SQLPassword"));
        
        
        readDrivePriority();
        readDBTables();
        scanDrives();

        
        if ((args.length==0) || (args[0].equals("make")))
        {
            lookForMissingImages(true);
            lookForMissingParitySets(true);
            makeParity();
        }
        else if (args[0].equals("report"))
        {
            lookForMissingImages(false);
            lookForMissingParitySets(false);
            statusReport();
        }

        
    }
    public void readDrivePriority() throws FileNotFoundException
    {
        Scanner scan=new Scanner(new FileInputStream("/argus/config/prilist"));
        
        while(scan.hasNext())
        {
            DriveID did=new DriveID(scan.next());
            double pri=scan.nextDouble();
            drivePriority.put(did, pri);
        }
        
    }
    
    private void readDBTables() throws SQLException
    {
        SQLConnection c=dbcm.getConn();
        
        {
            ResultSet r=c.doSingleQuery("select * from images");
            while(r.next())
            {
                ImageData img=new ImageData(r);
                images.put(img.getID(),img);                
            }
            r.close();
        
        }
        {
            ResultSet r=c.doSingleQuery("select * from paritysets");
            while(r.next())
            {
                ParitySet ps=new ParitySet(r);
                paritySets.put(ps.getSetID(), ps);
            }
            r.close();
            
        }
        {
            ResultSet r=c.doSingleQuery("select * from paritymembers");
            while(r.next())
            {
                ParitySetID psID=new ParitySetID(r.getString("setid"));
                ImageID imgID=new ImageID(r.getString("image_name"));
                
                ParitySet ps=paritySets.get(psID);
                if (ps==null)
                {
                    System.out.println("Unknown set in paritymembers: " + psID);
                }
                else
                {
                    ps.addMember(imgID);
                }
                
                ImageData imgData=images.get(imgID);
                if (imgData==null)
                {
                    System.out.println("Unknown image in paritymembers: " + imgID);
                }
                else
                {
                    imgData.addParitySet(psID);
                }
            }
            
        }
        
        dbcm.relConn(c);
        
        
    }
    
    
    private void scanDrives() throws FileNotFoundException, IOException, SQLException
    {
        File mntRoot=new File("/argus/mnt");
        for(File sub : mntRoot.listFiles())
        {
            if (sub.isDirectory())
            {
                if (sub.getTotalSpace() != mntRoot.getTotalSpace())
                {
                    DriveID did=new DriveID(sub.getName());
                    drives.put(did, sub);
                    //System.out.println("Drive: " + did);
                    scanSingleDrive(did,sub);
                }
            }
        }
    }
    
    private void scanSingleDrive(DriveID did, File drvFile) throws FileNotFoundException, IOException, SQLException
    {
        for(File sub : drvFile.listFiles())
        {
            if (sub.getName().endsWith(".ISO"))
            {
                ImageID imgID=new ImageID(sub.getName());
                if (!images.containsKey(imgID))
                {
                    System.out.println("Found new image: " + imgID);
                    ImageData imgData=new ImageData(sub);
                    
                    saveImageDB(imgData);
                    images.put(imgID,imgData);                 
                }
                images.get(imgID).addDrive(did);
                
            }
            if (sub.getName().endsWith(".PAR"))
            {
                ParitySetID setID=new ParitySetID(sub.getName());
                
                if (paritySets.containsKey(setID))
                {
                    paritySets.get(setID).addDrive(did);
                }
                else
                {
                    System.out.println("Unknown set found on " + did.getID() + ": " + setID.getID());
                }
                
                
                
            }
            if (sub.getName().equals("noplace"))
            {
                noplace_drives.add(did);
               
            }
            
        }
    }
    private void saveImageDB(ImageData imgData) throws SQLException
    {
        SQLConnection c=dbcm.getConn();
        
        ResultSet r=c.doSingleUpdatableQuery("select * from images limit 1");
        r.moveToInsertRow();
        
        r.updateString("name", imgData.getID().getID());
        r.updateString("md5", imgData.getMd5());
        r.updateLong("size", imgData.getSize());
        
        r.insertRow();
        r.close();
        
        dbcm.relConn(c);
        
        
    }
    
    private DriveID findDriveWithMostFreeSpace(Set<DriveID> exclusionList)
    {
       
        DriveID bestDrv=null;
        long mostFree=0;
        for(DriveID did : drives.keySet())
        {
            if ((exclusionList==null) || (!exclusionList.contains(did)))
            {
                long free=drives.get(did).getFreeSpace();
               
                if (!noplace_drives.contains(did))
                if (free > mostFree)
                {
                    mostFree=free;
                    bestDrv=did;
                }
            }
            
        }
        return bestDrv;
        
        
    }
    
    
   
    
    private void makeParity()
    {
        while(true)
        {
            TreeSet<DriveID> exclusionSet=new TreeSet<DriveID>();
            
            Random rng=new Random();
            
            TreeMap<Double,ImageData> potentialImages=new TreeMap<Double,ImageData>();
            for(ImageData imgData : images.values())
            {
                if (imgData.getParityCount()==0)
                {
                    double sz=rng.nextDouble() + imgData.getSize();
                    if (imgData.getDrives().size()>1)
                    {
                        sz=sz * -1.0;
                    }
                    potentialImages.put(sz,imgData);
                }
            }
            System.out.println("Non-protected files: " + potentialImages.size());
            
            TreeMap<ImageID,DriveID> deleteImages=new TreeMap<ImageID,DriveID>();
            TreeMap<ImageID,DriveID> selectedImages=new TreeMap<ImageID,DriveID>();
            
            while(selectedImages.size()<4)
            {
                if (potentialImages.size() ==0) return;
                
                double last=potentialImages.lastKey();
                ImageData img=potentialImages.get(last);
                potentialImages.remove(last);
                
                TreeSet<DriveID> availDrvs=img.getDrives();
                availDrvs.removeAll(exclusionSet);
                if (availDrvs.size() > 0)
                {
                    DriveID selectedDrv=selectHighestPriority(availDrvs);
                    
                    TreeSet<DriveID> totalDrv=img.getDrives();
                    totalDrv.remove(selectedDrv);
                    for(DriveID did : totalDrv)
                    {
                        deleteImages.put(img.getID(), did);
                    }
                    selectedImages.put(img.getID(),selectedDrv);
                    exclusionSet.add(selectedDrv);
                }
                
            }
            DriveID targetDrv=findDriveWithMostFreeSpace(exclusionSet);
            if (targetDrv==null) return;
            
            exclusionSet.add(targetDrv);
            
            
            
            System.out.println("Building parity of: ");
            for(ImageID i : selectedImages.keySet())
            {
                DriveID did=selectedImages.get(i);
                System.out.println("  " + did.getID() + "/" + i.getID() + " " + images.get(i).getSize());
            }
            
            System.out.println("Target: " + targetDrv.getID());
            
            if (deleteImages.size()>0)
            {
                System.out.println("Will be able to remove: ");
                for(ImageID i : deleteImages.keySet())
                {
                    DriveID did=deleteImages.get(i);
                    System.out.println("  " + did.getID() + "/" + i.getID() + " " + images.get(i).getSize());
                }
            }
            
             
            
            if (createParity(selectedImages,targetDrv))
            {
                for(ImageID i : deleteImages.keySet())
                {
                    DriveID did=deleteImages.get(i);
                    File f=getFile(did,i);
                    System.out.println("  Deleting file: " + f);
                    f.delete();
                }
            }
                    
            
            
        }
        
        
    }    
    
    private boolean createParity(TreeMap<ImageID,DriveID> selectedImages, DriveID targetDrv)
    {
        long len=0;
        LinkedList<File> inputFiles=new LinkedList<File>();
        for(ImageID imgID : selectedImages.keySet())
        {
            DriveID did=selectedImages.get(imgID);
            File f=getFile(did,imgID);
            
            inputFiles.add(f);
            
            len = Math.max(len, images.get(imgID).getSize());            
        }
        
        ParitySetID setID=genParitySetID();        
        File outputFile=getFile(targetDrv,new ParitySetID(setID.getID() + ".tmp"));
        
        ParityFileTool tool=new ParityFileTool();
        TreeMap<File,String> md5Sums=new TreeMap<File,String>();
        
        outputFile.deleteOnExit();
        tool.createParityFile(inputFiles, outputFile, md5Sums,0);
        outputFile.deleteOnExit();
        
        
        for(File f : md5Sums.keySet())
        {
            if (f.getName().endsWith(".ISO"))
            {
                String md5=md5Sums.get(f);
                ImageID img=new ImageID(f.getName());
                if (!md5.equals(images.get(img).getMd5()))
                {
                    System.out.println("MD5 read mismatch on " + img.getID());
                    System.out.println("Expected: " + images.get(img).getMd5());
                    System.out.println("Got:      " + md5);
                    outputFile.delete();
                    return false;               
                }
                else
                {
                    System.out.println("MD5 match on " + img.getID());
                }
            }
        }
        
        try
        {
            String outputMD5=Util.MD5File(outputFile);
            if (!outputMD5.equals(md5Sums.get(outputFile)))
            {
                System.out.println("MD5 write mismatch on " + outputFile);
                System.out.println("Expected: " + md5Sums.get(outputFile));
                System.out.println("Got:      " + outputMD5);
                outputFile.delete();
                return false;               
                
            }
            else
            {
                System.out.println("MD5 match on " + setID.getID());
            }
            
        }
        catch (IOException e1)
        {
            System.out.println("Error while reading " + outputFile);
            // TODO Auto-generated catch block
            e1.printStackTrace();
            outputFile.delete();
            return false;
        }
        
        
        
        ParitySet ps=new ParitySet(setID,selectedImages.keySet(),len,md5Sums.get(outputFile));
        ps.addDrive(targetDrv);
        
        try
        {
            SQLConnection c=dbcm.getConn();
            ps.saveDB(c);
            dbcm.relConn(c);
        }
        catch(SQLException e)
        {
            e.printStackTrace();
            outputFile.delete();
            return false;
            
        }
        
        try
        {
            PrintStream prnts = new PrintStream(new FileOutputStream("/argus/md5/" + ps.getSetID().getID() + ".md5"));
            prnts.println(ps.getMd5());
            prnts.flush();
            prnts.close();
        }
        catch(FileNotFoundException e)
        {
            e.printStackTrace();
            outputFile.delete();
            return false;
        }
        
        outputFile.renameTo(getFile(targetDrv,setID));
        paritySets.put(ps.getSetID(),ps);
        for(ImageID imgID : selectedImages.keySet())
        {
            images.get(imgID).addParitySet(ps.getSetID());
        }
        
        return true;
        
        
    }
    
    private ParitySetID genParitySetID()
    {
        Random rng=new Random();
        String id=Util.SHA1("" + rng.nextDouble());
        id=id.substring(0, 10);
        
        
        return new ParitySetID(id + ".PAR");
    }
    
    private File getFile(DriveID drv, StringID id)
    {
        return new File("/argus/mnt/" + drv.getID() + "/" + id.getID());
    }
    
    private double getDrivePriority(DriveID drv)
    {
        if (drivePriority.containsKey(drv)) return drivePriority.get(drv);
        
        return 100.0;
    }
    private DriveID selectHighestPriority(Set<DriveID> drvLst)
    {
        double pri=0.0;
        DriveID did=null;
        for(DriveID d : drvLst)
        {
            double p =getDrivePriority(d);
            if (p>pri)
            {
                pri=p;
                did=d;
            }
        }
        return did;
        
    }
    
    private void lookForMissingImages(boolean doSomethingAboutIt)
    {
        for(ImageData img : images.values())
        {
            if (img.getDrives().size()==0)
            {
                System.out.println("Image " + img.getID().getID() + " is not anywhere.");
                
            
            
                if (attemptRecover(img,doSomethingAboutIt))
                {
                    if (doSomethingAboutIt)
                    {
                        System.out.println("  Image " + img.getID().getID() + " successfully recovered.");
                    }
                    else
                    {
                        System.out.println("  Image " + img.getID().getID() + " should be recoverable.");
                    }
                    
                    
                }
                else
                {
                    System.out.println("  Unable to recover image " + img.getID().getID());
                    
                }
            
                
            }
            
        }
    }
    
    private boolean attemptRecover(ImageData img,boolean doSomethingAboutIt)
    {
        for(ParitySetID psid : img.getParitySet())
        {
            
            if (attemptRecover(img,psid,doSomethingAboutIt)) return true;
            
        }
        return false;
        
    }
    private boolean attemptRecover(ImageData restoreImg, ParitySetID psid,boolean doSomethingAboutIt)
    {
        if (doSomethingAboutIt)
            System.out.println("  Trying to recover " + restoreImg.getID().getID() + " using " + psid.getID());
        ParitySet ps=paritySets.get(psid);
        
        if (ps.getDrives().size()==0)
        {
            System.out.println("    " + psid.getID() +" is not on any drives.");
            return false;
        }
        
        TreeSet<DriveID> excludeList=new TreeSet<DriveID>();
        excludeList.addAll(ps.getDrives());
        
        for(ImageID id : ps.getImages())
        {
            if (!id.equals(restoreImg.getID()))
            {
                ImageData other=images.get(id);
                if (other.getDrives().size()==0)
                {
                    System.out.println("    " + psid.getID() + " cannot run.  Missing " + id.getID());
                    return false;
                }
                excludeList.addAll(other.getDrives());
            }
                    
        }
        
        if (!doSomethingAboutIt) return true;
        
        DriveID targetDrv=findDriveWithMostFreeSpace(excludeList);
        if (targetDrv==null)
        {
            System.out.println("    " + "unable to find target drive.");
            return false;
            
        }
        
        long len=0;
        LinkedList<File> inputFiles=new LinkedList<File>();
        
        inputFiles.add(getFile(selectHighestPriority(ps.getDrives()),psid));
        
        
        for(ImageID imgID : ps.getImages())
        {
            if (!imgID.equals(restoreImg.getID()))
            {
                ImageData dat=images.get(imgID);
                DriveID did=selectHighestPriority(dat.getDrives());
                File f=getFile(did,imgID);                
                inputFiles.add(f);
            }
            
                   
        }
        
              
        File outputFile=getFile(targetDrv,new ImageID(restoreImg.getID().getID() + ".tmp"));
        
        ParityFileTool tool=new ParityFileTool();
        TreeMap<File,String> md5Sums=new TreeMap<File,String>();
        
        outputFile.deleteOnExit();
        System.out.println("input: " + inputFiles);
        System.out.println("outputFile: " +  outputFile);
        System.out.println("size: " + restoreImg.getSize());
        tool.createParityFile(inputFiles, outputFile, md5Sums,restoreImg.getSize());
        outputFile.deleteOnExit();
        
        
        for(File f : md5Sums.keySet())
        {
            if (f.getName().endsWith(".ISO"))
            {
                String md5=md5Sums.get(f);
                ImageID img=new ImageID(f.getName());
                if (!md5.equals(images.get(img).getMd5()))
                {
                    System.out.println("    MD5 read mismatch on " + img.getID());
                    System.out.println("    Expected: " + images.get(img).getMd5());
                    System.out.println("    Got:      " + md5);
                    outputFile.delete();
                    return false;               
                }
                else
                {
                    System.out.println("    MD5 match on " + img.getID());
                }
            }
            if (f.getName().endsWith(".PAR"))
            {
                String md5=md5Sums.get(f);
                ParitySetID sub_psid=new ParitySetID(f.getName());
                
                if (!md5.equals(paritySets.get(sub_psid).getMd5()))
                {
                    System.out.println("    MD5 read mismatch on " + sub_psid.getID());
                    System.out.println("    Expected: " + images.get(sub_psid).getMd5());
                    System.out.println("    Got:      " + md5);
                    outputFile.delete();
                    return false;               
                }
                else
                {
                    System.out.println("    MD5 match on " + sub_psid.getID());
                }
            }            
        }
        
        
        try
        {
            String outputMD5=Util.MD5File(outputFile);
            if (!outputMD5.equals(restoreImg.getMd5()))
            {
                System.out.println("    MD5 write mismatch on " + outputFile);
                System.out.println("    Expected: " + restoreImg.getMd5());
                System.out.println("    Got:      " + outputMD5);
                outputFile.delete();
                return false;               
                
            }
            else
            {
                System.out.println("    MD5 match on " + restoreImg.getID().getID());
            }
            
        }
        catch (IOException e1)
        {
            System.out.println("Error while reading " + outputFile);
            // TODO Auto-generated catch block
            e1.printStackTrace();
            outputFile.delete();
            return false;
        }
        
        
        outputFile.renameTo(getFile(targetDrv,restoreImg.getID()));
        
        restoreImg.addDrive(targetDrv);
        
        
        return true;
        
        
    }
    
    
    
    
    private void lookForMissingParitySets(boolean doSomethingAboutIt)
    {
        for(ParitySet ps : paritySets.values())
        {
            if (ps.getDrives().size()==0)
            {
                System.out.println("ParitySet " + ps.getSetID().getID() + " is not anywhere.");
            }
            
            
            TreeSet<DriveID> usedDrives=new TreeSet<DriveID>();
            
            usedDrives.addAll(ps.getDrives());
            
            for(ImageID i : ps.getImages())
            {
                if (images.get(i)!=null)
                for(DriveID d : images.get(i).getDrives())
                {
                    if (usedDrives.contains(d))
                    {
                        System.out.println("ParitySet " + ps.getSetID().getID() + " has multiple elements on drive " + d.getID());
                    }
                    usedDrives.add(d);
                            
                }
            }
            
        }
            
            
        
    }
    
    public void statusReport()
    {
        for(ImageData img : images.values())
        {
            if ((img.getParityCount()!=1) || (img.getDrives().size() != 1))
            {
                System.out.println(img.getStatus());
            }
            
        }
        
        
    }


}