#!/usr/bin/env python
#This program is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with this program.  If not, see <http://www.gnu.org/licenses/>.
import os
import argparse
import signal
from collections import namedtuple


#Define a namedtuple to cleanly handle symbolic link data.
#Because we are performing potentially dangerous operations on the filesystem
#it is necessary to store data in a tuple. E.G. Write once, examine, execute.  
#... Can't afford to screw up.
Link = namedtuple('Link', ['path', 'destination'])


def readlinkabs(l):
    """Return an absolute path for the destination 
    of a symlink
    """
    if not os.path.islink(l):
        return ''
    p = os.readlink(l)
    if os.path.isabs(p):
        return p
    return os.path.abspath(os.path.join(os.path.dirname(l), p))

def generate_link_map(path):
    """Walk through "path" directory structure and record all symbolic links,
    and their destinations.
    
    return: list of Link namedtuples
    """
    original_paths = []
    link_paths = []
    links = []
    
    for root, dirnames, filenames in path:
        # For each directory
        for dirname in dirnames:
            dir_path = os.path.join(root, dirname)
            if readlinkabs(dir_path):
                original_paths.append(dir_path)
                link_paths.append(readlinkabs(dir_path))

            #For each file
            for filename in filenames:
                file_path = os.path.join(root, filename)
                if readlinkabs(file_path):
                    # The way os.walk nests filenames; ignore replicated entities
                    if file_path in original_paths or file_path in link_paths:
                        continue
                    #Store symbolic link paths
                    original_paths.append(file_path)
                    link_paths.append(readlinkabs(file_path))
    #Generate a list of Links
    for original_path, link_path in zip(original_paths, link_paths):
        links.append(Link(original_path, link_path))
    return links

def get_link_map(link_map):
    """return: yield Link namedtuples
    """
    for link in link_map:
        yield link

def generate_replacement_link_map(link_map, search, replace):
    """Find and replace occurrence of search parameter in link. 
    return: list of Link namedtuples
    """
    links = []
    for link in link_map:
        if search in link.destination:
            linkr = link.destination.replace(search, replace)
            if linkr:
                #This should never happen... but if it does, we're ready.
                if link.path == linkr:
                    print("Ignoring (dangerous) duplicated path: {}".format(link.path))
                    continue
                links.append(Link(link.path, linkr))
    return links

def replace_links(link_map):
    """Perform physical link replacement using link_map data.
    In order to update the symbolic links in the filesystem they
    must be removed and regenerated.
    return: None
    """
    for link in link_map:
        if not os.path.exists(link.path):
            print("WARNING :: Destination {0} does not exist.  May generate a dead link.".format(link.destination))
        
        print("Re-linking: {0} -> {1}".format(link.path, link.destination))
        os.unlink(link.path)
        os.symlink(link.destination, link.path)

def user_choice():
    """Would you like to play a game?
    """
    YES = ['YES', 'Y']
    choice = raw_input().upper()
    if choice not in YES:
        return False
    return True

def trap_sigint(sig, frame):
    """Perform any cleanup tasks (most likely none)
    """
    print("Caught signal {0}.  Exiting.".format(sig))
    exit(0)
    
if __name__ == "__main__":
    signal.signal(signal.SIGINT, trap_sigint)
    parser = argparse.ArgumentParser(description='Replace symbolic links')
    parser.add_argument('--force', action="store_true", help='Do not ask for confirmation')
    parser.add_argument('search', type=str, action="store", help='original path pattern')
    parser.add_argument('replace', type=str, action="store", help='replacement string in path pattern')
    parser.add_argument('path', type=str, action="store", help='path to perform search')
    args = parser.parse_args()
    
    if args.search == args.replace or \
        args.replace == args.search:
        print("Search and replace patterns cannot be the same.")
        exit(1)
    
    lmap = generate_link_map(os.walk(os.path.abspath(args.path), followlinks=False))
    if not lmap:
        print("No symbolic links detected under {0}".format(args.path))
        exit(0)
    
    print("### LINKS ####")
    for link in get_link_map(lmap):
        print("Link: {0} -> {1}".format(link.path, link.destination))

    rmap = generate_replacement_link_map(lmap, args.search, args.replace)
    if not rmap:
        print("Nothing to replace.".format(args.search, args.path))
        exit(0)
    
    print("")
    print("### PREVIEW ####")
    for link in get_link_map(rmap):
        print("Replacement: {0} -> {1}".format(link.path, link.destination))
    
    link_total = len(lmap)
    match_total = len(rmap)
    print("")
    print("{0} symbolic links total, {1} pattern matches".format(link_total, match_total))
    print("")
    if args.force:
        replace_links(rmap)
    else:
        print("Replace symbolic link(s)? [y/N]"),
        if user_choice():
            print("")
            replace_links(rmap)
        else:
            print("Abort.")
            exit(254)            
    exit(0)