#!/usr/bin/env python3
"""
PTI LEDY - Raspberry Pi 4 port
Requires: sudo pip3 install rpi_ws281x flask
Run with: sudo python3 main.py
"""

import time
import random
import socket
import sys

from rpi_ws281x import PixelStrip, ws

import config
import low_level
import server
from fonts import FONT_7X5

# ---------------------------------------------------------------------------
# Strip init
# ---------------------------------------------------------------------------
strip = PixelStrip(
    config.NUMPIXELS,
    config.PIN,
    freq_hz=800000,
    dma=10,
    invert=False,
    brightness=255,
    channel=0,
    strip_type=ws.WS2811_STRIP_GRB,
)
strip.begin()
low_level.strip = strip

_udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
_udp_addr = ('127.0.0.1', 4210)


def _send_frame():
    buf = bytearray(48 * 16 * 3)
    off = 0
    for y in range(16):
        for x in range(48):
            r, g, b = low_level.get_pixel_color(x, y)
            buf[off] = r
            buf[off + 1] = g
            buf[off + 2] = b
            off += 3
    _udp_sock.sendto(buf, _udp_addr)


# ---------------------------------------------------------------------------
# Global state
# ---------------------------------------------------------------------------
global_node_id = 0

text_nodes = [None] * config.MAX_TEXT_NODES_COUNT
multi_color_text_nodes = [None] * config.MAX_TEXT_NODES_COUNT

for i in range(config.MAX_TEXT_NODES_COUNT):
    text_nodes[i] = {
        'global_id': 0, 'content': '', 'color': (0, 0, 0),
        'char_height': config.MEDIUM_TEXT_HEIGHT, 'char_width': config.MEDIUM_TEXT_WIDTH,
        'pos_x': 0, 'pos_y': 0, 'character_count': 0,
        'scroll_slowness': 0, 'scroll_progress': 0,
        'disappear_time': 0, 'is_scrolling': True, 'is_repeating': False,
    }
    multi_color_text_nodes[i] = {
        'global_id': 0, 'content': '', 'colors': [],  # list of (r,g,b,start_index)
        'char_height': config.MEDIUM_TEXT_HEIGHT, 'char_width': config.MEDIUM_TEXT_WIDTH,
        'pos_x': 0, 'pos_y': 0, 'character_count': 0, 'color_count': 0,
        'scroll_slowness': 0, 'scroll_progress': 0,
        'disappear_time': 0, 'is_scrolling': True, 'is_repeating': False,
    }

main_animation_started = False
bottom_text_state = 0
scroll_node_global_id = 0
scroll_node_active = False

top_text_node_global_id = 0
top_text_cycle_count = 0

program2_active = False
bottom_text_iteration_count = 0
program2_vehicles_spawned = 0
program2_drain_counter = 0
program2_done = False
program2_scroll_progress = 0
program2_spacing_counter = 0
program2_draw_col = 0
program2_current_vehicle = None

dvd_pending = False
dvd_active = False
dvd_x = 0
dvd_y = 0
dvd_dx = 1
dvd_dy = 1
dvd_iteration = 0
dvd_move_counter = 0

DVD_TEXT = "PTI"
DVD_TEXT_W = 17  # 3 chars * 5px + 2 gaps
DVD_TEXT_H = 7
DVD_MAX_X = config.DISPLAY_MAX_X - DVD_TEXT_W + 1  # 31
DVD_MAX_Y = config.DISPLAY_MAX_Y - DVD_TEXT_H + 1  # 9
DVD_COLOR = (0, 0, 50)

VEHICLES = [
    [[0,0,0,0,0,0,0,0],[0,0,1,1,1,1,1,0],[1,1,1,1,1,1,1,1],[0,1,0,0,0,0,1,0]],
    [[0,0,1,1,0,0,0,0],[0,1,1,1,0,0,0,0],[1,1,1,1,1,1,1,1],[0,1,0,0,0,0,1,0]],
    [[0,0,1,1,0,1,1,1],[0,1,1,1,1,1,1,1],[1,1,1,1,1,1,1,1],[0,1,0,0,0,0,1,0]],
    [[0,0,0,0,1,1,0,0],[0,0,1,1,1,1,1,0],[0,1,1,1,1,1,1,1],[0,0,1,0,0,0,1,0]],
    [[0,0,0,0,0,0,0,0],[0,0,0,1,1,1,1,1],[0,0,1,1,1,1,1,1],[0,0,0,1,0,0,1,0]],
    [[1,1,1,1,1,1,1,1],[0,1,0,1,0,1,0,1],[0,1,1,1,1,1,1,1],[0,0,1,0,0,0,1,0]],
    [[0,0,0,0,0,0,0,0],[0,0,1,1,0,0,0,0],[1,1,1,1,1,1,1,1],[0,1,0,0,0,0,1,0]],
    [[1,1,0,1,1,1,1,1],[1,1,0,1,1,1,1,1],[1,1,1,1,1,1,1,1],[0,1,0,0,0,0,1,0]],
    [[0,0,0,1,1,1,1,1],[0,0,1,1,1,1,1,1],[1,1,1,1,1,1,1,1],[0,1,0,0,0,0,1,0]],
    [[0,0,0,0,0,0,0,0],[0,1,1,1,0,0,0,1],[1,1,1,1,1,1,1,0],[0,1,0,0,0,1,0,0]],
    [[0,1,0,1,1,1,1,1],[1,1,0,1,1,1,1,1],[1,1,1,1,1,1,1,1],[0,0,1,0,0,0,1,0]],
    [[0,0,0,0,0,0,0,0],[0,0,0,0,0,1,1,0],[0,0,0,1,1,1,1,1],[0,0,0,0,1,0,1,0]],
]

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _get_text_node_x2(node):
    if node['character_count'] == 0:
        return node['pos_x']
    return node['pos_x'] + (node['char_width'] * node['character_count']) + (node['character_count'] - 1) - 1

def _get_text_node_y2(node):
    return node['pos_y'] + node['char_height'] - 1

def _color_for_char(r, g, b):
    return (r, g, b)

def _add_text_node(content, color, pos_x, pos_y, scroll_slowness, is_scrolling, is_small, disappear_time, is_repeating):
    global global_node_id
    if not content:
        return None
    h = config.SMALL_TEXT_HEIGHT if is_small else config.MEDIUM_TEXT_HEIGHT
    w = config.SMALL_TEXT_WIDTH if is_small else config.MEDIUM_TEXT_WIDTH
    for node in text_nodes:
        if node['disappear_time'] == 0:
            node['content'] = content
            node['color'] = color
            node['pos_x'] = pos_x
            node['pos_y'] = pos_y
            node['char_height'] = h
            node['char_width'] = w
            node['character_count'] = len(content)
            node['scroll_slowness'] = scroll_slowness
            node['scroll_progress'] = 0
            node['is_scrolling'] = is_scrolling
            node['is_repeating'] = is_repeating
            node['disappear_time'] = disappear_time
            node['global_id'] = global_node_id
            global_node_id += 1
            return node
    return None

def _add_multi_color_node(content, colors, pos_x, pos_y, scroll_slowness, is_scrolling, is_small, disappear_time, is_repeating):
    global global_node_id
    if not content:
        return None
    h = config.SMALL_TEXT_HEIGHT if is_small else config.MEDIUM_TEXT_HEIGHT
    w = config.SMALL_TEXT_WIDTH if is_small else config.MEDIUM_TEXT_WIDTH
    for node in multi_color_text_nodes:
        if node['disappear_time'] == 0:
            node['content'] = content
            node['colors'] = list(colors)
            node['color_count'] = len(colors)
            node['pos_x'] = pos_x
            node['pos_y'] = pos_y
            node['char_height'] = h
            node['char_width'] = w
            node['character_count'] = len(content)
            node['scroll_slowness'] = scroll_slowness
            node['scroll_progress'] = 0
            node['is_scrolling'] = is_scrolling
            node['is_repeating'] = is_repeating
            node['disappear_time'] = disappear_time
            node['global_id'] = global_node_id
            global_node_id += 1
            return node
    return None

def _is_node_existing(gid):
    for node in text_nodes:
        if node['global_id'] == gid and node['disappear_time'] != 0:
            return True
    for node in multi_color_text_nodes:
        if node['global_id'] == gid and node['disappear_time'] != 0:
            return True
    return False

def _reset_global_id():
    global global_node_id
    global_node_id = 2

# ---------------------------------------------------------------------------
# Drawing
# ---------------------------------------------------------------------------

def _draw_char(ch, height, width, r, g, b, cursor):
    if ord(ch) < ord('!') or ord(ch) > ord('~'):
        ch = ' '
    glyph = FONT_7X5[ord(ch) - ord(' ')]
    low_level.draw_character(glyph, height, width, r, g, b, cursor)

def scroll_all_texts(split_scroll_mode=False):
    for node in text_nodes:
        if node['disappear_time'] == 0:
            continue
        if node['is_scrolling']:
            if node['scroll_slowness'] > node['scroll_progress']:
                node['scroll_progress'] += 1
            else:
                node['scroll_progress'] = 0
                x2 = _get_text_node_x2(node)
                if split_scroll_mode or node['pos_y'] >= 7:
                    if x2 < 0:
                        if node['is_repeating']:
                            node['pos_x'] = config.DISPLAY_MAX_X
                        else:
                            node['disappear_time'] = 1
                        continue
                    node['pos_x'] -= 1
                else:
                    if node['pos_x'] > config.DISPLAY_MAX_X:
                        if node['is_repeating']:
                            node['pos_x'] = -x2 + node['pos_x']
                        else:
                            node['disappear_time'] = 1
                        continue
                    node['pos_x'] += 1

        cursor = {'x': node['pos_x'], 'y': node['pos_y']}
        r, g, b = node['color']
        for ch in node['content']:
            _draw_char(ch, node['char_height'], node['char_width'], r, g, b, cursor)

def scroll_all_multi_color_texts(split_scroll_mode=False):
    for node in multi_color_text_nodes:
        if node['disappear_time'] == 0:
            continue
        if node['is_scrolling']:
            if node['scroll_slowness'] > node['scroll_progress']:
                node['scroll_progress'] += 1
            else:
                node['scroll_progress'] = 0
                x1 = node['pos_x']
                x2 = x1 + (node['char_width'] * node['character_count']) + (node['character_count'] - 1) - 1
                if split_scroll_mode or node['pos_y'] >= 7:
                    if x2 < 0:
                        if node['is_repeating']:
                            node['pos_x'] = config.DISPLAY_MAX_X
                        else:
                            node['disappear_time'] = 1
                        continue
                    node['pos_x'] -= 1
                else:
                    if x1 > config.DISPLAY_MAX_X:
                        if node['is_repeating']:
                            node['pos_x'] = -x2 + x1
                        else:
                            node['disappear_time'] = 1
                        continue
                    node['pos_x'] += 1

        cursor = {'x': node['pos_x'], 'y': node['pos_y']}
        color_index = 0
        r, g, b, _ = node['colors'][0]
        for j, ch in enumerate(node['content']):
            if color_index < node['color_count'] - 1 and node['colors'][color_index + 1][3] == j:
                color_index += 1
                r, g, b, _ = node['colors'][color_index]
            _draw_char(ch, node['char_height'], node['char_width'], r, g, b, cursor)

def handle_disappear_timers():
    for node in text_nodes:
        if node['disappear_time'] > 0:
            node['disappear_time'] -= 1
            if node['disappear_time'] == 0:
                low_level.fill_pixels(
                    node['pos_x'], node['pos_y'],
                    _get_text_node_x2(node), _get_text_node_y2(node),
                    0, 0, 0
                )

def handle_multi_color_disappear_timers():
    for node in multi_color_text_nodes:
        if node['disappear_time'] > 0:
            node['disappear_time'] -= 1
            if node['disappear_time'] == 0:
                x2 = node['pos_x'] + (node['char_width'] * node['character_count']) + (node['character_count'] - 1) - 1
                low_level.fill_pixels(
                    node['pos_x'], node['pos_y'],
                    x2, node['pos_y'] + node['char_height'] - 1,
                    0, 0, 0
                )

# ---------------------------------------------------------------------------
# Programs
# ---------------------------------------------------------------------------

def handle_program1():
    global main_animation_started, scroll_node_active, scroll_node_global_id
    global bottom_text_state, bottom_text_iteration_count, program2_active
    global top_text_node_global_id, top_text_cycle_count
    global dvd_pending, dvd_active

    s = server.state

    # --- DVD active: handled entirely in loop() ---
    if dvd_active:
        return

    # --- DVD pending: wait for bottom area to clear ---
    if dvd_pending:
        if scroll_node_active and not _is_node_existing(scroll_node_global_id):
            scroll_node_active = False
        if not scroll_node_active and not program2_active:
            dvd_pending = False
            dvd_active = True
            reset_dvd()
            low_level.fill_pixels(0, 0, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
        return

    # --- Top text management ---
    if not s['program_top_text_enabled']:
        for node in multi_color_text_nodes:
            if node['disappear_time'] != 0:
                node['disappear_time'] = 1
        main_animation_started = False
    elif not main_animation_started:
        if not _is_node_existing(0) and not _is_node_existing(1):
            main_animation_started = True
            node = _add_multi_color_node(
                "Witamy w PTI",
                [(40, 40, 40, 0), (0, 0, 50, 9)],
                0, 0, 3, True, True, -1, False
            )
            if node:
                top_text_node_global_id = node['global_id']
    elif not _is_node_existing(top_text_node_global_id):
        # Top text cycle ended
        top_text_cycle_count += 1

        force_dvd = s.get('force_dvd', False)
        if force_dvd:
            with server._lock:
                server.state['force_dvd'] = False

        if force_dvd or (top_text_cycle_count >= config.DVD_TOP_TEXT_CYCLES
                         and random.randint(0, 99) < config.DVD_TRIGGER_CHANCE):
            top_text_cycle_count = 0
            dvd_pending = True
            main_animation_started = False
            return

        # Re-add top text from right edge
        node = _add_multi_color_node(
            "Witamy w PTI",
            [(40, 40, 40, 0), (0, 0, 50, 9)],
            config.DISPLAY_MAX_X, 0, 3, True, True, -1, False
        )
        if node:
            top_text_node_global_id = node['global_id']

    # --- Bottom text management ---
    if not s['program_bottom_text_enabled']:
        if scroll_node_active:
            for node in text_nodes:
                if node['disappear_time'] != 0:
                    node['disappear_time'] = 1
            scroll_node_active = False
        if program2_active:
            program2_active = False
        return

    if program2_active:
        return

    if (main_animation_started or not s['program_top_text_enabled']) and \
       (not scroll_node_active or not _is_node_existing(scroll_node_global_id)):

        if scroll_node_active:
            bottom_text_iteration_count += 1
            scroll_node_active = False
            force = s.get('force_vehicles', False)
            if s['program_vehicles_enabled'] and (force or bottom_text_iteration_count >= config.PROGRAM2_TEXT_ITERATIONS):
                if force or random.randint(0, 99) < config.PROGRAM2_TRIGGER_CHANCE:
                    if force:
                        with server._lock:
                            server.state['force_vehicles'] = False
                    bottom_text_iteration_count = 0
                    program2_active = True
                    reset_program2()
                    low_level.fill_pixels(0, 9, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
                    return

        texts = [
            ("Technik informatyk + ai, robotyka", (0, 0, 50)),
            ("Technik programista",               (50, 40, 0)),
            ("Technik poligrafii i grafiki komputerowej", (0, 45, 0)),
            ("Technik reklamy",                   (60, 20, 0)),
        ]
        text, color = texts[bottom_text_state]
        node = _add_text_node(text, color, config.DISPLAY_MAX_X, 9, 1, True, True, -1, False)
        if node is not None:
            scroll_node_global_id = node['global_id']
            scroll_node_active = True
            bottom_text_state = (bottom_text_state + 1) % 4


def reset_dvd():
    global dvd_x, dvd_y, dvd_dx, dvd_dy, dvd_iteration, dvd_move_counter
    dvd_x = random.randint(0, DVD_MAX_X)
    dvd_y = random.randint(0, DVD_MAX_Y)
    dvd_dx = random.choice([-1, 1])
    dvd_dy = random.choice([-1, 1])
    dvd_iteration = 0
    dvd_move_counter = 0


def handle_dvd():
    global dvd_x, dvd_y, dvd_dx, dvd_dy, dvd_iteration, dvd_move_counter
    global dvd_active, main_animation_started

    if dvd_iteration >= config.DVD_ITERATIONS:
        dvd_active = False
        main_animation_started = False
        return

    dvd_iteration += 1

    dvd_move_counter += 1
    if dvd_move_counter >= 2:
        dvd_move_counter = 0

        new_x = dvd_x + dvd_dx
        new_y = dvd_y + dvd_dy

        if new_x < 0 or new_x > DVD_MAX_X:
            dvd_dx = -dvd_dx
            new_x = dvd_x + dvd_dx
        if new_y < 0 or new_y > DVD_MAX_Y:
            dvd_dy = -dvd_dy
            new_y = dvd_y + dvd_dy

        dvd_x = new_x
        dvd_y = new_y

    cursor = {'x': dvd_x, 'y': dvd_y}
    for ch in DVD_TEXT:
        _draw_char(ch, config.SMALL_TEXT_HEIGHT, config.SMALL_TEXT_WIDTH, *DVD_COLOR, cursor)


def draw_vehicle(index, pos_x, pos_y):
    vr, vg, vb = config.VEHICLE_COLOR
    for row in range(4):
        for col in range(8):
            if VEHICLES[index][row][col]:
                low_level.set_pixel(pos_x + col, pos_y + row, vr, vg, vb)


def reset_program2():
    global program2_scroll_progress, program2_spacing_counter
    global program2_vehicles_spawned, program2_drain_counter, program2_done
    global program2_draw_col, program2_current_vehicle
    program2_scroll_progress = 0
    program2_spacing_counter = 0
    program2_vehicles_spawned = 0
    program2_drain_counter = 0
    program2_done = False
    program2_draw_col = 0
    program2_current_vehicle = None


def handle_program2(spacing=4, scroll_slowness=1, y=0, max_vehicles=0):
    global program2_scroll_progress, program2_spacing_counter
    global program2_vehicles_spawned, program2_drain_counter, program2_done
    global program2_draw_col, program2_current_vehicle

    if scroll_slowness > program2_scroll_progress:
        program2_scroll_progress += 1
        return
    program2_scroll_progress = 0

    low_level.shift_rectangle_left(0, y, config.DISPLAY_MAX_X, y + 3, 1)

    if max_vehicles > 0 and program2_vehicles_spawned >= max_vehicles:
        program2_drain_counter += 1
        if program2_drain_counter >= config.DISPLAY_MAX_X + 1:
            program2_done = True
        return

    if program2_spacing_counter > 0:
        program2_spacing_counter -= 1
        return

    veh = VEHICLES[random.randint(0, 11)] if program2_draw_col == 0 else program2_current_vehicle
    if program2_draw_col == 0:
        program2_current_vehicle = veh

    vr, vg, vb = config.VEHICLE_COLOR
    for row in range(4):
        if veh[row][program2_draw_col]:
            low_level.set_pixel(config.DISPLAY_MAX_X, y + row, vr, vg, vb)

    program2_draw_col += 1
    if program2_draw_col >= 8:
        program2_draw_col = 0
        program2_spacing_counter = spacing
        if max_vehicles > 0:
            program2_vehicles_spawned += 1

# ---------------------------------------------------------------------------
# Main loop
# ---------------------------------------------------------------------------

def setup():
    strip.setBrightness(255)
    for i in range(config.NUMPIXELS):
        strip.setPixelColor(i, 0)
    strip.show()
    server.start_server(port=80)
    print("Server started. Running display loop.")


def loop():
    global program2_active

    s = server.state

    if dvd_active:
        low_level.fill_pixels(0, 0, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
        handle_dvd()
        strip.show()
        _send_frame()
        return

    low_level.fill_pixels(0, 0, config.DISPLAY_MAX_X, 7, 0, 0, 0)
    handle_program1()

    if dvd_active:
        strip.show()
        _send_frame()
        return

    if s['program_top_text_enabled']:
        scroll_all_multi_color_texts(split_scroll_mode=True)
        handle_multi_color_disappear_timers()

    if not s['program_bottom_text_enabled']:
        if program2_active:
            program2_active = False
        low_level.fill_pixels(0, 9, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
    elif program2_active and not s['program_vehicles_enabled']:
        program2_active = False
        low_level.fill_pixels(0, 9, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
    elif program2_active:
        handle_program2(spacing=4, scroll_slowness=1, y=12, max_vehicles=config.PROGRAM2_VEHICLE_COUNT)
        if program2_done:
            program2_active = False
            low_level.fill_pixels(0, 9, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
    else:
        low_level.fill_pixels(0, 9, config.DISPLAY_MAX_X, config.DISPLAY_MAX_Y, 0, 0, 0)
        scroll_all_texts(split_scroll_mode=True)
        handle_disappear_timers()

    strip.show()
    _send_frame()


if __name__ == '__main__':
    setup()
    try:
        while True:
            loop()
            time.sleep(0.02)  # ~50 fps
    except KeyboardInterrupt:
        for i in range(config.NUMPIXELS):
            strip.setPixelColor(i, 0)
        strip.show()
        print("\nStopped.")
        sys.exit(0)