#!/usr/bin/env python3
"""
SiemaszkoLab Logic Analyzer Bridge
===================================
Most miedzy analizatorem logicznym USB (sigrok/Saleae clone) a strona SiemaszkoLab.
Przechwytuje dane z analizatora przez sigrok-cli i streamuje je na serwer.

Wymagania:
  - Python 3.8+
  - sigrok-cli (https://sigrok.org/wiki/Downloads)
  - requests (pip install requests)

Uzycie:
  python slab_bridge.py                        # domyslne ustawienia
  python slab_bridge.py --channels 4 --rate 8M # 4 kanaly, 8 MHz
  python slab_bridge.py --list-devices         # lista podlaczonych urzadzen

Autor: SiemaszkoLab
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time
import threading
import signal
from pathlib import Path

try:
    import requests
except ImportError:
    print("[!] Brak modulu 'requests'. Zainstaluj: pip install requests")
    sys.exit(1)


# ============================================================================
# Konfiguracja
# ============================================================================
DEFAULT_SERVER   = "https://siemaszko.info"
API_ENDPOINT     = "/api/logic-analyzer/push"
DEFAULT_RATE     = "1M"       # 1 MHz domyslnie (bezpieczne dla power sequence)
DEFAULT_CHANNELS = 8          # 8 kanalow
DEFAULT_SAMPLES  = "100K"     # 100K probek na batch
DEFAULT_DRIVER   = "fx2lafw"  # Cypress FX2 (Saleae clone)
PUSH_INTERVAL    = 0.5        # co 0.5s wysylaj batch
SESSION_TOKEN    = None       # generowany automatycznie


# ============================================================================
# Kolorowy output
# ============================================================================
class C:
    CYAN   = "\033[96m"
    GREEN  = "\033[92m"
    YELLOW = "\033[93m"
    RED    = "\033[91m"
    DIM    = "\033[2m"
    RESET  = "\033[0m"
    BOLD   = "\033[1m"


def log(msg, color=C.CYAN):
    print(f"{color}[SLAB]{C.RESET} {msg}")


def log_ok(msg):
    log(msg, C.GREEN)


def log_warn(msg):
    log(msg, C.YELLOW)


def log_err(msg):
    log(msg, C.RED)


# ============================================================================
# Wykrywanie sigrok-cli
# ============================================================================
def find_sigrok_cli():
    """Szuka sigrok-cli w PATH i typowych lokalizacjach."""
    # Sprawdz PATH
    for cmd in ["sigrok-cli", "sigrok-cli.exe"]:
        try:
            result = subprocess.run([cmd, "--version"],
                                    capture_output=True, text=True, timeout=5)
            if result.returncode == 0:
                version = result.stdout.strip().split('\n')[0]
                log_ok(f"Znaleziono: {cmd} ({version})")
                return cmd
        except (FileNotFoundError, subprocess.TimeoutExpired):
            continue

    # Typowe lokalizacje Windows
    if sys.platform == "win32":
        common_paths = [
            Path(os.environ.get("PROGRAMFILES", "C:\\Program Files")) / "sigrok" / "sigrok-cli" / "sigrok-cli.exe",
            Path(os.environ.get("LOCALAPPDATA", "")) / "sigrok" / "sigrok-cli" / "sigrok-cli.exe",
            Path("C:\\sigrok") / "sigrok-cli.exe",
        ]
        for p in common_paths:
            if p.exists():
                log_ok(f"Znaleziono: {p}")
                return str(p)

    return None


def list_devices(sigrok_cmd):
    """Listuje podlaczone urzadzenia sigrok."""
    log("Skanuje urzadzenia...")
    try:
        result = subprocess.run(
            [sigrok_cmd, "--scan"],
            capture_output=True, text=True, timeout=10
        )
        if result.stdout.strip():
            print(result.stdout)
        else:
            log_warn("Nie znaleziono zadnych urzadzen. Sprawdz podlaczenie USB.")
            if sys.platform == "win32":
                log_warn("Na Windows moze byc potrzebny sterownik Zadig (WinUSB).")
                log_warn("Pobierz: https://zadig.akeo.ie/")
    except subprocess.TimeoutExpired:
        log_err("Timeout skanowania — moze brak sterownika.")
    except FileNotFoundError:
        log_err(f"Nie znaleziono {sigrok_cmd}")


# ============================================================================
# Przechwytywanie danych
# ============================================================================
def capture_batch(sigrok_cmd, driver, rate, channels, samples):
    """Przechwytuje batch probek i zwraca jako liste list bitow per kanal."""
    channel_spec = ",".join([f"D{i}" for i in range(channels)])

    cmd = [
        sigrok_cmd,
        "--driver", driver,
        "--config", f"samplerate={rate}",
        "--channels", channel_spec,
        "--samples", str(samples),
        "--output-format", "csv:header=false"
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode != 0:
            stderr = result.stderr.strip()
            if "No devices found" in stderr or "Failed" in stderr:
                return None, f"Brak urzadzenia: {stderr[:200]}"
            return None, f"Blad sigrok: {stderr[:200]}"

        lines = result.stdout.strip().split('\n')
        if not lines:
            return None, "Brak danych"

        # Parse CSV: kazda linia to "0,1,0,1,..." per kanal
        data = [[] for _ in range(channels)]
        for line in lines:
            bits = line.strip().split(',')
            for ch in range(min(channels, len(bits))):
                data[ch].append(int(bits[ch]))

        return data, None

    except subprocess.TimeoutExpired:
        return None, "Timeout przechwytywania"
    except Exception as e:
        return None, str(e)


# ============================================================================
# Wyslanie do serwera
# ============================================================================
def push_to_server(server, token, data, rate, channels, channel_labels=None):
    """Wysyla batch danych na serwer SiemaszkoLab."""
    payload = {
        "token": token,
        "rate": rate,
        "channels": channels,
        "channel_labels": channel_labels or [f"CH{i}" for i in range(channels)],
        "data": data,
        "ts": time.time()
    }

    try:
        resp = requests.post(
            server + API_ENDPOINT,
            json=payload,
            timeout=5,
            headers={"Content-Type": "application/json"}
        )
        if resp.status_code == 200:
            return True, None
        return False, f"HTTP {resp.status_code}: {resp.text[:100]}"
    except requests.exceptions.ConnectionError:
        return False, "Brak polaczenia z serwerem"
    except requests.exceptions.Timeout:
        return False, "Timeout"
    except Exception as e:
        return False, str(e)


# ============================================================================
# Glowna petla
# ============================================================================
def run_continuous(sigrok_cmd, args):
    """Ciagle przechwytywanie i wysylanie."""
    token = args.token or f"slab-{int(time.time())}-{os.getpid()}"
    batch_num = 0
    errors = 0
    max_errors = 10

    log(f"{C.BOLD}SiemaszkoLab Logic Analyzer Bridge{C.RESET}")
    log(f"Serwer:    {args.server}")
    log(f"Token:     {token}")
    log(f"Driver:    {args.driver}")
    log(f"Rate:      {args.rate}")
    log(f"Kanaly:    {args.channels}")
    log(f"Batch:     {args.samples} probek")
    print()
    log_ok("Rozpoczynam przechwytywanie. Ctrl+C aby zatrzymac.")
    print()

    # Channel labels z argumentow
    labels = None
    if args.labels:
        labels = [l.strip() for l in args.labels.split(',')]

    while True:
        try:
            # Przechwyc batch
            data, err = capture_batch(
                sigrok_cmd, args.driver, args.rate, args.channels,
                args.samples
            )

            if err:
                errors += 1
                log_err(f"Blad #{errors}: {err}")
                if errors >= max_errors:
                    log_err(f"Zbyt wiele bledow ({max_errors}). Koncze.")
                    break
                time.sleep(2)
                continue

            errors = 0
            batch_num += 1

            # Wyslij na serwer
            ok, send_err = push_to_server(
                args.server, token, data, args.rate, args.channels, labels
            )

            if ok:
                # Pokaz kroki info o danych
                active = sum(1 for ch in data if any(b == 1 for b in ch))
                total_samples = len(data[0]) if data else 0
                sys.stdout.write(
                    f"\r{C.GREEN}[#{batch_num:04d}]{C.RESET} "
                    f"Wyslano {total_samples} probek | "
                    f"Aktywne kanaly: {active}/{args.channels} | "
                    f"Token: {token[:16]}..."
                )
                sys.stdout.flush()
            else:
                log_warn(f"Blad wysylania: {send_err}")

            time.sleep(args.interval)

        except KeyboardInterrupt:
            print()
            log_ok("Zatrzymano przez uzytkownika.")
            break


# ============================================================================
# Pojedyncze przechwycenie (snapshot)
# ============================================================================
def run_single(sigrok_cmd, args):
    """Jednokrotne przechwycenie i wyslanie."""
    token = args.token or f"slab-{int(time.time())}-{os.getpid()}"

    log(f"{C.BOLD}Pojedyncze przechwycenie{C.RESET}")
    log(f"Rate: {args.rate} | Kanaly: {args.channels} | Probki: {args.samples}")

    data, err = capture_batch(
        sigrok_cmd, args.driver, args.rate, args.channels, args.samples
    )

    if err:
        log_err(f"Blad: {err}")
        return

    labels = None
    if args.labels:
        labels = [l.strip() for l in args.labels.split(',')]

    ok, send_err = push_to_server(
        args.server, token, data, args.rate, args.channels, labels
    )

    if ok:
        log_ok(f"Wyslano! Token: {token}")
        log_ok(f"Otworz: {args.server}/logic-analyzer  i wpisz token: {token}")
    else:
        log_err(f"Blad wysylania: {send_err}")

    # Pokaz lokalnie tez
    print()
    for ch_idx, ch_data in enumerate(data):
        name = labels[ch_idx] if labels and ch_idx < len(labels) else f"CH{ch_idx}"
        ones = sum(ch_data)
        total = len(ch_data)
        pct = (ones / total * 100) if total > 0 else 0
        bar_len = 40
        bar_fill = int(pct / 100 * bar_len)
        bar = f"{'█' * bar_fill}{'░' * (bar_len - bar_fill)}"
        print(f"  {name:>8s}: [{bar}] HIGH {pct:5.1f}% ({ones}/{total})")


# ============================================================================
# Tryb demo (bez analizatora)
# ============================================================================
def run_demo(args):
    """Generuje przykladowe dane i wysyla na serwer (do testowania bez sprzetu)."""
    import random
    token = args.token or f"demo-{int(time.time())}"

    log(f"{C.BOLD}Tryb DEMO (symulowane dane){C.RESET}")
    log(f"Token: {token}")
    log_ok("Generuje symulowane przebiegi SPI...")

    batch_num = 0
    while True:
        try:
            samples = int(args.samples.replace('K', '000').replace('M', '000000'))
            data = []

            # CH0: CS# — aktywny niski, cyklicznie
            cs = []
            for i in range(samples):
                block = (i // 64) % 3
                cs.append(0 if block < 2 else 1)
            data.append(cs)

            # CH1: CLK — zegar gdy CS# aktywny
            clk = []
            for i in range(samples):
                if cs[i] == 0:
                    clk.append((i // 2) % 2)
                else:
                    clk.append(0)
            data.append(clk)

            # CH2: MOSI — dane pseudolosowe gdy CS# aktywny
            mosi = []
            for i in range(samples):
                if cs[i] == 0 and clk[i] == 1:
                    mosi.append(random.randint(0, 1))
                else:
                    mosi.append(0 if cs[i] == 0 else 0)
            data.append(mosi)

            # CH3: MISO — odpowiedz pseudolosowa
            miso = []
            for i in range(samples):
                if cs[i] == 0 and clk[i] == 0 and i > 32:
                    miso.append(random.randint(0, 1))
                else:
                    miso.append(0)
            data.append(miso)

            # Dodatkowe kanaly (nieaktywne)
            for ch in range(4, args.channels):
                data.append([0] * samples)

            labels = ["CS#", "CLK", "MOSI", "MISO"] + [f"CH{i}" for i in range(4, args.channels)]

            ok, err = push_to_server(
                args.server, token, data, args.rate, min(args.channels, len(data)), labels
            )

            batch_num += 1
            if ok:
                sys.stdout.write(
                    f"\r{C.GREEN}[DEMO #{batch_num:04d}]{C.RESET} "
                    f"Wyslano {samples} probek | Token: {token}"
                )
                sys.stdout.flush()
            else:
                log_warn(f"Blad: {err}")

            time.sleep(args.interval)

        except KeyboardInterrupt:
            print()
            log_ok("Demo zatrzymane.")
            break


# ============================================================================
# Main
# ============================================================================
def main():
    parser = argparse.ArgumentParser(
        description="SiemaszkoLab Logic Analyzer Bridge — wysyla dane z analizatora USB na strone",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Przyklady:
  %(prog)s --demo                          Tryb demo (bez analizatora)
  %(prog)s --rate 8M --channels 4          4 kanaly, 8 MHz
  %(prog)s --labels "CLK,MOSI,MISO,CS#"   Nazwij kanaly
  %(prog)s --list-devices                  Pokaz podlaczone urzadzenia
  %(prog)s --single                        Jednorazowe przechwycenie
        """
    )

    parser.add_argument("--server", default=DEFAULT_SERVER,
                        help=f"URL serwera (domyslnie: {DEFAULT_SERVER})")
    parser.add_argument("--token", default=None,
                        help="Token sesji (generowany automatycznie)")
    parser.add_argument("--driver", default=DEFAULT_DRIVER,
                        help=f"Sterownik sigrok (domyslnie: {DEFAULT_DRIVER})")
    parser.add_argument("--rate", default=DEFAULT_RATE,
                        help=f"Czestotliwosc probkowania (domyslnie: {DEFAULT_RATE})")
    parser.add_argument("--channels", type=int, default=DEFAULT_CHANNELS,
                        help=f"Liczba kanalow (domyslnie: {DEFAULT_CHANNELS})")
    parser.add_argument("--samples", default=DEFAULT_SAMPLES,
                        help=f"Probki na batch (domyslnie: {DEFAULT_SAMPLES})")
    parser.add_argument("--labels", default=None,
                        help='Nazwy kanalow, rozdzielone przecinkami (np. "CLK,MOSI,MISO,CS#")')
    parser.add_argument("--interval", type=float, default=PUSH_INTERVAL,
                        help=f"Interwal miedzy batchami w sekundach (domyslnie: {PUSH_INTERVAL})")
    parser.add_argument("--list-devices", action="store_true",
                        help="Pokaz podlaczone urzadzenia i wyjdz")
    parser.add_argument("--single", action="store_true",
                        help="Jednorazowe przechwycenie (nie ciagle)")
    parser.add_argument("--demo", action="store_true",
                        help="Tryb demo — symulowane dane bez analizatora")

    args = parser.parse_args()

    print()
    print(f"{C.CYAN}{'='*60}{C.RESET}")
    print(f"{C.BOLD}  SiemaszkoLab Logic Analyzer Bridge{C.RESET}")
    print(f"{C.DIM}  Analizator USB → sigrok-cli → SiemaszkoLab{C.RESET}")
    print(f"{C.CYAN}{'='*60}{C.RESET}")
    print()

    # Tryb demo — nie wymaga sigrok
    if args.demo:
        run_demo(args)
        return

    # Szukaj sigrok-cli
    sigrok_cmd = find_sigrok_cli()
    if not sigrok_cmd:
        log_err("Nie znaleziono sigrok-cli!")
        log_err("Pobierz ze strony: https://sigrok.org/wiki/Downloads")
        if sys.platform == "win32":
            log_err("Na Windows pobierz instalator 'sigrok NSIS installer'")
            log_err("i zainstaluj sterownik Zadig: https://zadig.akeo.ie/")
        sys.exit(1)

    # Lista urzadzen
    if args.list_devices:
        list_devices(sigrok_cmd)
        return

    # Ciagle lub pojedyncze przechwytywanie
    if args.single:
        run_single(sigrok_cmd, args)
    else:
        run_continuous(sigrok_cmd, args)


if __name__ == "__main__":
    main()