import socket
import os
from picamera2 import Picamera2, Preview
import sys


def save_image(host, port):
    tuning = Picamera2.load_tuning_file("/home/USER/test_python/imx219_70d.json")
    picam2 = Picamera2(tuning=tuning)
    camera_config = picam2.create_still_configuration(
        main={"size": (3280, 2464)},
        controls={
            "AwbEnable": False,
            "ColourGains": (1.0, 1.0)
        }
    )
    picam2.configure(camera_config)
    picam2.start_preview(Preview.NULL)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind((host, port))
    server_socket.listen()

    while True:
        print("Waiting for a connection...")
        try:
            connection, client_address = server_socket.accept()
        except Exception as e:
            print("An error occurred while accepting a connection: ", e)
            continue
        print(connection)
        print("Connection from:", client_address)

        if os.path.exists("/home/USER/test_python/captured_image.jpg"):
            # In order to make "transfer_client" fail later, in case of a capture fail
            os.remove("/home/USER/test_python/captured_image.jpg")

        capture_success = True
        try:
            picam2.start_and_capture_file(
                "/home/USER/test_python/captured_image.jpg", show_preview=False
            )
        except Exception as e:
            print("An error occurred while capturing an image:", e)
            capture_success = False

        with connection:
            try:
                connection.send(
                    b"success" if capture_success else b"fail"
                )
            except Exception as e:
                print("An error occurred while sending result: ", e)
                continue


def main(argv):
    host = argv[1]
    port = int(argv[2])
    print(host, port)
    save_image(host, port)


if __name__ == "__main__":
    main(sys.argv)