import enum
import asyncio
import traceback

from atds.protocol.packets import TDSPacket, TDSPacketType, TDSStatus
from atds.protocol.packets.prelogin import TDS_PRELOGIN, PreLoginOption, EncryptMode, PreLoginOptionToken
from atds.protocol.packets.login import TDS_LOGIN, TDS_LOGIN_FLAGS2
from atds.network.packetizer import TDSPacketizer


class TDSConnectionStatus(enum.Enum):
    PRELOGIN = 'PRELOGIN'
    LOGIN = 'LOGIN'
    SSPI = 'SSPI'
    RUNNING = 'RUNNING'
    CLOSED = 'CLOSED'


class TDSRelayServerConnection:
    def __init__(self, settings, connection, connection_id):
        self.settings = settings
        self.gssapi = self.settings.create_new_gssapi(connection)
        self.connection_id = connection_id
        self.connection = connection
        self.incoming_task = None
        self.status = TDSConnectionStatus.PRELOGIN

    async def run(self):
        self.incoming_task = asyncio.create_task(self.__handle_tds_in())
        return self.incoming_task

    async def stop(self):
        try:
            await self.log('[INF] Stopping connection')
            if self.incoming_task is not None:
                self.incoming_task.cancel()
            if self.connection is not None:
                await self.connection.close()
        except Exception as e:
            traceback.print_exc()
            await self.log('[ERR] %s' % e)

    async def log(self, msg):
        if self.settings.log_callback is not None:
            await self.settings.log_callback(f'[TDSRELAYCONN][{self.connection_id}] {msg}')

    async def __handle_tds_in(self):
        try:
            async for packet in self.connection.read():
                if packet is None:
                    break

                await self.log(f'[PACKET] Type={packet.Type} Status={packet.Status}')

                if self.status == TDSConnectionStatus.PRELOGIN:
                    if packet.Type == TDSPacketType.PRELOGIN:
                        await self.handle_prelogin(packet)
                    else:
                        await self.log(f'[ERR] Unexpected packet type in PRELOGIN state: {packet.Type}')
                        return

                elif self.status == TDSConnectionStatus.LOGIN:
                    if packet.Type == TDSPacketType.TDS7_LOGIN:
                        await self.handle_login(packet)
                    else:
                        await self.log(f'[ERR] Unexpected packet type in LOGIN state: {packet.Type}')
                        return

                elif self.status == TDSConnectionStatus.SSPI:
                    # Client sends follow-up SSPI data
                    # Could be packet type SSPI (17) per spec, or TABULAR_RESULT (4)
                    await self.handle_sspi(packet)

                elif self.status == TDSConnectionStatus.RUNNING:
                    # Authentication completed, relay is done
                    return

        except asyncio.CancelledError:
            return
        except:
            traceback.print_exc()
        finally:
            await self.stop()

    async def handle_prelogin(self, packet: TDSPacket):
        try:
            client_prelogin = TDS_PRELOGIN.from_bytes(packet.Data)
            await self.log('[INF] Got PRELOGIN from client')

            # Build server prelogin response with encryption not supported
            # This avoids TLS negotiation on the relay server
            response = TDS_PRELOGIN()
            response.options = [
                PreLoginOption.create_version(
                    major=self.settings.sql_version_major,
                    minor=self.settings.sql_version_minor,
                    build=self.settings.sql_version_build
                ),
                PreLoginOption.create_encryption(EncryptMode.ENCRYPT_NOT_SUP),
                PreLoginOption.create_instance(''),
                PreLoginOption.create_threadid(0),
            ]

            reply = TDSPacket()
            reply.Type = TDSPacketType.TABULAR_RESULT
            reply.Status = TDSStatus.NORMAL | TDSStatus.EOM
            reply.Data = response.to_bytes()
            reply.SPID = 0
            reply.PacketID = 1

            await self.send_packet(reply)
            self.status = TDSConnectionStatus.LOGIN

        except Exception as e:
            traceback.print_exc()
            await self.log(f'[ERR] PRELOGIN handling failed: {e}')

    async def handle_login(self, packet: TDSPacket):
        try:
            data = packet.Data
            await self.log(f'[INF] Got LOGIN packet ({len(data)} bytes)')

            # Parse flags2 to check for integrated security
            # flags2 is at byte 25 in the LOGIN7 structure
            if len(data) < 86:
                await self.log('[ERR] LOGIN packet too short')
                await self.send_error_response(18456, 'Login failed.')
                return

            flags2 = data[25]
            is_integrated = bool(flags2 & TDS_LOGIN_FLAGS2.INTEGRATED_SECURITY)

            if is_integrated:
                # Extract SSPI data using spec-correct fixed offsets
                # ibSSPI at bytes 78-79, cbSSPI at bytes 80-81
                sspi_offset = int.from_bytes(data[78:80], byteorder='little')
                sspi_length = int.from_bytes(data[80:82], byteorder='little')

                if sspi_length > 0 and sspi_offset + sspi_length <= len(data):
                    sspi_data = data[sspi_offset:sspi_offset + sspi_length]
                    await self.log(f'[INF] Extracted SSPI data ({len(sspi_data)} bytes)')
                    await self.process_sspi(sspi_data)
                else:
                    await self.log('[ERR] Integrated security flag set but no valid SSPI data')
                    await self.send_error_response(18456, 'Login failed.')
            else:
                # Plain text login - not relayable
                await self.log('[INF] Plain text login received (not relayable)')
                await self.send_error_response(18456, 'Login failed. Only integrated authentication is supported for relay.')

        except Exception as e:
            traceback.print_exc()
            await self.log(f'[ERR] LOGIN handling failed: {e}')

    async def handle_sspi(self, packet: TDSPacket):
        try:
            sspi_data = packet.Data
            await self.log(f'[INF] Got follow-up SSPI data ({len(sspi_data)} bytes)')
            await self.process_sspi(sspi_data)
        except Exception as e:
            traceback.print_exc()
            await self.log(f'[ERR] SSPI handling failed: {e}')

    async def process_sspi(self, sspi_data: bytes):
        try:
            await self.log('[INF] Calling authenticate_relay_server')
            response_data, to_continue, err = await self.gssapi.authenticate_relay_server(sspi_data)
            await self.log(f'[INF] Auth result: data={response_data is not None}, to_continue={to_continue}, err={err}')

            if err is not None:
                raise err

            if to_continue is False and response_data is None:
                # Auth completed on the relay side, get final response for client
                response_data, err = await self.gssapi.authenticate_relay_server_finished()
                if err is not None:
                    raise err
                await self.send_login_success()
                self.status = TDSConnectionStatus.RUNNING
                return

            if to_continue is True:
                # Send SSPI challenge back to client
                await self.send_sspi_response(response_data)
                self.status = TDSConnectionStatus.SSPI
                return

            if to_continue is False:
                # Auth completed successfully
                await self.send_login_success()
                self.status = TDSConnectionStatus.RUNNING
                return

        except Exception as e:
            await self.log(f'[ERR] SSPI processing failed: {e}')
            traceback.print_exc()
            await self.send_error_response(18456, 'Login failed.')

    async def send_sspi_response(self, sspi_data: bytes):
        """Send SSPI challenge token wrapped in a TABULAR_RESULT packet"""
        token_data = self.build_sspi_token(sspi_data)

        reply = TDSPacket()
        reply.Type = TDSPacketType.TABULAR_RESULT
        reply.Status = TDSStatus.NORMAL | TDSStatus.EOM
        reply.Data = token_data
        reply.SPID = 0
        reply.PacketID = 1

        await self.send_packet(reply)

    async def send_login_success(self):
        """Send LOGINACK + ENVCHANGE + DONE tokens"""
        token_data = bytearray()
        token_data.extend(self.build_loginack_token())
        token_data.extend(self.build_envchange_database('master', ''))
        token_data.extend(self.build_envchange_packetsize(str(self.settings.packet_size), str(self.settings.packet_size)))
        token_data.extend(self.build_done_token())

        reply = TDSPacket()
        reply.Type = TDSPacketType.TABULAR_RESULT
        reply.Status = TDSStatus.NORMAL | TDSStatus.EOM
        reply.Data = bytes(token_data)
        reply.SPID = 0
        reply.PacketID = 1

        await self.send_packet(reply)

    async def send_error_response(self, error_number: int, message: str):
        """Send ERROR + DONE tokens"""
        token_data = bytearray()
        token_data.extend(self.build_error_token(error_number, message))
        token_data.extend(self.build_done_token(status=0x0002))  # DONE_ERROR

        reply = TDSPacket()
        reply.Type = TDSPacketType.TABULAR_RESULT
        reply.Status = TDSStatus.NORMAL | TDSStatus.EOM
        reply.Data = bytes(token_data)
        reply.SPID = 0
        reply.PacketID = 1

        await self.send_packet(reply)

    async def send_packet(self, packet: TDSPacket):
        await self.connection.write(packet.to_bytes())

    # ---- Token builders for server responses ----

    @staticmethod
    def build_sspi_token(sspi_data: bytes) -> bytes:
        """Build SSPI token (0xED) for sending challenge to client"""
        result = bytearray()
        result.append(0xED)  # SSPI token type
        result.extend(len(sspi_data).to_bytes(2, byteorder='little'))
        result.extend(sspi_data)
        return bytes(result)

    @staticmethod
    def build_loginack_token(tds_version: int = 0x71000001, prog_name: str = "Microsoft SQL Server") -> bytes:
        """Build LOGINACK token (0xAD) for successful login response"""
        result = bytearray()
        result.append(0xAD)  # LOGINACK token type

        body = bytearray()
        body.append(0x01)  # Interface: SQL_TSQL
        body.extend(tds_version.to_bytes(4, byteorder='little'))
        prog_name_encoded = prog_name.encode('utf-16-le')
        body.append(len(prog_name))  # prog_name length in chars
        body.extend(prog_name_encoded)
        body.extend(bytes([15, 0, 41, 23]))  # Version: 15.0.41.23

        result.extend(len(body).to_bytes(2, byteorder='little'))
        result.extend(body)
        return bytes(result)

    @staticmethod
    def build_done_token(status: int = 0, cur_cmd: int = 0, row_count: int = 0) -> bytes:
        """Build DONE token (0xFD)"""
        result = bytearray()
        result.append(0xFD)  # DONE token type
        result.extend(status.to_bytes(2, byteorder='little'))
        result.extend(cur_cmd.to_bytes(2, byteorder='little'))
        # TDS 7.1: row count is 4 bytes (LONG)
        result.extend(row_count.to_bytes(4, byteorder='little'))
        return bytes(result)

    @staticmethod
    def build_error_token(number: int = 18456, message: str = "Login failed.",
                          state: int = 1, severity: int = 14, server_name: str = "",
                          proc_name: str = "", line_number: int = 1) -> bytes:
        """Build ERROR token (0xAA) for login failure response"""
        result = bytearray()
        result.append(0xAA)  # ERROR token type

        body = bytearray()
        body.extend(number.to_bytes(4, byteorder='little'))
        body.append(state)
        body.append(severity)
        msg_encoded = message.encode('utf-16-le')
        body.extend(len(message).to_bytes(2, byteorder='little'))
        body.extend(msg_encoded)
        srv_encoded = server_name.encode('utf-16-le')
        body.append(len(server_name))
        body.extend(srv_encoded)
        proc_encoded = proc_name.encode('utf-16-le')
        body.append(len(proc_name))
        body.extend(proc_encoded)
        # TDS 7.1: line number is 2 bytes (USHORT)
        body.extend(line_number.to_bytes(2, byteorder='little'))

        result.extend(len(body).to_bytes(2, byteorder='little'))
        result.extend(body)
        return bytes(result)

    @staticmethod
    def build_envchange_database(new_db: str = "master", old_db: str = "") -> bytes:
        """Build ENVCHANGE token (0xE3) for database change"""
        result = bytearray()
        result.append(0xE3)  # ENVCHANGE token type

        body = bytearray()
        body.append(0x01)  # Type: DATABASE
        new_encoded = new_db.encode('utf-16-le')
        body.append(len(new_db))
        body.extend(new_encoded)
        old_encoded = old_db.encode('utf-16-le')
        body.append(len(old_db))
        body.extend(old_encoded)

        result.extend(len(body).to_bytes(2, byteorder='little'))
        result.extend(body)
        return bytes(result)

    @staticmethod
    def build_envchange_packetsize(new_size: str = "4096", old_size: str = "4096") -> bytes:
        """Build ENVCHANGE token (0xE3) for packet size change"""
        result = bytearray()
        result.append(0xE3)  # ENVCHANGE token type

        body = bytearray()
        body.append(0x04)  # Type: PACKET_SIZE
        new_encoded = new_size.encode('utf-16-le')
        body.append(len(new_size))
        body.extend(new_encoded)
        old_encoded = old_size.encode('utf-16-le')
        body.append(len(old_size))
        body.extend(old_encoded)

        result.extend(len(body).to_bytes(2, byteorder='little'))
        result.extend(body)
        return bytes(result)

