#-*- coding: utf-8 -*-

# Copyright 2010-2013 Calculate Ltd. http://www.calculate-linux.org
#
#  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.

import os
import re
import sys
import time
from os import path
from StringIO import StringIO
from random import choice
import string
from time import sleep
from subprocess import PIPE,STDOUT
from shutil import copy2
from calculate.core.server.task_set import TaskSet
from calculate.lib.utils.files import (runOsCommand,pathJoin,
                                       isMount,process,listDirectory,STDOUT,
                                       checkUtils,readFile,find,copyWithPath,
                                       readLinesFile)
from calculate.lib.utils.common import (appendProgramToEnvFile,
                        removeProgramToEnvFile, getTupleVersion,
                        cmpVersion,getUserPassword,
                        getSupportArch, getInstalledVideo )
from calculate.lib.utils.device import (detectDeviceForPartition,
                        getUdevDeviceInfo, getLvmPartitions, refreshLVM,
                        refreshUdev,countPartitions)
from calculate.lib.cl_vars_share import varsShare
from calculate.lib import cl_overriding
from calculate.lib.utils import ip
from calculate.lib.datavars import VariableError
from datavars import DataVarsInstall, __version__,__app__
from calculate.install.variables.autopartition import (AutopartitionError,
    AutoPartition)
from calculate.install.fs_manager import FileSystemManager

from calculate.lib.variables.locale import Locale
from calculate.lib.cl_template import Template,TemplatesError,ProgressTemplate
from calculate.lib.datavars import DataVars
from distr import (PartitionDistributive,
                DistributiveError, ScratchPartitionDistributive,
                MultiPartitions, FlashDistributive, 
                Distributive)
from calculate.lib.utils.text import tableReport
from calculate.lib.server.utils import dialogYesNo
from subprocess import Popen,PIPE,STDOUT
from itertools import *

class InstallError(Exception):
    """Installation Error"""

from migrate_users import migrate, currentUsers, MigrationError
from calculate.lib.encrypt import encrypt

_ = lambda x:x
from calculate.lib.cl_lang import setLocalTranslate,getLazyLocalTranslate
setLocalTranslate('cl_install3',sys.modules[__name__])
__ = getLazyLocalTranslate(_)

class Install(TaskSet):
    """Primary class for templates appling and system installation"""

    def init(self):
        # refresh information about LVM
        refreshLVM()
        # refresh information about device in udevadm info
        refreshUdev()

    def initVars(self,datavars=None):
        """Primary initialization of variables"""
        if not datavars:
            self.clVars = DataVarsInstall()
            self.clVars.importInstall()
            self.clVars.flIniFile()
        else:
            self.clVars = datavars

    def canInstallGrub2(self,target):
        """Check that system has grub2 in current and installed system"""
        if self.clVars.Get('os_grub2_path'):
            return bool(
                   filter(lambda x:x.startswith('grub-1.99') or \
                                   x.startswith('grub-2'),
                   listDirectory('/var/db/pkg/sys-boot')))
        return False

    def prepareBoot(self,targetDistr):
        """Prepare system for boot"""
        if self.clVars.Get('os_install_root_type') == "flash":
            self.installSyslinuxBootloader(targetDistr)
        else:
            if self.canInstallGrub2(targetDistr):
                self.installGrub2Bootloader(targetDistr)
            else:
                self.installLegacyGrubBootloader(targetDistr)
        return True

    def setActivePartition(self,partition):
        """
        Установка активного раздела для dos и gpt таблицы разделов
        """
        deviceName = detectDeviceForPartition(partition)
        if deviceName is None:
            raise DistributiveError(
                _("Failed to determine the parent device for %s")%partition)
        # device hasn't any partition
        elif deviceName == "":
            return True
        fdiskProg, gdiskProg, partedProg = checkUtils('/sbin/fdisk',
                    '/usr/sbin/gdisk','/usr/sbin/parted')

        disk = self.clVars.Select('os_install_disk_parent',
                                  where='os_install_disk_dev',eq=partition,
                                  limit=1)
        parttable = self.clVars.Select('os_device_table',where='os_device_dev',
                                  eq=disk,limit=1)
        partitionNumber = \
            getUdevDeviceInfo(name=partition).get('ID_PART_ENTRY_NUMBER','') or \
            getUdevDeviceInfo(name=partition).get('UDISKS_PARTITION_NUMBER','')

        devicePartitionCount = countPartitions(deviceName)
        if deviceName and not partitionNumber:
            raise DistributiveError(
                _("Failed to determine the partition number for %s")%partition)
        bootFlag = "boot" if parttable == "dos" else "legacy_boot"
        if parttable == "dos":
            fdisk = process(fdiskProg, "-l",deviceName)
            DEVICENUM,AFLAG = 0,1
            changeActive = \
                map(lambda x:x[DEVICENUM],
                filter(lambda x:x[DEVICENUM] != partitionNumber and \
                                x[AFLAG] == "*" or \
                                x[DEVICENUM] == partitionNumber and \
                                not x[AFLAG] == "*",
                list(map(lambda x:[str(x[0]),x[1][1].strip()],
                # enumerate partitions
                enumerate(filter(None,
                map(lambda x:x.split()[:2],
                # drop string before information about partitions
                dropwhile(lambda x:not x.lstrip().startswith("Device"),
                fdisk.readlines()))))))[1:]))
        else:
            parted = process(partedProg, "-m",deviceName,"print")
            DEVICENUM,FLAGS = 0,6
            changeActive = \
                map(lambda x:x[DEVICENUM],
                filter(lambda x:x[DEVICENUM] != partitionNumber and \
                                bootFlag in x[FLAGS].strip(';').split(', ') or \
                                x[DEVICENUM] == partitionNumber and \
                                not bootFlag in x[FLAGS].strip(';').split(', '),
                filter(lambda x:len(x)>=7,
                map(lambda x:x.split(':'),
                parted.readlines()[2:]))))
        if not changeActive:
            return True
        if parttable == "dos":
            pipe = Popen([fdiskProg,deviceName],
                stdin=PIPE, stdout=PIPE,stderr=PIPE)
            for partnum in changeActive:
                pipe.stdin.write("a\n%s\n"%partnum)
            pipe.stdin.write("w\n")
            pipe.stdin.close()
            pipe.wait()
        elif parttable == "gpt":
            pipe = Popen([gdiskProg,deviceName],
                stdin=PIPE, stdout=PIPE,stderr=PIPE)
            if devicePartitionCount > 1:
                pipe.stdin.write("x\n")
                for partnum in changeActive:
                    pipe.stdin.write("a\n%s\n2\n\n"%partnum)
                pipe.stdin.write("w\nY\n")
            else:
                pipe.stdin.write("x\na\n2\n\nw\nY\n")
            pipe.stdin.close()
            pipe.wait()
        for waittime in (0.1,0.2,0.5,1,2,4):
            if path.exists(partition):
                return True
            else:
                sleep(waittime)
        raise InstallError(
            _("Failed to find partition %s after changing the activity")%
                partition)

    def installSyslinuxBootloader(self,target):
        """
        Установить syslinux загрузчик (используется для flash)
        """
        if not self.clVars.Get('os_install_mbr'):
            return
        # прописать MBR
        ddProcess = process("/bin/dd","if=/usr/share/syslinux/mbr.bin",
                            "of=%s"%self.clVars.Get('os_install_mbr')[0],
                            stderr=STDOUT)
        if ddProcess.failed():
            raise DistributiveError(
                _("Failed to write the master boot record\n%s")%
                            ddProcess.read())
        target.close()
        # выполнить установку syslinux загрузчика
        installRootDev = self.clVars.Get('os_install_root_dev')
        syslinuxProcess = process("/usr/bin/syslinux",
                                  installRootDev, stderr=STDOUT)
        if syslinuxProcess.failed():
            raise DistributiveError(_("Failed to install syslinux\n%s")%
                            syslinuxProcess.read())
        # установить загрузочный раздел активным
        return self.setActivePartition(self.clVars.Get('os_install_root_dev'))

    def installGrub2Bootloader(self,target):
        """
        Установка GRUB2 загрузчика
        """
        # проверить наличие grub2
        cmdGrubInstall = self.clVars.Get('os_grub2_path')
        if not cmdGrubInstall:
            raise DistributiveError(_("Failed to install the bootloader"))
        process("sync").success()
        # если установка GRUB2 производится на текущую систему
        # загруженную в builder режиме
        if self.clVars.Get('os_install_scratch') == "on" and \
            self.clVars.Get('cl_action') != "system":
            prefixBoot = "/mnt/scratch"
        else:
            prefixBoot = "/"
        # установка UEFI
        if self.clVars.GetBool('os_install_uefi_set'):
            grubParams = ["--boot-directory=%s"%pathJoin(prefixBoot,
                              target.getBootDirectory()),
                              "--target=x86_64-efi",
                              "--efi-directory=%s"%
                              target.getEfiDirectory(),
                              "-f"]
            # в случае установки на usb-hdd EFI загрузчик не прописывается
            # в efivars
            if self.clVars.Get('os_install_root_type') == 'usb-hdd':
                grubParams.append("--removable")
            if self.clVars.Get('cl_action') != "system" and \
                not isMount('/boot/efi'):
                raise DistributiveError(_("Failed to install the bootloader. "
                                          "/boot/efi is not mounted."))
            grubProcess = process(cmdGrubInstall,
                              *grubParams, stderr=STDOUT,envdict=os.environ)
            if grubProcess.failed():
                raise DistributiveError(_("Failed to install the bootloader"))
            # проверяем успешность создания загрузочной записи
            # если среди загрузочных записей отсутствует запись
            # calculate и dmesg содержит сообщение об ошибке efivars -
            # запись создать не удалось
            efiBootMgr = varsShare().getProgPath('/usr/sbin/efibootmgr')
            dmesg = varsShare().getProgPath('/bin/dmesg')
            if efiBootMgr and dmesg:
                if not re.search('Boot.*calculate',
                        process(efiBootMgr).read(),re.M) and \
                   re.search('efivars.*set_variable.*failed',
                        process(dmesg).read(),re.M):
                    raise DistributiveError(
                        _("Failed to create the UEFI boot record"))
        # не UEFI установка
        else:
            # получить загрузочный раздел (если есть /boot, то 
            # он является загрузочным иначе корень)
            for bootPath in ("/boot","/"):
                bootDisk = self.clVars.Select("os_install_disk_dev",
                    where="os_install_disk_mount",eq=bootPath,limit=1)
                if bootDisk:
                    self.setActivePartition(bootDisk)
                    break

            # если GRUB2 версии 2.00 и выше, обычная установка требует
            # параметра --target=i386-pc, иначе GRUB2 может попытаться
            # прописать себя как UEFI
            if filter(lambda x:"2.00" in x,
                           process(cmdGrubInstall,'-v')):
                platform = ["--target=i386-pc"]
            else:
                platform = []
            # прописать GRUB2 на все указанные диски
            for mbrDisk in self.clVars.Get('os_install_mbr'):
                grubProcess = process(cmdGrubInstall,
                                  "--boot-directory=%s"%pathJoin(prefixBoot,
                                  target.getBootDirectory()),
                                  mbrDisk, "-f", *platform,
                                  stderr=STDOUT,envdict=os.environ)
                if grubProcess.failed():
                    raise DistributiveError(_("Failed to install the bootloader"))

    def installLegacyGrubBootloader(self,target):
        """
        Install legecy grub boot loader

        Perform grub installation to disk, which has root partition
        """
        cmdGrub = varsShare().getProgPath('/sbin/grub')
        if not cmdGrub:
            raise DistributiveError(_("Failed to install the bootloader"))
        grubProcess = process(cmdGrub,
               "--device-map=%s/boot/grub/device.map"%target.getDirectory(),
               "--batch",stderr=STDOUT)
        bootDisk = self.clVars.Select('os_install_disk_grub',
                                      where='os_install_disk_mount',
                                      _in=('/', '/boot'),
                                      sort='DESC', limit=1)
        if not bootDisk:
            raise DistributiveError(_("Failed to determine the boot disk"))
        self.setActivePartition(bootDisk)
        for mbrDisk in self.clVars.Get('os_install_mbr'):
            mbrDiskNum = self.clVars.Select("os_device_map",
                                            where="os_device_dev",
                                            eq=mbrDisk)
            if not mbrDiskNum and mbrDiskNum != 0:
                raise DistributiveError(_("Failed to determine mbr"))
            for line in ("root (hd%s)"%bootDisk,
                         "setup (hd%d)"%mbrDiskNum,
                         "quit"):
                grubProcess.write("%s\n"%line)
            if grubProcess.failed():
                raise DistributiveError(_("Failed to install the bootloader"))
    def setupOpenGL(self):
        """
        Выполнить выбор opengl для текущего видеодрайвера
        """
        defaultGL = "xorg-x11"
        pathGlModules = path.join(self.clVars.Get('cl_chroot_path'),
                                 'usr/lib/opengl')
        openGLenv = path.join(self.clVars.Get('cl_chroot_path'),
                                 'etc/env.d/03opengl')

        openGlMods = filter(lambda x:x != "global",
                               listDirectory(pathGlModules))
        mapGL_drivers = {'fglrx':"ati" if "ati" in openGlMods
                                 else defaultGL,
                         'nvidia':"nvidia" if "nvidia" in openGlMods
                                  else defaultGL}
        x11_driver = self.clVars.Get('os_install_x11_video_drv')
        if x11_driver in mapGL_drivers:
            newModuleName = mapGL_drivers[x11_driver]
        else:
            newModuleName = defaultGL
        curModuleName = map(lambda x:x.strip().rpartition('=')[-1].strip('"\''),
                   filter(lambda x: x.startswith("OPENGL_PROFILE="),
                          readLinesFile(openGLenv)))
        curModuleName = curModuleName[-1] if curModuleName else ""
        if curModuleName == newModuleName:
            return True
        return process('/usr/bin/eselect','opengl','set',newModuleName).success()

    def checkVideoDriver(self):
        """
        Проверить видео драйвер, и если это nvidia, то
        обновить маску на пакет видеодрайвера
        """
        if self.clVars.Get('hr_video') != 'nvidia':
            return True
        maskFile = '/etc/portage/package.mask'
        nvidiaMaskFile = path.join(maskFile,'nvidia')
        # если package.mask является файлом - делаем его директорией
        if path.isfile(maskFile):
            os.rename(maskFile,maskFile+"2")
            os.mkdir(maskFile,mode=0755)
            os.rename(maskFile+"2",path.join(maskFile,"default"))
        curNvidiaMask = readFile(nvidiaMaskFile).strip()
        maskNvidia = self.clVars.Get('os_nvidia_mask')
        if maskNvidia == curNvidiaMask:
            return True
        open(nvidiaMaskFile,'w').write(maskNvidia)
        return True

    def changeScheduler(self,scheduler):
        """
        Изменить текущий IO планировщик
        """
        root_dev = self.clVars.Select('os_disk_parent',
                               where='os_disk_mount',
                               eq='/',limit=1)
        try:
            schedpath = ("/sys%s/queue/scheduler"%
              (getUdevDeviceInfo(name=root_dev).get('DEVPATH','')))
            if path.exists(schedpath):
                open(schedpath,'w').write(scheduler)
        except:
            raise InstallError(("Unable to change the I/O scheduler"))
        return True

    def autopartition(self,table,devices,data,lvm,lvm_vgname,bios_grub,
                      bios_grub_size):
        """
        Авторазметка диска с таблицей разделов 'table', диски указываются
        'device', параметры таблицы 'data', 'lvm' использование LVM,
        'lvm_vgname' название группы томов LVM, bios_grub - создавать
        bios_grub раздел, bios_grub_size - раздел bios grub раздела в байтах
        """
        ap = AutoPartition()
        ap.clearLvm(devices,self.clVars)
        ap.clearRaid(devices,self.clVars)
        ap.recreateSpace(table,devices,data,lvm,
                         lvm_vgname,bios_grub,bios_grub_size)
        return True

    def format(self,target):
        """
        Форматировать разделы для 'target' дистрибутива
        """
        target.performFormat()
        return True

    def unpack(self,source,target,filesnum):
        """
        Распаковать 'source' в 'target', 'filesnum' количество копируемых файлов
        """
        self.addProgress()
        if filesnum.isdigit():
            filesnum = int(filesnum)
        else:
            filesnum = 0
        target.installFrom(source, callbackProgress=self.setProgress,
                           filesnum=filesnum)
        return True

    def copyClt(self,source,target,cltpath):
        """
        Скопировать clt шаблоны из 'cltpath' в 'target' дистрибутив из
        'source' дистрибутива
        """
        targetdir = target.getDirectory()
        sourcedir = source.getDirectory()
        for f in filter(lambda x:x.endswith('.clt'),
                 chain(*map(lambda x:find(pathJoin(sourcedir,x),filetype="f"),
                        cltpath))):
            copyWithPath(f,targetdir,prefix=sourcedir)
        return True

    def copyOther(self,source,target):
        """
        Скопировать прочие настройки из текущей системы в новую
        """
        fileMask = re.compile("(/etc/ssh/ssh_host_.*|"
                               "/root/.ssh/(id_.*|known_hosts))")
        targetdir = target.getDirectory()
        sourcedir = source.getDirectory()
        for f in filter(fileMask.search,
                 chain(*map(lambda x:find(pathJoin(sourcedir,x),filetype="f"),
                        ["/etc","/root/.ssh"]))):
            copyWithPath(f,targetdir,prefix=sourcedir)
        return True


    def rndString(self):
        """
        Получить произвольную строку из 8 символов
        """
        """Get random string with len 8 char"""
        return "".join([choice(string.ascii_letters+string.digits)
                    for i in xrange(0,8)])

    def _getFreeDirectory(self,directory):
        """
        Получить название директории
        """
        newDirectoryName = directory
        while path.exists(newDirectoryName):
            newDirectoryName = "%s.%s"%(directory,self.rndString())
        return newDirectoryName

    def remountNTFS(self):
        """
        Перемонтировать NTFS разделы для работы os-prober
        """
        res = True
        for disk in self.clVars.Select('os_disk_dev',
                                       where='os_disk_format', like='ntfs'):
            mountDir = self._getFreeDirectory('/var/lib/calculate/mount.ntfs')
            try:
                os.mkdir(mountDir)
            except:
                continue
            if process('/bin/mount',disk,mountDir).success():
                for i in (0.2,0.5,1,2,4,5):
                    if process('/bin/umount',mountDir).success():
                        break
                    time.sleep(i)
                else:
                    self.printWARNING(_("Unable to umount %s")%mountDir)
                    res = False
            try:
                os.rmdir(mountDir)
            except:
                self.printWARNING(_("Unable to remove directory %s")%mountDir)
                return False
        return res

    def mountBind(self,target):
        """
        Подключить bind точки монтирования у дистрибутива
        """
        target.postinstallMountBind()
        return True

    def userMigrate(self,target,migrate_data,root_pwd):
        """
        Перенос текущих пользователей в новую систему,
        установка пароля пользователя root
        """
        objMigrate = migrate(target.getDirectory())
        if not objMigrate.migrate(migrate_data,root_pwd,[],[],):
            raise InstallError(_("Failed to migrate users onto the new system"))
        return True

    def umount(self,distr):
        """
        Отключить дистрибутив
        """
        distr.close()
        return True