from time import time, gmtime, mktime,sleep
import network, machine, ntptime, sockets, json, requests

def time_from_iso(datetime_str):
    [yyyy, mm, dd] = [int(i) for i in datetime_str.split('T')[0].split('-')]
    [h,m,s] = [int(i) for i in datetime_str.split('T')[1].split(':')]
    time_u = mktime((yyyy, mm, dd, h, m, s, 0, 0))
    return time_u

def start_ap():
    # Sets up initial captive portal for configuring wifi and basic options
    ap = network.WLAN(network.AP_IF)
    ap.active(true)

    ap.config(
        essid="CTA-Setup",
        authmode=network.AUTH_OPEN
    )

    print("AP active for initial setup: ", ap.active())
    print("AP IP address:", ap.ifconfig()[0])

    return ap

def parse_form_data(body):
    """Parse URL-encoded form data"""
    data = {}
    pairs = body.split("&")
    for pair in pairs:
        if "=" in pair:
            key, value = pair.split("=", 1)
            # URL decode
            data[key] = value.replace("+", " ").replace("%40", "@")
    return data

def save_config(data):
    """Save Wi-Fi credentials to config.json"""
    if not data.get("SSID"):
        raise ValueError("SSID is required")
    
    config = {
        "SSID": data.get("SSID"),
        "PASS": data.get("PASS", ""),
        "APIKEY": data.get("APIKEY", "")
    }
    
    with open("config.json", "w") as f:
        json.dump(config, f)
    
    print("Configuration saved:", config)

def start_http_server():
    addr = socket.getaddrinfo("0.0.0.0", 80)[0][-1]
    s = socket.socket()
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(addr)
    s.listen(1)

    print("HTTP server listening on port 80")

    while True:
        try:
            conn, addr = s.accept()
            request = conn.recv(2048).decode()

            if request.startswith("POST /save"):
                # Handle form submission
                body = request.split("\r\n\r\n", 1)[1]
                form_data = parse_form_data(body)
                save_config(form_data)

                response_body = """<!DOCTYPE html>
<html>
<head>
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Setup Complete</title>
    <style>
        body { font-family: sans-serif; padding: 16px; text-align: center; }
        h2 { color: #28a745; }
    </style>
</head>
<body>
    <h2>Configuration saved</h2>
    <p>Device will reboot and connect to Wi-Fi.</p>
    <p><small>You can close this window.</small></p>
</body>
</html>"""
            else:
                # Serve the setup page for any other request
                try:
                    with open("portal.html") as f:
                        response_body = f.read()
                except:
                    response_body = "<h1>Setup Page Not Found</h1><p>Please upload portal.html</p>"

            response = (
                "HTTP/1.1 200 OK\r\n"
                "Content-Type: text/html\r\n"
                "Connection: close\r\n\r\n"
                + response_body
            )

            conn.send(response.encode())

        except Exception as e:
            print("HTTP error:", e)
            try:
                error = f"<h2>Error</h2><p>{e}</p>"
                conn.send(
                    "HTTP/1.1 500 Internal Server Error\r\n"
                    "Content-Type: text/html\r\n"
                    "Connection: close\r\n\r\n"
                    .encode() + error.encode()
                )
            except:
                pass
        finally:
            try:
                conn.close()
            except:
                pass

def reconnect_wifi():
    while not wlan.isconnected():
        try:
            wlan.connect(SSID, PASS)
            max_wait = 10
            while max_wait > 0:
                if wlan.isconnected():
                    break
                max_wait -= 1
                print('Waiting for connection...')
                sleep(1)
            if wlan.isconnected():
                print('Connected to Wi-Fi')
                print('Network config:', wlan.ifconfig())
                break
        except Exception as e:
            print('Connection error:', e)
        print('Retrying connection...')
        sleep(2)

def set_LEDs(stations):
    # Sets up LEDs and their ports.
    for station in stations:
        led = machine.Pin(station['led_index'], machine.Pin.OUT)
        led.value(station['state'])

# Synchronize time of RTC with NTP
last_ntp_sync = -3700 # Init as more than an hour ago
last_api_query = mktime([2000,1,1,0,0,0,0,0]) # Set last query default (will run on first run)
query_period = 120 # Query the API at most every 2 minutes
distance_on = 10 # Time from station before turning on LED

stations = {}

# Check if already configured and load config for Wifi, API Key etc.
try:
    with open("config.json", "r") as f:
        config = json.load(f)
    APIKEY = config.APIKEY
    SSID = config.SSID
    PASS = config.PASS
    # Won't scale well for infinite LEDs
    stations["L"] = { 'id':config.LStation, 'eta':[], 'led_index':15, 'state':0 }
    stations["R"] = { 'id':config.RStation, 'eta':[], 'led_index':16, 'state':0 }
    # Function sets up LEDs.
    set_LEDs()
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    reconnect_wifi()
    start_http_server()   # Still launch server to allow changing SSID and config
    demo = False
except:
    # Failed to open file, launch AP mode
    start_ap()
    start_http_server()

    # Halt waiting for config
    while(True):
        sleep(5)

while(True):
    # Synchronise NTP time if hasn't been done in 1 hour
    if (time() - last_ntp_sync) > 3600:
        reconnect_wifi()
        try:
            ntptime.settime()
            print("Time synchronized successfully")
            last_ntp_sync = time()
        except Exception as e:
            print("Error synchronizing time:", e)
            
    # Have we queried API in the query_period seconds?
    if (time() - last_api_query) > query_period:
        # Query updated API
        #trains = (requests.get("https://lapi.transitchicago.com/api/1.0/ttpositions.aspx?key="+APIKEY+"&rt=Blue&outputType=JSON")).json()['ctatt']
        # Query all stations
        for station,dict in stations.item():
            trains_s = (requests.get(f"https://lapi.transitchicago.com/api/1.0/ttarrivals.aspx?key={APIKEY}&mapid={dict['id']}outputType=JSON")).json()['ctatt']
        
            # Timestamps are served in local Chicago time, convert to GMT for simplicity on this side.
            #print(f"Time of poll is: "+str(gmtime(time_from_iso(trains_s['tmst'])+(5*60*60))))
        
            # Loop through all trains and update matching stations
            for train in trains_s['eta']:
                eta_time = time_from_iso(train['arrT'])+(5*60*60)
                
                # Append the ETA, we rely on the station loop to remove old trains
                stations[station]['eta'].append(eta_time)
        
        
        last_api_query = time() # uptime polling window
    
    # Process the ETA for each station 
    for station, dict in stations.items(): # loop through all stations          
        # Set LED state tracker to 0, if no trains, it'll stay 0.
        stations[station]['state'] = 0
        for index, eta in enumerate(dict['eta']):   # multiple trains. Enumerate is used so we have the index to update the dict itself.
            time_from_now = eta - time()

            print('Train will arrive at'+dict['name']+' in '+ str(time_from_now) + ' seconds')
            if time_from_now < distance_on: 
                stations[station]['state'] = 1
        
            if time_from_now < -distance_on: # Remove old trains
                stations[station]['eta'].pop(index)
        
        # Set Station LED state
        set_LEDs()

    sleep(1)