#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
exemple usage: ./mirrorlist.py -o mirrorlist.txt uncomment $(./mirrorlist.py sections | dmenu -i -l 3)
"""

import sys
import argparse
import os
import contextlib


@contextlib.contextmanager
def output_context(output):
    if output == '-':
        yield sys.stdout
    else:
        yield open(output, 'w')


def split_head(lines):
    head = []
    body = []

    target = head

    for line in lines:
        if target is head:
            if not line.startswith('##'):
                target = body

        target.append(line)

    return head, body


def list_sections(args):
    with open(args.mirrorlist_path, 'r') as mirrorlist:
        lines = mirrorlist.readlines()

    head, body = split_head(lines)

    with output_context(args.output) as output:
        for line in body:
            if line.startswith('##'):
                section = line.replace('#', '').strip()
                output.write('{}\n'.format(section))


def uncomment_sections(args):
    with open(args.mirrorlist_path, 'r') as mirrorlist:
        lines = mirrorlist.readlines()

    head, body = split_head(lines)

    with output_context(args.output) as output:
        for line in head:
            output.write(line)

        uncomment = False
        for line in body:
            if line.startswith('##'):
                section = line.replace('#', '').strip()

                if section in args.sections:
                    uncomment = True

                output.write(line)
            elif line.startswith('#'):
                if uncomment:
                    output.write(line.replace('#', ''))
                else:
                    output.write(line)
            elif not line.strip():
                uncomment = False
                output.write(line)
            else:
                output.write(line)


def parse_args(argv=sys.argv):
    parser = argparse.ArgumentParser(description='Manage archlinux mirrorlist')
    parser.add_argument('-p', '--path',
                        dest='mirrorlist_path',
                        default='/etc/pacman.d/mirrorlist',
                        help='mirrorlist path')
    parser.add_argument('-o', '--output',
                        dest='output',
                        default='-',
                        help='result output (default: stdout)')

    subparsers = parser.add_subparsers(help='sub-command help')

    sections_parser = subparsers.add_parser('sections', help='sections help')
    sections_parser.set_defaults(action=list_sections)

    uncomment_parser = subparsers.add_parser('uncomment', help='uncomment help')
    uncomment_parser.set_defaults(action=uncomment_sections)
    uncomment_parser.add_argument('sections', nargs='+',
                                  help='sections to uncomment')

    return parser.parse_args(argv[1:])


if __name__ == "__main__":
    args = parse_args()
    args.action(args)