from enum import Enum
import enum
import time
import RPi.GPIO as GPIO
from AlphaBot2 import AlphaBot2
from TRSensors import TRSensor
import datetime

class LINE_MODE(Enum):
    UNSET = -1
    WHITE_LINE_MODE = 0
    BLACK_LINE_MODE = 1

class STATE(Enum):
    INIT_SUCCESS = 0
    INIT_FAIL = -1

BTN_PIN = 7
GPIO.setmode(GPIO.BCM)
GPIO.setwarnings(False)
GPIO.setup(BTN_PIN, GPIO.IN, pull_up_down=GPIO.PUD_UP)

MAX_SPEED = 40
OVERDRIVE_SPEED = 50
MAX_REVERSE_SPEED = 20
MAX_CORRECTION_STRENGTH = 5

LineColorDescriptor: LINE_MODE = LINE_MODE.UNSET
CurrentCorrectionStrength: int = 1
logging: bool = True
LineSearchingDirection: int = 0 # -1 -> left, 1 -> right

TR = TRSensor()
Ab = AlphaBot2()
Ab.stop()

def log(msg: str) -> None:
    if not logging: return;
    print("-> Logger")
    print(f"Timestamp: {datetime.datetime.now()}")
    print(f"Message: {msg}")
    print("<- Logger")

def LineModeBoolConverter(mode: LINE_MODE) -> bool:
    return True if mode == LINE_MODE.WHITE_LINE_MODE else False

def setup() -> int:
    log("Initializing TR!")
    for i in range(1, 4):
        TR.calibratedMin[i] = 1023
        TR.calibratedMax[i] = 0

    log("Calibration sensors!")
    for _ in range(100):
        values = TR.AnalogRead()
        for i in range(1, 5):  # Index 1 bis 3 → Sensor 2–4
            TR.calibratedMin[i] = min(TR.calibratedMin[i], values[i])
            TR.calibratedMax[i] = max(TR.calibratedMax[i], values[i])
    Ab.stop()
    time.sleep(0.5)

    # Determine line color
    sensor_values = TR.readCalibrated()
    center_avg = sum(sensor_values[1:4]) / 3
    outer_avg = (sensor_values[0] + sensor_values[4]) / 2
    LineColorDescriptor = LINE_MODE.WHITE_LINE_MODE if outer_avg < center_avg else LINE_MODE.BLACK_LINE_MODE # If outer sensors are dark → white line on dark
    log("Detected line type: " + "WHITE on BLACK" if LineColorDescriptor == LINE_MODE.WHITE_LINE_MODE else "BLACK on WHITE")

    # Wait for button press to start
    print("Press the button to start line following")
    while GPIO.input(BTN_PIN):
        time.sleep(0.1)

    LineSearchingDirection = 0

    log("Setup done!")
    return 0;

def loop() -> None:
    Ab.forward()
    while True:
        position, sensors = TR.readLine(white_line=LineModeBoolConverter(LineColorDescriptor))
        background_level = (sensors[0] + sensors[1]) / 2

        correction_strength = (background_level / 1023.0) * 2 * MAX_CORRECTION_STRENGTH - MAX_CORRECTION_STRENGTH

        if sensors[2] < background_level:
            Ab.backward()
            Ab.setPWMA(MAX_REVERSE_SPEED)
            Ab.setPWMB(MAX_REVERSE_SPEED)
            while True:
                position, sensors = TR.readLine(white_line=LineModeBoolConverter(LineColorDescriptor))
                if sensors[2] >= background_level:
                                    break

            Ab.forward()
            continue

        if sensors[1] < background_level:
            Ab.setPWMA(OVERDRIVE_SPEED)
            Ab.setPWMB(MAX_SPEED)

        elif sensors[3] < background_level:
            Ab.setPWMB(OVERDRIVE_SPEED)
            Ab.setPWMA(MAX_SPEED)

        else:
            if correction_strength < 0:
                Ab.setPWMB(int(MAX_SPEED + correction_strength)) # decrease power from right motor if sensor indicates missing line left
            elif correction_strength > 0:
                Ab.setPWMA(int(MAX_SPEED - correction_strength))

        # time.sleep(0.01); # we drive fast so handle steering as fast as possible

if __name__ == '__main__':
    if setup():
        loop()
    else:
        log("Setup failed!")