# This is a compatible server for the TEF6686 firmware, but instead using the i2c control mode
# This fully works with FM-DX-Webserver (no plugins) + Spectrum plugin
# Released under the Unlicense (however both FM-DX-Webserver and the firmware are under the tyranny of GPL-3)

import socket, os, time
import hashlib, secrets, string
from libtimer2 import Timer
from driver.tef import TEF6686
from driver.protocol import I2CPCClient
from functools import wraps
from typing import Callable
from dataclasses import dataclass

STEREO_BLENDING = False

INITIAL_FREQ = 9500
INITIAL_EQ = True
INITIAL_IMS = True

HOST = os.getenv("HOST") or '0.0.0.0'
PORT = int(os.getenv("PORT") or 0) or 7373
DEVICE = os.getenv("DEV", "COM6")

DEFAULT_PW = "test"
PASSWORD = os.getenv("PW", DEFAULT_PW)
if PASSWORD == DEFAULT_PW: print("Password left at default. Please change if neccesary")

CLOCK_ENV = os.getenv("CLK", "dp").strip().lower()
if CLOCK_ENV == "dp": CLOCK = 12000000 # DP-666
elif CLOCK_ENV == "x": CLOCK = 9216000
else: CLOCK = int(CLOCK_ENV)

SALT_LENGTH = 16
SS_UPDATE_INTERVAL = 0.1
RDS_UPDATE_INTERVAL = 0.086

@dataclass
class State:
    last_tune: int = INITIAL_FREQ
    last_eqims: int = (INITIAL_EQ << 1) | INITIAL_IMS
    forced_mono: bool = False
    deemp: int = 0
    bw: int = 0
    scan_start: int = 8750
    scan_stop: int = 10800
    scan_step: int = 10
    scan_bw: int = 0
    last_bw_extend: bool = False

def init_tef():
    tef = TEF6686(I2CPCClient(DEVICE, int(os.getenv("BAUD") or 0) or 115200))
    tef.init(clock=CLOCK)
    tef.AUDIO_Set_Mute(False)
    tef.FM_Tune_To(TEF6686.TuneTo_Mode.Preset, INITIAL_FREQ)
    tef.FM_Set_RDS(1)
    tef.FM_Set_ChannelEqualizer(INITIAL_EQ)
    tef.FM_Set_MphSuppression(INITIAL_IMS)

    tef.FM_Set_LevelStep(-1, -1, -1, -1, -4, -8, 0)
    tef.FM_Set_Highcut_Mph(0)
    tef.FM_Set_Highcut_Max(False)
    tef.FM_Set_Lowcut_Max(True, 60)
    tef.FM_Set_Stereo_Time(60, 120, 100, 200)

    tef.FM_Set_Stereo_Max(STEREO_BLENDING)
    tef.APPL_Set_OperationMode(True) # Turn off
    return tef

def authenticate(conn: socket.socket):
    salt = "".join(secrets.choice(string.ascii_lowercase) for _ in range(SALT_LENGTH))
    conn.sendall(salt.encode() + b"\n")
    expected_hash = hashlib.sha1((salt + PASSWORD).encode()).hexdigest().encode()

    while True:
        data = conn.recv(len(expected_hash))
        if not data: break

        if data.strip() == expected_hash.strip(): 
            conn.sendall(b"o1,0\nOK\n")
            return True
        else: 
            conn.sendall(b"a0\n") # You failed
            break
    conn.close()
    return False

def process_command(tef: TEF6686, data: bytes, state: State, conn: socket.socket):
    out = b""
    for cmd in data.splitlines():
        if cmd.startswith(b"T"):
            freq = int(cmd.decode().removeprefix("T").strip()) // 10
            if freq < 6500 or freq > 10800: continue
            tef.FM_Tune_To(TEF6686.TuneTo_Mode.FM_Jump, freq)
            state.last_tune = freq
            out += f"T{freq*10}\n".encode()
        elif cmd.startswith(b"G"):
            eqims = int(cmd.decode().removeprefix("G").strip(), 2)
            tef.FM_Set_ChannelEqualizer((eqims & 1) == 1)
            tef.FM_Set_MphSuppression((eqims & 2) == 2)
            out += f"G{bin(eqims).removeprefix('0b').zfill(2)}\n".encode()
            state.last_eqims = eqims
        elif cmd.startswith(b"B"):
            state.forced_mono = bool(int(cmd.decode().removeprefix("B").strip(), 2))
            tef.FM_Set_Stereo_Min(2 if state.forced_mono else 0)
            out += f"B{int(state.forced_mono)}\n".encode()
        elif cmd.startswith(b"D"):
            state.deemp = int(cmd.decode().removeprefix("D").strip())
            dtime = 500 if state.deemp == 0 else (750 if state.deemp == 1 else 0)
            tef.FM_Set_Deemphasis(dtime)
            out += f"D{state.deemp}\n".encode()
        elif cmd.startswith(b"x"): tef.FM_Tune_To(TEF6686.TuneTo_Mode.Preset, state.last_tune)
        elif cmd.startswith(b"X"): tef.APPL_Set_OperationMode(True) # turn off
        elif cmd.startswith(b"W"):
            state.bw = int(cmd.decode().removeprefix("W").strip())
            auto = (state.bw == 0)
            tef.FM_Set_Bandwidth(auto, 2360 if auto else (state.bw // 100))
            out += f"W{state.bw}\n".encode()
        elif cmd.startswith(b"?"): out += b">XRD Python driver\n"
        elif cmd.startswith(b"S"):
            cmd = cmd[1:]
            if cmd.strip():
                arg = int(cmd.decode()[1:].strip())
                match cmd[0]:
                    case 97: state.scan_start = (arg + 5) // 10
                    case 98: state.scan_stop = (arg + 5) // 10
                    case 99: state.scan_step = (arg + 5) // 10
                    case 119: state.scan_bw = arg
            else:
                start = True

                tef.FM_Set_Bandwidth((state.scan_bw == 0), state.scan_bw)
                conn.sendall(b"U")
                for freq in range(state.scan_start, state.scan_stop + state.scan_step, state.scan_step):
                    if not start: conn.sendall(b", ") # Prevent trailing comma, because the FM-DX-Webserver spectrum plugin treats us as actual firmware and throws as harder api, without it we're treated as a module
                    start = False

                    tef.FM_Tune_To(TEF6686.TuneTo_Mode.Search, freq) # Auto mutes, less commands sent
                    time.sleep(0.0067) # sick seven
                    _, level, *_ = d if (d := tef.FM_Get_Quality_Data()) else (None, None)
                    if level is None: continue
                    conn.sendall(str(freq * 10).encode() + b" = " + str((level / 10) + 11.25).encode())
                conn.sendall(b"\n")

                tef.FM_Tune_To(TEF6686.TuneTo_Mode.Preset, state.last_tune)
                tef.FM_Set_Bandwidth((state.bw == 0), 2360 if (state.bw == 0) else (state.bw // 100))

    return out

PERIODIC_FUNCTIONS: list[tuple[Callable, float, Timer]] = []

def periodic(t: float):
    def decorator(func):
        PERIODIC_FUNCTIONS.append((func, t, Timer()))
        @wraps(func)
        def wrapper(*args, **kwargs): return func(*args, **kwargs)
        return wrapper
    return decorator

def reset_periodic():
    for (_, _, timer) in PERIODIC_FUNCTIONS: timer.reset()

def run_periodic(*args, **kwargs):
    for (func, t, timer) in PERIODIC_FUNCTIONS:
        max_catchup = 10
        count = 0
        while timer.get_time() > t and count < max_catchup:
            func(*args, **kwargs)
            timer.subtract(t)
            count += 1

@periodic(SS_UPDATE_INTERVAL)
def send_signal_status(tef: TEF6686, conn: socket.socket, state: State):
    res = tef.FM_Get_Quality_Data()
    if res is None: return

    status, level, usn, wam, _, bandwidth, *_ = res

    ms_since_tune = (status & 0x3ff) / 10
    if ms_since_tune < 45: return # Give only "quality data"

    level = level / 10

    data = b"S"

    if STEREO_BLENDING:
        stereo_pilot, _ = d if (d := tef.FM_Get_Signal_Status()) else (None, None)
        if state.forced_mono: data += b"M"
        elif stereo_pilot: data += b"s"
        else: data += b"m"
    else:
        if state.forced_mono: data += b"M"
        else: data += b"S"

    if not state.last_bw_extend and level > 38 and usn < 7:
        tef.FM_Set_Bandwidth_Options(400)
        state.last_bw_extend = True
    if state.last_bw_extend and (level < 33 or usn > 10):
        tef.FM_Set_Bandwidth_Options()
        state.last_bw_extend = False

    conn.sendall(data + f"{level + 11.25},{wam//10},{usn//10},{bandwidth}\n\n".encode())

@periodic(RDS_UPDATE_INTERVAL)
def send_rds_data(tef: TEF6686, conn: socket.socket, state: State):
    res = tef.FM_Get_RDS_Data__decoder()
    if res is None: return
    status, A, B, C, D, dec_error = res
    dec_error >>= 8

    if (status & (1 << 9) == 0) or (status & (1 << 15) == 0) or (status & (1 << 13) == (1 << 13)): return

    data = b""
    if (status & (1 << 13) == 0):
        err = 0
        err |= (dec_error & 0x30) >> 4
        err |= (dec_error & 0xC)
        err |= (dec_error & 3) << 4

        data = b"R"
        data += f"{B:04X}{C:04X}{D:04X}{err:02X}\n".encode()

    pi_error = (dec_error >> 6) & 0b11
    if pi_error < 3:
        data += b"P" + f"{A:04X}".encode()
        data += b"?" * pi_error + b"\n"
    elif status & (1 << 12):
        pi_error = (dec_error >> 2) & 0b11
        if pi_error < 3:
            data += b"P" + f"{C:04X}".encode()
            data += b"?" * pi_error + b"\n"
    conn.sendall(data)

def run_server():
    with init_tef() as tef:
        state = State()

        device, hw_version, sw_version = d if (d := tef.APPL_Get_Identification()) else (None, None, None)
        if device and hw_version and sw_version:
            variant = device & 127
            hw_major = hw_version >> 8
            hw_minor = hw_version & 127
            sw_major = sw_version >> 8
            sw_minor = sw_version & 127

            variant_str = "TEF6686"
            if variant == 1: variant_str = "TEF6687"
            elif variant == 9: variant_str = "TEF6688"
            elif variant == 3: variant_str = "TEF6689"

            print(f"{variant_str} (V{hw_major}{hw_minor}{sw_major}.{sw_minor})")
        else: print("Could not fetch variant")

        with socket.socket() as s:
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            s.bind((HOST, PORT))
            s.listen()
            print(f"Server listening on {HOST}:{PORT}")

            while True:
                conn, addr = s.accept()
                with conn:
                    reset_periodic()
                    print(f"Connected by {addr}")
                    if not authenticate(conn):
                        print("Authentication failed.")
                        continue
                    conn.setblocking(False)

                    # Send initial state
                    conn.sendall(f"T{state.last_tune*10}\n".encode())
                    conn.sendall(f"G{bin(state.last_eqims).removeprefix('0b').zfill(2)}\n".encode())
                    conn.sendall(f"B{int(state.forced_mono)}\n".encode())
                    conn.sendall(f"D{state.deemp}\n".encode())
                    conn.sendall(f"W{state.bw}\n".encode())
                    conn.sendall(f"M0\n".encode())

                    while True:
                        try:
                            if not (data := conn.recv(1024)): break
                            try:
                                resp = process_command(tef, data, state, conn)
                                if resp: conn.sendall(resp)
                            except Exception as e:
                                try: conn.sendall(b"!" + str(e).encode() + b"\n")
                                except Exception: pass
                        except ConnectionResetError: break
                        except ConnectionAbortedError: break
                        except BlockingIOError:
                            try: run_periodic(tef, conn, state)
                            except Exception as e:
                                try: conn.sendall(b"!" + str(e).encode() + b"\n")
                                except Exception: pass
                            time.sleep(0.03)

import traceback
if __name__ == "__main__": 
    try: run_server()
    except Exception as e: traceback.print_exception(e)