from typing import AsyncIterator

import trio
import ftplib


class AFTP(ftplib.FTP):
    def __init__(self, host='', user='', passwd='', acct='',
                 timeout=ftplib._GLOBAL_DEFAULT_TIMEOUT, source_address=None):
        super().__init__(acct=acct, timeout=timeout, source_address=source_address)
        self._host = host
        self._user = user
        self._passwd = passwd

    async def connect(self, host='', port=0, timeout=-999, source_address=None):
        if host == '':
            host = self._host
        welcome = await trio.to_thread.run_sync(
            super().connect,
            host,
            port,
            timeout,
            source_address
        )

        if self._user:
            await self.login()
        return welcome

    async def __aenter__(self):
        if self.sock is None:
            await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await trio.to_thread.run_sync(self.__exit__)

    async def login(self, user='', passwd='', acct=''):
        if not user:
            user = self._user
        if not passwd:
            passwd = self._passwd
        await trio.to_thread.run_sync(super().login, user, passwd, acct)

    async def cwd(self, dirname):
        return await trio.to_thread.run_sync(self.cwd, dirname)

    async def iretrlines(self, cmd, transf_type='A') -> AsyncIterator[bytes]:
        """
        Retrieve data in line mode.
        The argument is a RETR or LIST command.
        Yields each line.
        """
        await trio.to_thread.run_sync(self.sendcmd, f'TYPE {transf_type}')
        conn = await trio.to_thread.run_sync(self.transfercmd, cmd)
        fp = conn.makefile('rb')

        def _finalize():
            fp.close()
            conn.close()

        while True:
            try:
                line = await trio.to_thread.run_sync(fp.readline)
            except Exception:
                await trio.to_thread.run_sync(_finalize)
                raise
            if not line:
                await trio.to_thread.run_sync(_finalize)
                await trio.to_thread.run_sync(self.voidresp)  # check the response code
                break
            try:
                yield line
            except GeneratorExit:
                await trio.to_thread.run_sync(_finalize)
                raise