from socket import *
import sys
import os
from enum import Enum
from pathlib import Path
from ftp_utils import *
import traceback
HOST = '127.0.0.1'
PORT = 12000
AUTOCONNECT = True

sock = None

def require_sock():
	global sock
	if sock is None:
		raise FTPError("Not connected! Try connecting to a server by running 'open <port>'")

def ftp_client_OPEN(*args):
	global sock
	port = get_argument(args, "port", 0)
	try:
		print(f"Establishing connection to {HOST}:{port}")
		sock = socket(AF_INET, SOCK_STREAM)
		sock.connect((HOST, int(port)))
		return (None, None)
	except TypeError:
		raise FTPError("No port specified")
	except ValueError:
		raise FTPError("Invalid port number. Ports must be given as an integer")
	except ConnectionRefusedError:
		raise FTPError(f"Connection on {HOST}:{port} refused")
	except OverflowError:
		raise FTPError("A valid port is between 0 and 65535", 7783)

def ftp_client_CLOSE(*args):
	global sock
	resp = fetch("CLOSE")
	if resp == "bye<3":
		print("Closing connection")
		sock.close()
		sock = None
	else:
		raise FTPError("Failed to close connection")

def ftp_client_QUIT(*args):
	global sock

	if sock is not None:
		# We don't bother checking if the server closed
		# the connection; we're leaving, and if they can't
		# handle it, that's their problem!
		_resp = fetch("CLOSE")

		print("Closing connection")
		sock.close()

	print("Exiting...")
	sys.exit(0)

def ftp_client_GET(*args):
	global sock
	require_sock()

	file_name = get_argument(args, "file name", 0)
	file_size = int(fetch(f"SIZE {file_name}"))
	parent_path = os.path.join("./client", str(Path(file_name).parent))

	if not os.path.exists(parent_path):
		os.makedirs(parent_path)

	print(f"Downloading {file_name}")
	resp = fetch_bytes(f"GET {file_name}", file_size)
	with open(os.path.join("client", file_name), "wb") as new_file:
		new_file.write(resp)

	print(f"File written to {os.path.join('./client', file_name)}")

def ftp_client_PUT(*args):
	global sock
	require_sock()

	file_name = get_argument(args, "file name", 0)

	try:
		file_size = os.path.getsize(os.path.join("client", file_name))
	except FileNotFoundError:
		raise FTPError("File not found!", 8341)

	ack = fetch(f"PUT {file_name} {file_size}")

	if ack != "OK":
		raise FTPError(f"Server unwilling to accept file '{file_name}'", 2377)

	with open(os.path.join("./client", file_name), "rb") as fetched_file:
		file_bytes = fetched_file.read()
		sock.sendall(file_bytes)
	print(f"File uploaded to ./client/{file_name}")


CLIENT_COMMANDS = {
	"OPEN": ftp_client_OPEN,
	"CLOSE": ftp_client_CLOSE,
	"QUIT": ftp_client_QUIT,
	"GET": ftp_client_GET,
	"PUT": ftp_client_PUT,
}

def fetch_bytes(request, amount=1024):
	global sock
	require_sock()

	sock.sendall(request.encode("utf-8"))
	return sock.recv(amount)

def fetch(request):
	message = fetch_bytes(request).decode("utf-8")

	tokens = message.split()
	if message.split()[0] == "Error":
		raise FTPError(" ".join(tokens[2:]), int(tokens[1][0:-1]))

	return message

FTP_Client_Outcome = Enum("FTP_Client_Outcome", [])

# set up the tcp socket
if AUTOCONNECT:
	try:
		ftp_client_OPEN(PORT)
	except FTPError as e:
		print(f"Error {e.code}: {e.message}")
		print("You may reattempt a connection by running 'OPEN <port>'")

try:
	while (True):
		s = input("Message: ")

		command_name, arguments = parse_ftp_command(s)

		try:
			CLIENT_COMMANDS[command_name](*arguments)
		except KeyError:
			print(f"Error: Unknown command {command_name}")
			continue
		except FTPError as e:
			print(f"Error {e.code}: {e.message}")
except KeyboardInterrupt:
	ftp_client_QUIT()
except Exception as e:
	print(f"Fatal Error: {e}")
	traceback.print_exc()
finally:
	if sock is not None:
		sock.close()