"""Raadio muusikasalvesti: SDR++ IQ Exporter (baseband, TCP) -> FM stereo demod + RDS -> lood eraldi WAV-failideks.

Loo piir = RadioText muutus (RDS grupp 2A). Failid: music/<jaam>/<kuupäev kell> <RT>.wav + logi songs.csv.
Salvestab ainult siis, kui SDR++ keskmine sagedus (rigctl 'f') on jaamast < 1.1 MHz; muidu ootab.

Käivitus: python music_recorder.py [--station 101.6] [--name "Raadio 2"] [--port 1234]
Peatamine: Ctrl+C või music.ps1 stop (kirjutab pooleli oleva loo lõpuni).
"""
import argparse, csv, datetime as dt, os, re, socket, sys, threading, time, wave
from collections import deque
from pathlib import Path

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import iq as iqlib

ROOT = Path(__file__).resolve().parent.parent
BLOCK_S = 5.0                # töödeldav plokk
RDS_WINDOW_S = 10.0          # RDS dekodeeritakse viimase 10 s pealt
AUDIO_FS = 48000
MIN_SONG_S = 45              # lühemad lähevad _short kausta


def rigctl(cmd, port=4532):
    try:
        s = socket.create_connection(("127.0.0.1", port), timeout=1.0)
        s.sendall((cmd + "\n").encode()); time.sleep(0.15)
        s.settimeout(0.5)
        data = b""
        try:
            while True:
                chunk = s.recv(4096)
                if not chunk: break
                data += chunk
                if b"\n" in data: break
        except socket.timeout:
            pass
        s.close()
        return data.decode(errors="ignore").strip().splitlines()
    except OSError:
        return None


def current_center():
    r = rigctl("f")
    try:
        return int(r[0])
    except (TypeError, ValueError, IndexError):
        return None


class SongWriter:
    def __init__(self, outdir, station, log_path):
        self.outdir, self.station, self.log_path = outdir, station, log_path
        self.w = None; self.path = None; self.rt = None; self.start = None; self.frames = 0

    def open(self, rt):
        self.close()
        safe = re.sub(r'[\\/:*?"<>|]+', "_", rt).strip() or "tundmatu"
        stamp = dt.datetime.now().strftime("%Y-%m-%d %H-%M")
        self.path = self.outdir / f"{stamp} {safe}.wav"
        self.w = wave.open(str(self.path), "wb")
        self.w.setnchannels(2); self.w.setsampwidth(2); self.w.setframerate(AUDIO_FS)
        self.rt, self.start, self.frames = rt, dt.datetime.now(), 0
        print(f"[{stamp}] ALGAB  {rt}", flush=True)

    def write(self, stereo):
        if self.w is None:
            return
        self.w.writeframes((np.clip(stereo, -1, 1) * 32767).astype(np.int16).tobytes())
        self.frames += len(stereo)

    def close(self):
        if self.w is None:
            return
        self.w.close()
        dur = self.frames / AUDIO_FS
        final = self.path
        if dur < MIN_SONG_S:
            short = self.outdir / "_short"; short.mkdir(exist_ok=True)
            final = short / self.path.name
            os.replace(self.path, final)
        new = not self.log_path.exists()
        with open(self.log_path, "a", newline="", encoding="utf-8") as f:
            wr = csv.writer(f)
            if new: wr.writerow(["algus", "kestus_s", "jaam", "radiotext", "fail"])
            wr.writerow([self.start.strftime("%Y-%m-%d %H:%M:%S"), f"{dur:.0f}", self.station, self.rt, final.name])
        print(f"[{dt.datetime.now():%H:%M}] LÕPP   {self.rt}  ({dur:.0f} s) -> {final.name}", flush=True)
        self.w = None; self.path = None


def stereo_from_mpx(mpx, fs):
    """L+R (0-15k) ja L-R (38 kHz DSB, faas piloodist) -> (n, 2) float 48 kHz."""
    n = len(mpx)
    pil = iqlib.fft_bandpass_analytic(mpx, fs, 18900, 19100)
    c38 = (pil / (np.abs(pil) + 1e-9)) ** 2
    sub = iqlib.fft_bandpass_analytic(mpx, fs, 23000, 53000)
    lmr = np.real(sub * np.conj(c38))                        # L-R
    lpr, afs = iqlib.fft_lowpass_real_decimate(mpx, fs, 15e3, AUDIO_FS)
    lmr, _ = iqlib.fft_lowpass_real_decimate(lmr, fs, 15e3, AUDIO_FS)
    m = min(len(lpr), len(lmr))
    L, R = (lpr[:m] + lmr[:m]), (lpr[:m] - lmr[:m])
    out = np.stack([L, R], axis=1)
    # de-emphasis 50 us
    A = np.fft.rfft(out, axis=0)
    f = np.fft.rfftfreq(len(out), 1 / afs)[:, None]
    A /= (1 + 2j * np.pi * f * 50e-6)
    out = np.fft.irfft(A, len(out), axis=0) / 75e3
    pilot_ok = np.mean(np.abs(pil) ** 2) > 1e3               # stereo-pilot olemas?
    if not pilot_ok:
        out[:, 0] = out[:, 1] = (out[:, 0] + out[:, 1]) / 2
    return out.astype(np.float32)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--station", type=float, default=101.6)
    ap.add_argument("--name", default="Raadio 2")
    ap.add_argument("--port", type=int, default=1234)
    ap.add_argument("--fs", type=int, default=2400000)
    a = ap.parse_args()
    f_st = int(round(a.station * 1e6))
    outdir = ROOT / "music" / a.name
    outdir.mkdir(parents=True, exist_ok=True)
    writer = SongWriter(outdir, a.name, outdir / "songs.csv")

    fs = a.fs
    block_n = int(fs * BLOCK_S)
    stop = threading.Event()
    buf = deque()                                              # sissetulevad toorplokid (bytes)
    buf_lock = threading.Lock()

    def reader():
        while not stop.is_set():
            try:
                s = socket.create_connection(("127.0.0.1", a.port), timeout=5)
                s.settimeout(5)
                print(f"IQ Exporter ühendatud (port {a.port})", flush=True)
                need = block_n * 4                             # int16 I+Q
                while not stop.is_set():
                    blk = bytearray(need)                      # üks plokk, täidetakse otse (ilma kopeerimata)
                    view = memoryview(blk); got = 0
                    while got < need:
                        n = s.recv_into(view[got:], need - got)
                        if not n:
                            raise OSError("ühendus katkes")
                        got += n
                    with buf_lock:
                        buf.append(blk)
                        if len(buf) > 6:                       # töötlus jääb maha -> viska vanimad
                            buf.popleft(); print('HOIATUS: töötlus jääb maha, plokk visati ära', flush=True)
            except OSError as e:
                print(f"IQ Exporter: {e} – proovin 5 s pärast uuesti", flush=True)
                time.sleep(5)

    threading.Thread(target=reader, daemon=True).start()

    mpx_hist = deque(maxlen=int(RDS_WINDOW_S / BLOCK_S))
    pending_audio = deque()                                    # 15 s viivitus: RT tuleb loo algusest hiljem
    DELAY_BLOCKS = 3
    last_rt, stable_rt, rt_count = None, None, 0
    center = None; last_center_check = 0; nblk = 0
    print(f"salvestan {a.name} ({a.station} MHz) -> {outdir}", flush=True)
    try:
        while True:
            if time.time() - last_center_check > 3:
                c = current_center(); last_center_check = time.time()
                if c != center:
                    center = c
                    print(f"[{dt.datetime.now():%H:%M}] keskmine sagedus {center}", flush=True)
            with buf_lock:
                raw = buf.popleft() if buf else None
            if raw is None:
                time.sleep(0.2); continue
            if center is None or abs(center - f_st) > 1.1e6:
                while pending_audio: writer.write(pending_audio.popleft())
                if writer.w: writer.close()
                mpx_hist.clear(); stable_rt = None; rt_count = 0
                continue                                       # dongel on mujal
            t0 = time.time()
            d = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
            iqd = (d[0::2] + 1j * d[1::2]).astype(np.complex64)
            t1 = time.time()
            x, fs2 = iqlib.fft_decimate_complex(iqd, fs, f_st - center, 240e3)
            t2 = time.time()
            mpx = iqlib.fm_demod(x, fs2)
            mpx_hist.append(mpx)
            stereo = stereo_from_mpx(mpx, fs2)
            t3 = time.time()

            # RDS viimase 10 s pealt
            rt = None
            nblk += 1
            if len(mpx_hist) == mpx_hist.maxlen and nblk % 2 == 0:   # RDS iga 10 s, mitte iga 5 s
                r = iqlib.rds_decode(np.concatenate(mpx_hist), fs2)
                if r.get("ok") and r["rt"] and "?" not in r["rt"]:
                    rt = r["rt"].strip()
            if rt:
                if rt == last_rt:
                    rt_count += 1
                else:
                    last_rt, rt_count = rt, 1
                if rt_count >= 2 and rt != stable_rt:          # kaks järjestikust sama RT -> kindel
                    stable_rt = rt
                    writer.open(stable_rt)                     # vana fail kinni; puhvris olev 15 s läheb uude
            if time.time() - t0 > BLOCK_S * 0.8: print(f'  plokk {time.time()-t0:.1f} s (limiit {BLOCK_S} s): konv {t1-t0:.1f} decim {t2-t1:.1f} stereo {t3-t2:.1f} rds {time.time()-t3:.1f}, puhvris {len(buf)}', flush=True)
            pending_audio.append(stereo)
            while len(pending_audio) > DELAY_BLOCKS:
                blk = pending_audio.popleft()
                if writer.w is not None:
                    writer.write(blk)
                # enne esimest RT-d: heli läheb kaotsi (pole nime, kuhu panna)
    except KeyboardInterrupt:
        pass
    finally:
        stop.set()
        while pending_audio: writer.write(pending_audio.popleft())
        writer.close()
        print("lõpetatud", flush=True)


if __name__ == "__main__":
    main()