import 'dart:async';
import 'dart:convert';
import 'dart:io';
import 'dart:isolate';
import 'dart:typed_data';

import 'package:crypto/crypto.dart' as hashes;
import 'package:web_socket_channel/io.dart';

import 'edge_client.dart';
import 'identity_store.dart';
import 'models.dart';

const protocolMajor = 1;
const protocolMinor = 16;
const _maxDimension = 8192;
const _maxPixels = 33177600;
const _maxCompressedBytes = 64 * 1024 * 1024;

enum AgentPhase {
  connecting,
  switchingToRelay,
  authenticating,
  openingDesktop,
  waitingForFrame,
  connected,
  disconnected,
  failed,
}

enum AgentError {
  network,
  certificate,
  protocol,
  authentication,
  totp,
  permission,
  desktop,
  frame,
  unknown,
}

class AgentState {
  const AgentState(this.phase, [this.error]);

  final AgentPhase phase;
  final AgentError? error;
}

class RemoteFrame {
  const RemoteFrame({
    required this.sequence,
    required this.width,
    required this.height,
    required this.bgra,
    this.textureId,
  });

  final int sequence;
  final int width;
  final int height;
  final Uint8List bgra;
  final int? textureId;
}

class AgentException implements Exception {
  const AgentException(this.error);

  final AgentError error;
}

abstract interface class RemoteDesktopClient {
  Stream<AgentState> get states;
  Stream<RemoteFrame> get frames;

  Future<void> connect({String? pairingCode, required String totpCode});
  void sendPointer(int x, int y);
  void sendButton(String button, bool pressed);
  void sendKey(int keysym);
  void sendText(String value);
  void acknowledgeFrame(int sequence);
  Future<void> close();
  Future<void> dispose();
}

class AgentClient implements RemoteDesktopClient {
  AgentClient({
    required this.host,
    required this.settings,
    IdentityStore? identityStore,
  }) : _identityStore = identityStore ?? IdentityStore();

  final RemoteHost host;
  final AppSettings settings;
  final IdentityStore _identityStore;
  final StreamController<AgentState> _states =
      StreamController<AgentState>.broadcast(sync: true);
  final StreamController<RemoteFrame> _frames =
      StreamController<RemoteFrame>.broadcast(sync: true);

  IOWebSocketChannel? _channel;
  StreamIterator<dynamic>? _messages;
  HttpClient? _httpClient;
  _FrameAssembly? _assembly;
  bool _closed = false;
  bool _certificateRejected = false;
  bool _usingRelay = false;
  bool _desktopOpen = false;
  RelayTunnel? _relayTunnel;

  @override
  Stream<AgentState> get states => _states.stream;
  @override
  Stream<RemoteFrame> get frames => _frames.stream;

  @override
  Future<void> connect({String? pairingCode, required String totpCode}) async {
    _closed = false;
    _emit(AgentPhase.connecting);
    try {
      final expectedFingerprint = normalizeFingerprint(host.certificateSha256);
      if (!isValidFingerprint(expectedFingerprint)) {
        throw const AgentException(AgentError.certificate);
      }

      final identity = await _identityStore.loadOrCreate();
      late Map<String, dynamic> hello;
      if (host.connectionMode == ConnectionMode.relay) {
        final relayUri = await _openRelay(identity, pairingCode != null);
        hello = await _openTransport(
          relayUri,
          expectedFingerprint,
          usingRelay: true,
        );
      } else {
        try {
          hello = await _openTransport(
            host.uri,
            expectedFingerprint,
            usingRelay: false,
          );
        } on AgentException catch (error) {
          if (host.connectionMode != ConnectionMode.automatic ||
              error.error != AgentError.network) {
            rethrow;
          }
          _emit(AgentPhase.switchingToRelay);
          final relayUri = await _openRelay(identity, pairingCode != null);
          hello = await _openTransport(
            relayUri,
            expectedFingerprint,
            usingRelay: true,
          );
        }
      }
      _emit(AgentPhase.authenticating);

      final challenge = _decodeUnpaddedBase64(_string(hello, 'challenge'));
      if (challenge.length != 32) {
        throw const AgentException(AgentError.protocol);
      }
      final payload = <int>[
        ...utf8.encode('remotedesk-agent-auth-v1\u0000'),
        ...challenge,
      ];
      final signature = await _identityStore.sign(identity, payload);
      final publicKey = _encodeUnpaddedBase64(identity.publicKey.bytes);
      final command = pairingCode == null
          ? <String, Object>{
              'type': 'authenticate',
              'protocol_major': protocolMajor,
              'protocol_minor': protocolMinor,
              'client_public_key': publicKey,
              'signature': _encodeUnpaddedBase64(signature),
            }
          : <String, Object>{
              'type': 'pair',
              'protocol_major': protocolMajor,
              'protocol_minor': protocolMinor,
              'client_name': 'RemoteDesk Android',
              'client_public_key': publicKey,
              'signature': _encodeUnpaddedBase64(signature),
              'code': pairingCode,
            };
      _send(command);

      final totpRequired = await _nextMessage(const Duration(seconds: 8));
      if (totpRequired['type'] == 'error') {
        throw AgentException(_serverError(_string(totpRequired, 'code')));
      }
      if (totpRequired['type'] != 'totp_required' ||
          _integer(totpRequired, 'digits') != 6 ||
          _integer(totpRequired, 'period_seconds') != 30) {
        throw const AgentException(AgentError.protocol);
      }
      _send({'type': 'verify_totp', 'code': totpCode});

      final authenticated = await _nextMessage(const Duration(seconds: 8));
      _validateAuthentication(authenticated);
      _emit(AgentPhase.openingDesktop);
      _send({
        'type': 'open_desktop',
        'user': host.user,
        'max_width': settings.maxWidth,
        'max_height': settings.maxHeight,
        'frames_per_second': settings.framesPerSecond,
        'resume_token': null,
        'client_fingerprint': null,
        'edge_session_id': null,
        'webrtc_h264': false,
        'opus_audio': false,
        'clipboard_read': false,
        'clipboard_write': false,
      });

      final opened = await _nextMessage(const Duration(seconds: 15));
      if (opened['type'] == 'error') {
        throw const AgentException(AgentError.desktop);
      }
      if (opened['type'] != 'desktop_opened' ||
          opened['encoding'] != 'zlib_bgra') {
        throw const AgentException(AgentError.protocol);
      }
      _desktopOpen = true;
      _emit(AgentPhase.waitingForFrame);
      unawaited(_readLoop());
    } on AgentException catch (error) {
      await _fail(error.error);
      rethrow;
    } on HandshakeException {
      final error = _certificateRejected
          ? AgentError.certificate
          : AgentError.network;
      await _fail(error);
      throw AgentException(error);
    } on Object {
      final error = _certificateRejected
          ? AgentError.certificate
          : AgentError.network;
      await _fail(error);
      throw AgentException(error);
    }
  }

  Future<Uri> _openRelay(ClientIdentity identity, bool pairing) async {
    final tunnel = await EdgeClient(
      identityStore: _identityStore,
    ).authorize(host: host, identity: identity, pairing: pairing);
    _relayTunnel = tunnel;
    return Uri(scheme: 'wss', host: '127.0.0.1', port: tunnel.port, path: '/');
  }

  Future<Map<String, dynamic>> _openTransport(
    Uri uri,
    String expectedFingerprint, {
    required bool usingRelay,
  }) async {
    _certificateRejected = false;
    _usingRelay = usingRelay;
    final httpClient = HttpClient(
      context: SecurityContext(withTrustedRoots: false),
    );
    _httpClient = httpClient;
    httpClient.connectionTimeout = const Duration(seconds: 10);
    httpClient.badCertificateCallback = (certificate, _, _) {
      final actual = hashes.sha256.convert(certificate.der).toString();
      final accepted = actual == expectedFingerprint;
      _certificateRejected = !accepted;
      return accepted;
    };
    try {
      final channel = IOWebSocketChannel.connect(
        uri,
        customClient: httpClient,
        connectTimeout: const Duration(seconds: 12),
        pingInterval: const Duration(seconds: 20),
      );
      _channel = channel;
      await channel.ready.timeout(const Duration(seconds: 12));
      _messages = StreamIterator(channel.stream);
      final hello = await _nextMessage(const Duration(seconds: 8));
      _validateHello(hello, expectedFingerprint);
      return hello;
    } on AgentException {
      await _closeTransport();
      rethrow;
    } on HandshakeException {
      final error = _certificateRejected
          ? AgentError.certificate
          : AgentError.network;
      await _closeTransport();
      throw AgentException(error);
    } on Object {
      final error = _certificateRejected
          ? AgentError.certificate
          : AgentError.network;
      await _closeTransport();
      throw AgentException(error);
    }
  }

  Future<void> _closeTransport() async {
    final channel = _channel;
    _channel = null;
    _messages = null;
    try {
      await channel?.sink.close();
    } on Object {
      // The transport may already be closed.
    }
    _httpClient?.close(force: true);
    _httpClient = null;
  }

  @override
  void sendPointer(int x, int y) {
    if (!_desktopOpen) return;
    _send({
      'type': 'desktop_input',
      'event': {
        'kind': 'pointer_move',
        'x': x.clamp(0, 65535),
        'y': y.clamp(0, 65535),
      },
    });
  }

  @override
  void sendButton(String button, bool pressed) {
    if (!_desktopOpen) return;
    _send({
      'type': 'desktop_input',
      'event': {'kind': 'pointer_button', 'button': button, 'pressed': pressed},
    });
  }

  @override
  void sendKey(int keysym) {
    if (!_desktopOpen) return;
    for (final state in const ['pressed', 'released']) {
      _send({
        'type': 'desktop_input',
        'event': {'kind': 'key', 'keysym': keysym, 'state': state},
      });
    }
  }

  @override
  void sendText(String value) {
    for (final rune in value.runes) {
      sendKey(rune <= 0xff ? rune : 0x01000000 | rune);
    }
  }

  @override
  void acknowledgeFrame(int sequence) {
    if (!_desktopOpen) return;
    _send({'type': 'desktop_frame_ack', 'sequence': sequence});
  }

  @override
  Future<void> close() async {
    if (_closed) return;
    _closed = true;
    _desktopOpen = false;
    final channel = _channel;
    if (channel != null) {
      try {
        _send({'type': 'close'});
      } on Object {
        // The connection may already be closed.
      }
      try {
        await channel.sink.close();
      } on Object {
        // Nothing else is required during shutdown.
      }
    }
    _httpClient?.close(force: true);
    final relayTunnel = _relayTunnel;
    _relayTunnel = null;
    if (relayTunnel != null) await EdgeClient.stopTunnel(relayTunnel.id);
    _emit(AgentPhase.disconnected);
  }

  @override
  Future<void> dispose() async {
    await close();
    await _states.close();
    await _frames.close();
  }

  Future<void> _readLoop() async {
    try {
      while (!_closed) {
        final event = await _nextMessage(const Duration(seconds: 45));
        switch (event['type']) {
          case 'desktop_frame_start':
            _startFrame(event);
          case 'desktop_frame_chunk':
            _addFrameChunk(event);
          case 'desktop_frame_complete':
            await _completeFrame(event);
          case 'desktop_closed':
            await _fail(AgentError.desktop);
            return;
          case 'error':
            await _fail(_serverError(_string(event, 'code')));
            return;
          case 'desktop_pong':
            break;
          default:
            break;
        }
      }
    } on AgentException catch (error) {
      await _fail(error.error);
    } on Object {
      if (!_closed) await _fail(AgentError.network);
    }
  }

  void _startFrame(Map<String, dynamic> event) {
    if (_assembly != null) {
      throw const AgentException(AgentError.frame);
    }
    final metadata = event['metadata'];
    if (metadata is! Map<String, dynamic>) {
      throw const AgentException(AgentError.frame);
    }
    final sequence = _integer(metadata, 'sequence');
    final width = _integer(metadata, 'width');
    final height = _integer(metadata, 'height');
    final uncompressed = _integer(metadata, 'uncompressed_bytes');
    final compressed = _integer(metadata, 'compressed_bytes');
    final chunks = _integer(metadata, 'chunk_count');
    final pixels = width * height;
    if (metadata['encoding'] != 'zlib_bgra' ||
        width <= 0 ||
        height <= 0 ||
        width > _maxDimension ||
        height > _maxDimension ||
        pixels > _maxPixels ||
        uncompressed != pixels * 4 ||
        compressed <= 0 ||
        compressed > _maxCompressedBytes ||
        chunks <= 0 ||
        chunks > 65535) {
      throw const AgentException(AgentError.frame);
    }
    _assembly = _FrameAssembly(
      sequence: sequence,
      width: width,
      height: height,
      uncompressedBytes: uncompressed,
      compressedBytes: compressed,
      chunkCount: chunks,
    );
  }

  void _addFrameChunk(Map<String, dynamic> event) {
    final assembly = _assembly;
    if (assembly == null ||
        _integer(event, 'sequence') != assembly.sequence ||
        _integer(event, 'index') != assembly.nextChunk) {
      throw const AgentException(AgentError.frame);
    }
    final bytes = _decodeUnpaddedBase64(_string(event, 'data'));
    if (bytes.isEmpty ||
        assembly.length + bytes.length > assembly.compressedBytes) {
      throw const AgentException(AgentError.frame);
    }
    assembly.add(bytes);
  }

  Future<void> _completeFrame(Map<String, dynamic> event) async {
    final assembly = _assembly;
    _assembly = null;
    if (assembly == null ||
        _integer(event, 'sequence') != assembly.sequence ||
        assembly.nextChunk != assembly.chunkCount ||
        assembly.length != assembly.compressedBytes) {
      throw const AgentException(AgentError.frame);
    }
    final compressed = assembly.takeBytes();
    final actualHash = hashes.sha256.convert(compressed).toString();
    if (actualHash != _string(event, 'sha256').toLowerCase()) {
      throw const AgentException(AgentError.frame);
    }
    final decompressed = await Isolate.run(
      () => Uint8List.fromList(ZLibCodec().decode(compressed)),
    );
    if (decompressed.length != assembly.uncompressedBytes) {
      throw const AgentException(AgentError.frame);
    }
    if (_closed) return;
    _emit(AgentPhase.connected);
    _frames.add(
      RemoteFrame(
        sequence: assembly.sequence,
        width: assembly.width,
        height: assembly.height,
        bgra: decompressed,
      ),
    );
  }

  Future<Map<String, dynamic>> _nextMessage(Duration timeout) async {
    final messages = _messages;
    if (messages == null || !await messages.moveNext().timeout(timeout)) {
      throw const AgentException(AgentError.network);
    }
    final current = messages.current;
    if (current is! String) {
      throw const AgentException(AgentError.protocol);
    }
    final decoded = jsonDecode(current);
    if (decoded is! Map<String, dynamic>) {
      throw const AgentException(AgentError.protocol);
    }
    return decoded;
  }

  void _validateHello(Map<String, dynamic> hello, String fingerprint) {
    final relayIdentityMatches =
        !_usingRelay ||
        _string(hello, 'device_public_key') == host.agentPublicKey;
    if (hello['type'] != 'hello' ||
        _integer(hello, 'protocol_major') != protocolMajor ||
        _integer(hello, 'protocol_minor') < 16 ||
        hello['desktop'] != true ||
        _string(hello, 'tls_certificate_sha256').toLowerCase() != fingerprint ||
        !relayIdentityMatches) {
      throw const AgentException(AgentError.protocol);
    }
  }

  void _validateAuthentication(Map<String, dynamic> response) {
    if (response['type'] == 'error') {
      throw AgentException(_serverError(_string(response, 'code')));
    }
    if (response['type'] != 'authenticated' && response['type'] != 'paired') {
      throw const AgentException(AgentError.authentication);
    }
    final permissions = response['permissions'];
    if (permissions is! List || !permissions.contains('desktop')) {
      throw const AgentException(AgentError.permission);
    }
  }

  void _send(Map<String, Object?> value) {
    final channel = _channel;
    if (channel == null) throw const AgentException(AgentError.network);
    channel.sink.add(jsonEncode(value));
  }

  void _emit(AgentPhase phase, [AgentError? error]) {
    if (!_states.isClosed) _states.add(AgentState(phase, error));
  }

  Future<void> _fail(AgentError error) async {
    _desktopOpen = false;
    _assembly = null;
    _emit(AgentPhase.failed, error);
    _closed = true;
    try {
      await _channel?.sink.close();
    } on Object {
      // The transport is already unusable.
    }
    _httpClient?.close(force: true);
    final relayTunnel = _relayTunnel;
    _relayTunnel = null;
    if (relayTunnel != null) await EdgeClient.stopTunnel(relayTunnel.id);
  }

  AgentError _serverError(String code) {
    final value = code.toLowerCase();
    if (value.contains('pair') || value.contains('auth')) {
      return AgentError.authentication;
    }
    if (value.contains('totp')) return AgentError.totp;
    if (value.contains('permission') || value.contains('denied')) {
      return AgentError.permission;
    }
    if (value.contains('desktop') || value.contains('session')) {
      return AgentError.desktop;
    }
    return AgentError.unknown;
  }
}

class _FrameAssembly {
  _FrameAssembly({
    required this.sequence,
    required this.width,
    required this.height,
    required this.uncompressedBytes,
    required this.compressedBytes,
    required this.chunkCount,
  });

  final int sequence;
  final int width;
  final int height;
  final int uncompressedBytes;
  final int compressedBytes;
  final int chunkCount;
  final BytesBuilder _bytes = BytesBuilder(copy: false);
  int nextChunk = 0;

  int get length => _bytes.length;

  void add(Uint8List bytes) {
    _bytes.add(bytes);
    nextChunk += 1;
  }

  Uint8List takeBytes() => _bytes.takeBytes();
}

int _integer(Map<String, dynamic> value, String key) {
  final result = value[key];
  if (result is! int) throw const AgentException(AgentError.protocol);
  return result;
}

String _string(Map<String, dynamic> value, String key) {
  final result = value[key];
  if (result is! String) throw const AgentException(AgentError.protocol);
  return result;
}

String _encodeUnpaddedBase64(List<int> bytes) =>
    base64Encode(bytes).replaceAll('=', '');

Uint8List _decodeUnpaddedBase64(String value) {
  if (!RegExp(r'^[A-Za-z0-9+/]*$').hasMatch(value) || value.length % 4 == 1) {
    throw const AgentException(AgentError.protocol);
  }
  final padding = '=' * ((4 - value.length % 4) % 4);
  try {
    return Uint8List.fromList(base64Decode('$value$padding'));
  } on FormatException {
    throw const AgentException(AgentError.protocol);
  }
}