"""
FID Sequence tab - PySide6 port of seq_interp/src/fid_gui/gui_RF_adj(FID).py.

Generates a Free Induction Decay (FID) pulse sequence via seqgen_FID,
visualises it with the shared PlotPanel, and emits fid_seq_generated(path)
so LFMRIWindow can hand the .seq off to SeqInterpTab.
"""
from __future__ import annotations

import json
import os
import tempfile
from datetime import datetime
from math import ceil
from types import SimpleNamespace

import numpy as np
from PySide6.QtCore import Qt, QThread, Signal
from PySide6.QtGui import QFont
from PySide6.QtWidgets import (
    QWidget, QSplitter, QVBoxLayout, QHBoxLayout,
    QGroupBox, QFormLayout, QLabel, QDoubleSpinBox, QSpinBox,
    QPushButton, QRadioButton, QButtonGroup, QScrollArea,
    QSizePolicy, QFileDialog, QMessageBox, QProgressBar,
    QFrame,
)

from src.gui.plot_panel import PlotPanel
from src.gui.tr_widgets import TrGroupBox, TrPushButton
from src.gui.workers import SeqInterpHttpWorker
from src.clients.seq_interp_client import SeqInterpClient
from src import i18n, theme

_GAMMA = 42.576e6   # Hz/T, proton
_G_AMP_MAX_MT_M = 5.0
_G_SLEW_MAX_T_M_S = 45.0


# -----------------------------------------------------------------------------
class _FidWorker(QThread):
    """Generates a FID sequence and saves it to a .seq file off the UI thread."""
    finished = Signal(str)   # path to saved .seq
    error = Signal(str)
    log_msg = Signal(str)

    def __init__(self, param: SimpleNamespace, seq_path: str) -> None:
        super().__init__()
        self._param = param
        self._seq_path = seq_path

    def run(self) -> None:
        try:
            from src.fid.seqgen_FID import seqgen_FID
            self.log_msg.emit("Generating FID sequence...")
            seq = seqgen_FID(self._param)
            seq.write(self._seq_path)
            self.log_msg.emit(f"Sequence saved: {self._seq_path}")
            self.finished.emit(self._seq_path)
        except Exception as exc:
            self.error.emit(f"{type(exc).__name__}: {exc}")


# -----------------------------------------------------------------------------
class FidTab(QWidget):
    """FID pulse-sequence editor and visualiser."""

    fid_seq_generated = Signal(str)   # emitted when user clicks "Load in Sequence Tab"

    # -- defaults matching original Tkinter GUI ----------------------------
    _D = {
        "BW_per_point": 100.0,   # Hz
        "N_point":       512,
        "TR":            200.0,  # ms  (converted to s when building param)
        "NA":            1,
        "freq_offset":   0.0,    # Hz
        "t_ex":          1.0,    # ms  (converted to s)
        "rf_adc_delay":  0.1,    # ms  (rf_ringdown_time, converted to s via *1e-3)
        "scale_rf":      1.0,
    }

    def __init__(
        self,
        hw_config_path: str | None = None,
        output_dir: str | None = None,
        seq_interp_url: str = "http://localhost:7475",
        parent: QWidget | None = None,
    ) -> None:
        super().__init__(parent)
        self._hw_config_path = hw_config_path
        self._output_dir = output_dir
        self._last_seq_path: str | None = None
        self._last_param: SimpleNamespace | None = None
        self._gen_worker: QThread | None = None
        self._interp_worker: QThread | None = None
        # Interpretation runs in the seq_interp service (GUI is a pure client).
        self._seq_interp_client = SeqInterpClient(seq_interp_url)

        root = QVBoxLayout(self)
        root.setContentsMargins(0, 0, 0, 0)
        root.setSpacing(0)
        root.addWidget(self._build_toolbar())
        root.addWidget(self._build_splitter(), stretch=1)
        root.addWidget(self._build_status_bar())

    # ================================================================== #
    #  Public API                                                          #
    # ================================================================== #

    def set_hw_config(self, path: str) -> None:
        self._hw_config_path = path

    def set_output_dir(self, path: str) -> None:
        self._output_dir = path

    # ================================================================== #
    #  Toolbar                                                             #
    # ================================================================== #

    def apply_theme(self) -> None:
        p = theme.palette()
        self._toolbar.setStyleSheet(
            f"#FidToolBar {{ background: {p['surface']}; border-bottom: 1px solid {p['border2']}; }}"
        )
        if hasattr(self, "_plots"):
            self._plots.apply_theme()

    def retranslate_ui(self) -> None:
        # All buttons/groups auto-translate via TrPushButton/TrGroupBox.
        # Only the status label needs manual handling (shows runtime messages).
        if self._btn_generate.isEnabled() and not self._progress.isVisible():
            self._status_lbl.setText(i18n.tr("fid_status_init"))

    def _build_toolbar(self) -> QWidget:
        bar = QWidget()
        bar.setObjectName("FidToolBar")
        p = theme.palette()
        bar.setStyleSheet(
            f"#FidToolBar {{ background: {p['surface']}; border-bottom: 1px solid {p['border2']}; }}"
        )
        lay = QHBoxLayout(bar)
        lay.setContentsMargins(6, 4, 6, 4)
        lay.setSpacing(4)

        def btn(key: str, tip: str, slot=None, enabled: bool = True) -> TrPushButton:
            b = TrPushButton(key)
            b.setToolTip(tip)
            b.setEnabled(enabled)
            b.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Fixed)
            if slot:
                b.clicked.connect(slot)
            lay.addWidget(b)
            return b

        def sep() -> QFrame:
            f = QFrame()
            f.setFrameShape(QFrame.VLine)
            f.setFrameShadow(QFrame.Sunken)
            f.setFixedWidth(2)
            return f

        self._btn_generate    = btn("btn_generate",    "Generate FID sequence and visualise", self._generate)
        lay.addWidget(sep())
        self._btn_save        = btn("btn_save_seq",    "Save .seq + .json to output directory",
                                    self._save, enabled=False)
        self._btn_load_seq_tab = btn("btn_load_seq_tab",
                                     "Send this .seq to the Sequence interpreter tab",
                                     self._load_in_seq_tab, enabled=False)
        lay.addStretch()

        self._progress = QProgressBar()
        self._progress.setRange(0, 0)
        self._progress.setFixedWidth(120)
        self._progress.setVisible(False)
        lay.addWidget(self._progress)

        self._toolbar = bar
        return bar

    # ================================================================== #
    #  Main splitter                                                        #
    # ================================================================== #

    def _build_splitter(self) -> QSplitter:
        split = QSplitter(Qt.Horizontal)
        split.addWidget(self._build_left_panel())

        self._plots = PlotPanel()
        split.addWidget(self._plots)
        split.setSizes([320, 900])
        return split

    def _build_left_panel(self) -> QScrollArea:
        inner = QWidget()
        inner.setMinimumWidth(240)
        inner.setMaximumWidth(400)
        lay = QVBoxLayout(inner)
        lay.setContentsMargins(6, 6, 6, 6)
        lay.setSpacing(8)

        # -- parameter form --------------------------------------------
        self._param_grp = TrGroupBox("grp_fid_params")
        param_grp = self._param_grp
        form = QFormLayout(param_grp)
        form.setFieldGrowthPolicy(QFormLayout.AllNonFixedFieldsGrow)

        def dbl(val, mn, mx, step, dec=3) -> QDoubleSpinBox:
            sb = QDoubleSpinBox()
            sb.setRange(mn, mx)
            sb.setSingleStep(step)
            sb.setDecimals(dec)
            sb.setValue(val)
            return sb

        def integer(val, mn, mx) -> QSpinBox:
            sb = QSpinBox()
            sb.setRange(mn, mx)
            sb.setValue(val)
            return sb

        d = self._D
        self._sb_bw   = dbl(d["BW_per_point"], 1.0, 1e6, 10.0, 1)
        self._sb_npt  = integer(d["N_point"], 1, 65536)
        self._lbl_bwfull = QLabel("-")
        self._lbl_bwfull.setFont(QFont("Courier New", 9))
        self._sb_tr   = dbl(d["TR"], 1.0, 1e4, 10.0, 1)
        self._sb_na   = integer(d["NA"], 1, 65536)
        self._sb_foff = dbl(d["freq_offset"], -1e6, 1e6, 1.0, 1)
        self._sb_tex  = dbl(d["t_ex"], 0.001, 100.0, 0.1, 3)
        self._sb_rfadc = dbl(d["rf_adc_delay"], 0.0, 100.0, 0.01, 3)
        self._sb_scale = dbl(d["scale_rf"], 0.0, 10.0, 0.1, 3)

        form.addRow("BW per point (Hz):", self._sb_bw)
        form.addRow("N points:",           self._sb_npt)
        form.addRow("BW full (Hz):",       self._lbl_bwfull)
        form.addRow("TR (ms):",            self._sb_tr)
        form.addRow("NA (averages):",      self._sb_na)
        form.addRow("Freq offset (Hz):",   self._sb_foff)
        form.addRow("RF pulse dur (ms):",  self._sb_tex)
        form.addRow("RF-ADC delay (ms):",  self._sb_rfadc)
        form.addRow("RF scale:",           self._sb_scale)
        lay.addWidget(param_grp)

        # update BW full on change
        self._sb_bw.valueChanged.connect(self._update_bwfull)
        self._sb_npt.valueChanged.connect(self._update_bwfull)
        self._update_bwfull()

        # -- pulse type ------------------------------------------------
        self._pulse_grp = TrGroupBox("grp_pulse_type")
        pulse_grp = self._pulse_grp
        pulse_lay = QHBoxLayout(pulse_grp)
        self._rb_block = QRadioButton("block")
        self._rb_sinc  = QRadioButton("sinc")
        self._rb_block.setChecked(True)
        pulse_btn_grp = QButtonGroup(self)
        pulse_btn_grp.addButton(self._rb_block)
        pulse_btn_grp.addButton(self._rb_sinc)
        pulse_lay.addWidget(self._rb_block)
        pulse_lay.addWidget(self._rb_sinc)
        lay.addWidget(pulse_grp)

        # -- output filename -------------------------------------------
        from PySide6.QtWidgets import QLineEdit
        self._file_grp = TrGroupBox("grp_out_filename")
        file_grp = self._file_grp
        file_lay = QVBoxLayout(file_grp)
        self._le_filename = QLineEdit("FID_" + datetime.now().strftime("%d%m%y_%H%M"))
        file_lay.addWidget(self._le_filename)
        lay.addWidget(file_grp)

        lay.addStretch()

        scroll = QScrollArea()
        scroll.setWidget(inner)
        scroll.setWidgetResizable(True)
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        return scroll

    def _build_status_bar(self) -> QFrame:
        bar = QFrame()
        bar.setFrameShape(QFrame.StyledPanel)
        bar.setFixedHeight(22)
        lay = QHBoxLayout(bar)
        lay.setContentsMargins(6, 0, 6, 0)
        self._status_lbl = QLabel(i18n.tr("fid_status_init"))
        self._status_lbl.setFont(QFont("Arial", 8))
        lay.addWidget(self._status_lbl)
        lay.addStretch()
        return bar

    # ================================================================== #
    #  Helpers                                                             #
    # ================================================================== #

    def _update_bwfull(self) -> None:
        bw_full = self._sb_bw.value() * self._sb_npt.value()
        self._lbl_bwfull.setText(f"{ceil(bw_full * 1000) / 1000:.1f}")

    def _build_param(self) -> SimpleNamespace:
        gamma = _GAMMA
        G_amp_max = _G_AMP_MAX_MT_M * 1e-3 * gamma
        G_slew_max = _G_SLEW_MAX_T_M_S * gamma
        rf_raster_time = 1e-6
        grad_raster_time = 1e-6

        p = SimpleNamespace()
        p.G_amp_max = G_amp_max
        p.G_slew_max = G_slew_max
        p.gamma = gamma
        p.grad_raster_time = grad_raster_time
        p.rf_raster_time = rf_raster_time
        p.rf_dead_time = 0.0
        p.adc_dead_time = 0.0

        p.BW_per_point = self._sb_bw.value()
        p.N_point = float(self._sb_npt.value())
        p.BW_full = p.BW_per_point * p.N_point
        p.TR = self._sb_tr.value() * 1e-3
        p.average = float(self._sb_na.value())
        p.freq_offset = self._sb_foff.value()
        p.t_ex = self._sb_tex.value() * 1e-3
        p.rf_ringdown_time = self._sb_rfadc.value() * 1e-3
        p.scale_rf = self._sb_scale.value()
        p.is_sinc_pulse = self._rb_sinc.isChecked()

        p.flip_angle = 90.0
        p.apodization = 0.3
        p.t_BW_product_ex = 2.18
        return p

    def _tmp_seq_path(self) -> str:
        tmp = tempfile.gettempdir()
        name = self._le_filename.text().strip() or "fid_tmp"
        return os.path.join(tmp, f"{name}.seq")

    # ================================================================== #
    #  Actions                                                             #
    # ================================================================== #

    def _generate(self) -> None:
        param = self._build_param()
        seq_path = self._tmp_seq_path()
        self._last_param = param
        self._last_seq_path = seq_path
        self._set_busy(True, "Generating sequence...")

        self._gen_worker = _FidWorker(param, seq_path)
        self._gen_worker.log_msg.connect(self._status)
        self._gen_worker.finished.connect(self._on_generated)
        self._gen_worker.error.connect(self._on_error)
        self._gen_worker.start()

    def _on_generated(self, seq_path: str) -> None:
        self._status(f"Generated -> {seq_path}  |  Interpreting...")
        if not self._seq_interp_client.healthcheck():
            self._on_error(
                "seq_interp service is unavailable. Start the service and retry."
            )
            return
        self._interp_worker = SeqInterpHttpWorker(
            self._seq_interp_client, seq_path
        )
        self._interp_worker.progress.connect(self._status)
        self._interp_worker.finished.connect(self._on_interp_done)
        self._interp_worker.error.connect(self._on_error)
        self._interp_worker.start()

    def _on_interp_done(self, result: dict) -> None:
        self._set_busy(False)
        waveforms = result.get("waveforms", {})
        if waveforms:
            self._plots.plot_from_waveforms(waveforms)
        self._btn_save.setEnabled(True)
        self._btn_load_seq_tab.setEnabled(True)
        nb = result.get("metadata", {}).get("sync_block_count", "?")
        self._status(
            f"Done - {nb} sync blocks  |  "
            f"Click 'Save' to write .seq or '-> Load in Sequence Tab'"
        )

    def _save(self) -> None:
        if self._last_param is None:
            return
        out_dir = self._output_dir
        if not out_dir:
            out_dir = QFileDialog.getExistingDirectory(self, "Choose output directory")
            if not out_dir:
                return
            self._output_dir = out_dir

        name = self._le_filename.text().strip() or "fid_output"
        seq_path = os.path.join(out_dir, f"{name}.seq")
        json_path = os.path.join(out_dir, f"{name}.json")

        self._set_busy(True, "Saving...")
        worker = _FidWorker(self._last_param, seq_path)
        worker.log_msg.connect(self._status)
        worker.finished.connect(lambda p: self._on_saved(p, json_path))
        worker.error.connect(self._on_error)
        worker.start()
        self._gen_worker = worker

    def _on_saved(self, seq_path: str, json_path: str) -> None:
        try:
            with open(json_path, "w") as f:
                json.dump(vars(self._last_param), f, indent=4, default=str)
        except Exception as exc:
            self._on_error(f"JSON save failed: {exc}")
            return
        self._last_seq_path = seq_path
        self._set_busy(False)
        self._status(f"Saved -> {os.path.dirname(seq_path)}")
        QMessageBox.information(
            self, "Saved",
            f".seq and .json written to:\n{os.path.dirname(seq_path)}"
        )

    def _load_in_seq_tab(self) -> None:
        if self._last_seq_path and os.path.isfile(self._last_seq_path):
            self.fid_seq_generated.emit(self._last_seq_path)
        else:
            QMessageBox.warning(self, "No sequence",
                                "Generate and save a sequence first.")

    def _on_error(self, msg: str) -> None:
        self._set_busy(False)
        self._status(f"ERROR: {msg}")
        QMessageBox.critical(self, "Error", msg)

    # ================================================================== #
    #  UI state helpers                                                    #
    # ================================================================== #

    def _set_busy(self, busy: bool, tip: str = "") -> None:
        self._progress.setVisible(busy)
        self._btn_generate.setEnabled(not busy)
        if tip:
            self._status(tip)

    def _status(self, msg: str) -> None:
        self._status_lbl.setText(msg)
