# Credit to Michael Kaplan <http://blogs.msdn.com/b/michkap/archive/2010/04/07/9989346.aspx>
# and TZOmegaTZIOY
# <http://stackoverflow.com/questions/878972/windows-cmd-encoding-change-causes-python-crash/1432462#1432462>.

import codecs
import ctypes
import locale
import os
import sys
from ctypes import WINFUNCTYPE, windll, POINTER, byref, c_int
from ctypes.wintypes import BOOL, HANDLE, DWORD, LPCSTR, LPWSTR, LPCWSTR, LPVOID

KERNEL32 = windll.kernel32

INVALID_HANDLE_VALUE = DWORD(-1).value
STD_OUTPUT_HANDLE = DWORD(-11)
STD_ERROR_HANDLE  = DWORD(-12)
FILE_TYPE_CHAR   = 0x0002
FILE_TYPE_REMOTE = 0x8000
STDOUT_FILENO = 1
STDERR_FILENO = 2

GetStdHandle = WINFUNCTYPE(HANDLE, DWORD)(("GetStdHandle", KERNEL32))
GetFileType = WINFUNCTYPE(DWORD, DWORD)(("GetFileType", KERNEL32))
GetConsoleMode = WINFUNCTYPE(BOOL, HANDLE, POINTER(DWORD)) \
                     (("GetConsoleMode", KERNEL32))
WriteConsoleA = WINFUNCTYPE(BOOL, HANDLE, LPCSTR, DWORD, POINTER(DWORD), \
                            LPVOID)(("WriteConsoleA", KERNEL32))
WriteConsoleW = WINFUNCTYPE(BOOL, HANDLE, LPWSTR, DWORD, POINTER(DWORD), \
                            LPVOID)(("WriteConsoleW", KERNEL32))

class Base:
    def __init__(self, fileno):
        self._fileno = fileno
        self.closed = False
        self.softspace = False
        self.mode = 'w'

    def close(self):
        # don't really close the handle, that would only cause problems
        self.closed = True

    def fileno(self):
        return self._fileno

    # @abstractmethod
    def write(self, data):
        pass

    def writelines(self, lines):
        for line in lines:
            self.write(line)

class Console(Base):
    def __init__(self, hConsole, fileno):
        Base.__init__(self, fileno)
        self._hConsole = hConsole

    def isatty(self):
        return True

    def flush(self):
        pass

class Stream(Base):
    def __init__(self, stream, fileno):
        Base.__init__(self, fileno)
        self._stream = stream

    def isatty(self):
        return False

    def flush(self):
        self._stream.flush()

class BytesStream(Stream):
    def write(self, raw):
        self._stream.write(raw)

class BytesConsole(Console):
    def write(self, raw):
        while raw:
            # Issue #11395: the Windows console returns an error (12: not
            # enough space error) on writing into stdout if stdout mode is
            # binary and the length is greater than 66,000 bytes (or less,
            # depending on heap usage).
            chunk_len = min(len(raw), 32767)
            chunk = raw[:chunk_len]

            retval = WriteConsoleA(self._hConsole,
                                   chunk, len(chunk),
                                   None, None)
            if retval == 0:
                raise IOError("WriteConsoleA returned %r" % retval)
            raw = raw[len(chunk):]

class BaseUnicode:
    def __init__(self, buffer):
        self.buffer = buffer
        self.newline = '\r\n'

class UnicodeStream(BaseUnicode, Stream):
    def __init__(self, buffer, stream, fileno, encoding):
        BaseUnicode.__init__(self, buffer)
        Stream.__init__(self, stream, fileno)
        self.encoding = encoding

    def write(self, text):
        text = text.replace('\n', self.newline)
        text = text.encode(self.encoding)
        self._stream.write(text)

class UnicodeConsole(BaseUnicode, Console):
    def __init__(self, buffer, hConsole, fileno):
        BaseUnicode.__init__(self, buffer)
        Console.__init__(self, hConsole, fileno)

    def write(self, text):
        text = text.replace('\n', self.newline)
        while text:
            # Issue #11395: the Windows console returns an error (12: not
            # enough space error) on writing into stdout if stdout mode is
            # binary and the length is greater than 66,000 bytes (or less,
            # depending on heap usage).
            chunk_len = min(len(text), 32767)
            chunk = text[:chunk_len]

            retval = WriteConsoleW(self._hConsole,
                                   chunk, len(chunk),
                                   None, None)
            if retval == 0:
                raise IOError("WriteConsoleW returned %r" % retval)
            text = text[len(chunk):]

def is_a_console(handle):
    if handle == INVALID_HANDLE_VALUE or handle is None:
        return False
    if (GetFileType(handle) & ~FILE_TYPE_REMOTE) != FILE_TYPE_CHAR:
        return False
    if GetConsoleMode(handle, byref(DWORD())) == 0:
        return False
    return True

def install_unicode_console():
    # Make Unicode console output work independently of the current code page.
    if hasattr(sys.stdout, 'fileno'):
        old_stdout_fileno = sys.stdout.fileno()
    else:
        old_stdout_fileno = None
    if hasattr(sys.stderr, 'fileno'):
        old_stderr_fileno = sys.stderr.fileno()
    else:
        old_stderr_fileno = None

    real_stdout = (old_stdout_fileno == STDOUT_FILENO)
    real_stderr = (old_stderr_fileno == STDERR_FILENO)

    if real_stdout:
        hStdout = GetStdHandle(STD_OUTPUT_HANDLE)
        if not is_a_console(hStdout):
            real_stdout = False

    if real_stderr:
        hStderr = GetStdHandle(STD_ERROR_HANDLE)
        if not is_a_console(hStderr):
            real_stderr = False

    if not(real_stdout or real_stderr):
        print("Nothing to do")
        return

    cp = KERNEL32.GetConsoleOutputCP()
    encoding = "cp%s" % cp

    sys.stdout.flush()
    if real_stdout:
        buffer = BytesConsole(hStdout, STDOUT_FILENO)
        textio = UnicodeConsole(buffer, hStdout, STDOUT_FILENO)
    else:
        raw = sys.stdout.buffer.raw
        buffer = BytesStream(raw, old_stdout_fileno)
        textio = UnicodeStream(buffer, raw, old_stdout_fileno, encoding)
    sys.stdout = textio

    sys.stderr.flush()
    if real_stderr:
        buffer = BytesConsole(hStderr, STDERR_FILENO)
        textio = UnicodeConsole(buffer, hStderr, STDERR_FILENO)
    else:
        raw = sys.stderr.buffer.raw
        buffer = BytesStream(raw, old_stderr_fileno)
        textio = UnicodeStream(bufer, raw, old_stderr_fileno, encoding)
    sys.stderr = textio

def tests():
    print("ANSI CP: %s, OEM CP: %s"
          % (KERNEL32.GetACP(), KERNEL32.GetOEMCP()))
    print("GetConsoleCP: %s, GetConsoleOutputCP: %s"
          % (KERNEL32.GetConsoleCP(), KERNEL32.GetConsoleOutputCP()))
    print("device_encoding(0=stdin):  %s, tty? %s"
          % (os.device_encoding(0), os.isatty(0)))
    print("device_encoding(1=stdout): %s, tty? %s"
          % (os.device_encoding(1), os.isatty(1)))
    print("device_encoding(2=stderr): %s, tty? %s"
          % (os.device_encoding(2), os.isatty(2)))
    print("sys encoding: stdin=%s, stdout=%s, stderr=%s"
          % (sys.stdin.encoding, sys.stdout.encoding, sys.stderr.encoding))
    print("locale.getpreferredencoding(): %s"
          % (locale.getpreferredencoding()))

    cp = KERNEL32.GetConsoleOutputCP()
    encoding = "cp%s" % cp
    print("Use encoding %s" % encoding)

    chars = ('A', '\xE9', '\u0141', '\u20ac')

    install_unicode_console()

    print(ascii(sys.stdout))
    print(ascii(sys.stdout.buffer))
    print(ascii(sys.stderr))
    print(ascii(sys.stderr.buffer))
    print()

    for ch in chars:
        try:
            print('Write U+%04x as unicode: %s' % (ord(ch), ch))
        except UnicodeEncodeError:
            print('Write U+%04x as unicode: <unicode encode error>' % ord(ch))
        except IOError:
            print('Write U+%04x as unicode: <WriteConsoleW error>' % ord(ch))
        sys.stdout.flush()

    for ch in chars:
        try:
            data = ('Write U+%04x as bytes: %s' % (ord(ch), ch)).encode(encoding)
        except UnicodeEncodeError:
            data = ('Write U+%04x as bytes: <unicode encode error>' % ord(ch)).encode(encoding)
        except IOError:
            print('Write U+%04x as unicode: <WriteConsoleA error>' % ord(ch))
        sys.stdout.buffer.write(data + b'\n')
        sys.stdout.buffer.flush()
        sys.stdout.flush()

def main():
    KERNEL32 = ctypes.WinDLL('KERNEL32')
    old_console_cp = KERNEL32.GetConsoleCP()
    old_console_output_cp = KERNEL32.GetConsoleOutputCP()

    try:
        console_cp = console_output_cp = None
        #console_cp = 65001
        #console_output_cp = 65001

        if console_cp is not None:
            x = KERNEL32.SetConsoleCP(console_cp)
            if x!= 1:
                print("SetConsoleCP returned ", x )
                exit( 1 )
        if console_output_cp is not None:
            x = KERNEL32.SetConsoleOutputCP(console_output_cp)
            if x != 1:
                print("SetConsoleOutputCP returned ", x )
                exit( 1 )

        tests()
    finally:
        KERNEL32.SetConsoleCP(old_console_cp )
        KERNEL32.SetConsoleOutputCP(old_console_output_cp )

if __name__ == "__main__":
    main()