import base64
import io
import itertools
import json
import os
import random
import unicodedata
import uuid
import gradio as gr
from collections import defaultdict
import soundfile as sf
import numpy as np
import librosa
import librosa.display
from matplotlib.figure import Figure
from matplotlib.ticker import FuncFormatter
import matplotlib.patheffects as pe

# this builds exercises specifically to the structure of the original book

JSON_REL = os.path.join("skynjun_shsh", "skynjun_shsh_AlignDict__HF-MANUAL.json")
PAD_S = 0.5  # padding added to each side of an item's start/end - best to have a lot
LENGDIR = ["langt", "stutt"]
LENGD_SEL = [("löng", "langt"), ("stutt", "stutt")] # is this grammer idk at all.....



# nicer: hint printed either side of textbox which user fills in -- formatting maybe weird tho
# txtbox: backup plan - prefill the textbox with the hint e.g. b__la, user edits underscores
HINT_STYLE = "nicer"

# for nicer hints -
#  assuming 3.2em is good for abt 2 underscores & 9em around 5-10 chars, maybe adjust
BASE_CSS = """
.gradio-container { max-width: 1250px !important; }
.ex-row { gap: 6px !important; align-items: center; border-bottom: 1px solid var(--border-color-primary); padding: 2px 0; flex-wrap: nowrap !important; }
.ex-row .block { padding: 2px !important; }
.ex-num p { font-size: 1.2em; margin: 0; text-align: center; }
.play-btn { min-width: 44px !important; max-width: 52px; height: 40px; padding: 0 !important; font-size: 1.2em; }
.radio-row { gap: 12px !important; flex-wrap: nowrap !important; }
.fb p { margin: 0; line-height: 1.35; }
.hint-wrap { display: flex !important; flex-direction: row !important; flex-wrap: nowrap !important;
             align-items: center; justify-content: flex-start !important; gap: 0 !important;
             width: max-content !important; max-width: 100%; min-width: 0 !important; padding: 0 !important; }
.hint-wrap::before, .hint-wrap::after { font-size: 1.3em; white-space: nowrap; flex: 0 0 auto;
             margin: 0 !important; padding: 0 !important; }
.hint-wrap > .hint-box { flex: 0 0 auto !important; width: 3.2em !important; min-width: 0 !important;
             margin: 0 2px !important; padding: 0 !important; }
.hint-box input, .hint-box textarea { text-align: center; }
.plain-box { flex: 0 0 auto !important; width: 9em !important; min-width: 0 !important; }
.hint-box input, .plain-box input, .hint-box textarea, .plain-box textarea {
    padding: 4px 6px !important; min-height: 0 !important; height: 2em !important; line-height: 1.2 !important; }
.chapter-head h2 { font-size: 2em; margin: 0.2em 0 0.4em 0; }
"""


# without manual correction 18 tokens fail this
# (more are aligned to wrong/bad audio but not caught here)
def has_audio(d: dict) -> bool:
    return d.get("start") is not None and d.get("end") is not None


def load_book(folder: str) -> list[dict]:
    with open(os.path.join(folder, JSON_REL), encoding="utf-8") as f:
        return json.load(f)["bók"]



def kafli_choices(book):
    return [(f"{k['n_kaf']}. kafli", k["n_kaf"]) for k in book]


def aefing_choices(book, n_kaf):
    kaf = next(k for k in book if k["n_kaf"] == n_kaf)
    return [(f"{e['n_æf']}. {e['hljóðskynjun']}", e["n_æf"]) for e in kaf["kafli"]]


def chapter_header(book, n_kaf, n_aef) -> str:
    kaf = next(k for k in book if k["n_kaf"] == n_kaf)
    aef = next(e for e in kaf["kafli"] if e["n_æf"] == n_aef)
    # key name guessed: accepts either spelling
    lys = kaf.get("lýsing") or kaf.get("lysing") or ""
    txt = f"## **{n_kaf}. kafli**"
    if lys:
        txt += f"\n\n### {lys}"
    return txt + f"\n\n## {aef['hljóðskynjun']}"



def sound_choices(book):
    seen = {}
    for k in book:
        for e in k["kafli"]:
            sides = e["hljóðskynjun"].split(" — ")
            for x in e["æfingar"]:
                if len(sides) == len(x["hljóð"]):
                    for sound, side in zip(x["hljóð"], sides):
                        seen.setdefault(sound, side.split(" [")[0].strip())
    return [(f"{orth} [{s}]", s) for s, orth in seen.items()]


def build_pool(book) -> dict:
    pool = defaultdict(list)
    for k in book:
        for e in k["kafli"]:
            for x in e["æfingar"]:
                for d in x["dæmi"]:
                    if has_audio(d):
                        pool[(d["svar"], d["lengd"])].append(d)
    return pool


# make the user pick settings that generate a reasonable number of exercises
def validate_selection(sounds, lengdir) -> str | None:
    ns, nl = len(sounds or []), len(lengdir or [])
    if ns == 0 or nl == 0:
        return "Veldu að minnsta kosti 1 hljóð ásamt lengd(um)!"
    if nl == 2 and ns > 3:
        return f"Þú valdir of mörg hljóð! ({ns})"
    if nl == 1 and not 2 <= ns <= 4:
        return f"Þú valdir of mörg hljóð! ({ns})"
    return None


# if the user picked 3 sounds and both lengths, only generate 2 of each type
#   so they get 12 exercises instead of 18
# but otherwise generate 3 of each type
def items_per_condition(n_sounds: int, n_lengdir: int) -> int:
    return 2 if (n_sounds, n_lengdir) == (3, 2) else 3


# generates randomly -
# user can get a new exercise even for the same settings
def sample_exercises(pool, folder, sounds, lengdir, rng=random):
    k = items_per_condition(len(sounds), len(lengdir)) # usually 3 unless thats too many
    used, chosen, warnings = set(), [], []
    for sound, lengd in itertools.product(sounds, lengdir):
        cands = list(pool.get((sound, lengd), []))
        rng.shuffle(cands)
        got = 0
        for d in cands:
            w = _norm(d["orð"])
            if w not in used: # ensure unique word set, not just unique tokens
                used.add(w); chosen.append(d); got += 1
                if got == k:
                    break
        if got < k: # in case missing datas
            warnings.append(f"fann aðeins {got} / {k} orð fyrir [{sound}] * {lengd}")
    rng.shuffle(chosen)
    exercises = [make_exercise(d, folder) for d in chosen]
    for ex in exercises:
        ex["sounds"] = list(sounds)
    return exercises, warnings


def make_exercises_for_chapter(book, folder, n_kaf, n_aef, lengd) -> list[dict]:
    kaf = next(k for k in book if k["n_kaf"] == n_kaf)
    aef = next(e for e in kaf["kafli"] if e["n_æf"] == n_aef)
    ex = next(x for x in aef["æfingar"] if x["lengd"] == lengd)
    return [make_exercise(d, folder) for d in ex["dæmi"] if has_audio(d)]



def clip_audio(folder: str, item: dict):
    path = os.path.join(folder, item["audio_path"])
    sr = sf.info(path).samplerate
    start = max(0, int(round((item["start"] - PAD_S) * sr)))
    stop = int(round((item["end"] + PAD_S) * sr)) 
    data, sr = sf.read(path, start=start, stop=stop, dtype="int16")
    return sr, data





SPEC_FMAX = 5000
SPEC_CMAP = "viridis"
SPEC_TOP_DB = 50
TRIM_DB = 30 # silence threshold
SPEC_PEAK_PCT = 99.5 # bc my computer likes to record clicks even when i dont click anything

def _fmt_tick(v, _pos=None) -> str:
    return f"{round(v, 2):g}"  # for tiny graphs
 
def _spec_html(y, sr, num) -> str:
    y = np.asarray(y, dtype=np.float32)
    if y.ndim > 1:
        y = y.mean(axis=1)
    n_fft = 2 ** int(round(np.log2(0.02 * sr)))
    if y.size < n_fft or not np.any(y):
        return ""
    hop = n_fft // 4
    mag = np.abs(librosa.stft(y, n_fft=n_fft, hop_length=hop))
    S = librosa.amplitude_to_db(mag, ref=np.percentile(mag, SPEC_PEAK_PCT), top_db=None)
    fig = Figure(figsize=(3.0, 1.35))  # Figure (not pyplot) => no global state, thread-safe
    ax = fig.subplots()
    librosa.display.specshow(S, sr=sr, hop_length=hop, x_axis="time", y_axis="hz", ax=ax,
                             cmap=SPEC_CMAP, vmin=-SPEC_TOP_DB, vmax=0)
    fmax = min(SPEC_FMAX, sr / 2)
    ax.set_ylim(0, fmax)
    ax.set_yticks(np.arange(0, fmax + 1, 2000))
    ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _p: f"{v / 1000:g}"))
    ax.xaxis.set_major_formatter(FuncFormatter(_fmt_tick))
    ax.set_xlabel("")
    ax.set_ylabel("kHz", fontsize=7, labelpad=1)
    ax.tick_params(labelsize=6, length=2, pad=1)
    ax.text(0.04, 0.94, str(num), transform=ax.transAxes, color="white", fontsize=13,
            fontweight="bold", ha="left", va="top",
            path_effects=[pe.withStroke(linewidth=2, foreground="black")])
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=120, bbox_inches="tight", pad_inches=0.02)
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f'<img src="data:image/png;base64,{b64}" style="display:block;max-width:100%;height:auto;">'
 
 
def sample_spec_html(clip, num) -> str:
    sr, arr = clip
    return _spec_html(arr.astype(np.float32) / 32768.0, sr, num)
 
 
def user_spec_html(path, num) -> str:
    if not path:
        return ""
    try:
        y, sr = librosa.load(path, sr=None, mono=True)
        y, _ = librosa.effects.trim(y, top_db=TRIM_DB)
        return _spec_html(y, sr, num)
    except Exception as e:
        print("spectrogram problem:", e)
        return ""


# hopefully,
# this makes playable audio without a whole gradio audio component for it
# bc those are huge and cant be reduced enough in size.
# ¡¡ not entirely sure will work on all browsers/devices !!
def audio_data_uri(clip) -> str:
    sr, arr = clip
    buf = io.BytesIO()
    sf.write(buf, arr, sr, format="WAV", subtype="PCM_16")
    return "data:audio/wav;base64," + base64.b64encode(buf.getvalue()).decode()


def make_exercise(d: dict, folder: str) -> dict:
    pre, _, post = d["spurning"].partition("__")
    clip = clip_audio(folder, d)
    return dict(
        prefix=pre, suffix=post,
        word=d["orð"], answer=d["svar"], lengd=d["lengd"],
        sounds=list(d["original_target_sounds"]),
        audio_uri=audio_data_uri(clip),
        clip=clip,
    )


def _norm(s: str) -> str:
    return unicodedata.normalize("NFC", (s or "")).strip().casefold()

# check the user's typed word with and without hint format
def word_ok(item: dict, typed: str, hint: bool) -> bool:
    target = _norm(item["word"])
    if _norm(typed) == target:
        return True
    return hint and _norm(item["prefix"] + (typed or "").strip() + item["suffix"]) == target

# to point out non-answers as non-answers rather than incorrect --  
def _word_blank(item: dict, typed: str, hint: bool) -> bool:
    t = (typed or "").strip()
    if not t:
        return True
    return hint and HINT_STYLE != "nicer" and t == (item["prefix"] + "__" + item["suffix"]).strip()



def _cls(s: str) -> str:
    return s.encode("utf-8").hex()


def hint_css(book) -> str:
    pres, sufs = set(), set()
    for k in book:
        for e in k["kafli"]:
            for x in e["æfingar"]:
                for d in x["dæmi"]:
                    p, _, q = d["spurning"].partition("__")
                    if p: pres.add(p)
                    if q: sufs.add(q)
    esc = lambda t: t.replace("\\", "\\\\").replace('"', '\\"')
    rules = [f'.hint-wrap.pre-{_cls(p)}::before {{ content: "{esc(p)}"; }}' for p in pres]
    rules += [f'.hint-wrap.suf-{_cls(q)}::after {{ content: "{esc(q)}"; }}' for q in sufs]
    return "\n".join(rules)



# to listen without gradio audio component, hopefully it works on most devices...
def _play_js(uri: str) -> str:
    key = uuid.uuid4().hex
    return f"""() => {{
        const w = window; w.__clips = w.__clips || {{}};
        for (const [k, a] of Object.entries(w.__clips)) if (k !== '{key}') a.pause();
        let a = w.__clips['{key}'];
        if (!a) {{ a = new Audio('{uri}'); w.__clips['{key}'] = a; }}
        if (a.paused) {{ a.currentTime = 0; a.play(); }} else {{ a.pause(); }}
    }}"""

 
def new_exercise_set(items, hint, ask_lengd=False, ask_sound=True, header=None, spectro=False) -> dict:
    return {"items": items, "hint": bool(hint), "ask_lengd": ask_lengd,
            "ask_sound": ask_sound, "gen": uuid.uuid4().hex, "header": header,
            "spectro": bool(spectro)}


# to refresh exercises, destroy all current widgets then make new ones
#  (otherwise if the user recorded any audio that tends to be sticky)
def wire_generate(btn, state, inputs, make, extra_outputs=(), guard=None):
    def blank(*args):
        return gr.skip() if (guard and guard(*args)) else None
    btn.click(blank, inputs, state).then(make, inputs, [state, *extra_outputs])


def exercise_panel() -> gr.State:
    """Create the exercise area (call inside a tab). Returns the State that drives it."""
    state = gr.State(None)

    @gr.render(inputs=state)
    def show(st):
        if not st:
            gr.Markdown("Veldu hljóð og smelltu á **Byrja**")
            return
        items, use_hint = st["items"], st["hint"]
        ask_lengd, ask_sound, gen = st["ask_lengd"], st["ask_sound"], st["gen"]
        spectro = st.get("spectro", False)
        k = lambda name, i: f"{gen}-{name}-{i}"  # unique per set => fresh, empty widgets
        texts, sound_radios, lengd_radios, fbs = [], [], [], []
        if st.get("header"):
            gr.Markdown(st["header"], elem_classes="chapter-head")

        for i, it in enumerate(items, 1):
            with gr.Row(elem_classes="ex-row", equal_height=True):
                gr.Markdown(f"**{i}**", elem_classes="ex-num", scale=0, min_width=30)
                with gr.Column(scale=1 if spectro else 0, min_width=230 if spectro else 44):
                    play = gr.Button("🔊", size="sm", scale=0, min_width=44, elem_classes="play-btn")
                    play.click(None, None, None, js=_play_js(it["audio_uri"]))
                    if spectro:
                        gr.HTML(sample_spec_html(it["clip"], i), key=k("sspec", i))
                with gr.Column(scale=2, min_width=230 if spectro else 100):
                    rec = gr.Audio(sources=["microphone"], type="filepath", show_label=False,
                                   container=False, min_width=100, key=k("rec", i),
                                   waveform_options=gr.WaveformOptions(show_recording_waveform=False))
                    if spectro:
                        uspec = gr.HTML("", key=k("uspec", i))
                        rec.change(lambda p, n=i: user_spec_html(p, n), rec, uspec)
                with gr.Column(scale=0, min_width=150):
                    texts.append(_word_box(it, use_hint, k("txt", i)))
                if ask_sound or ask_lengd:
                    with gr.Column(scale=2, min_width=220):
                        with gr.Row(elem_classes="radio-row"):
                            if ask_sound:
                                sound_radios.append(gr.Radio(
                                    [(f"[{s}]", s) for s in it["sounds"]], show_label=False,
                                    container=False, min_width=0, key=k("snd", i)))
                            if ask_lengd:
                                lengd_radios.append(gr.Radio(
                                    LENGDIR, show_label=False, container=False,
                                    min_width=0, key=k("len", i)))
                fbs.append(gr.Markdown("", elem_classes="fb", scale=2, min_width=150, key=k("fb", i)))

        score_btn = gr.Button("Athuga svör", variant="primary", key=f"{gen}-score")
        summary = gr.Markdown("", key=f"{gen}-summary")

        def score(*vals):
            n = len(items)
            typed = vals[:n]
            rest = vals[n:]
            sel = rest[:n] if ask_sound else [None] * n
            rest = rest[n:] if ask_sound else rest
            lens = rest[:n] if ask_lengd else [None] * n
            msgs, n_s, n_l, n_w = [], 0, 0, 0
            for it, s, w, l in zip(items, sel, typed, lens):
                s_ok, l_ok = s == it["answer"], l == it["lengd"]
                w_ok = word_ok(it, w, use_hint)
                n_s += s_ok; n_w += w_ok; n_l += l_ok

                def line(label, ok, blank, answer):
                    if ok: return f"✅ {label}"
                    if blank: return f"🟣 {label}: Ekkert svar"
                    return f"❌ {label} → **{answer}**"
                lines = [line("orð", w_ok, _word_blank(it, w, use_hint), it["word"])]
                if ask_sound:
                    lines.append(line("hljóð", s_ok, not s, f"[{it['answer']}]"))
                if ask_lengd:
                    lines.append(line("lengd", l_ok, not l, it["lengd"]))
                msgs.append("  \n".join(lines))
            parts = ([f"Hljóð: {n_s}/{n}"] if ask_sound else []) \
                + ([f"Lengdir: {n_l}/{n}"] if ask_lengd else []) + [f"Orð: {n_w}/{n}"]
            perfect = n_w == n and (n_s == n or not ask_sound) and (n_l == n or not ask_lengd)
            total = "**" + " · ".join(parts) + "**" + ("  🎉" if perfect else "")
            return [gr.update(value=m) for m in msgs] + [total]

        score_btn.click(score, texts + sound_radios + lengd_radios, fbs + [summary])

    return state

# hopefully the css will work eventually... but....
def _word_box(it, use_hint, key):
    if use_hint and HINT_STYLE == "nicer":
        classes = ["hint-wrap"]
        if it["prefix"]: classes.append(f"pre-{_cls(it['prefix'])}")
        if it["suffix"]: classes.append(f"suf-{_cls(it['suffix'])}")
        with gr.Column(elem_classes=classes, min_width=0):
            return gr.Textbox(show_label=False, container=False, placeholder="__", key=key,
                              lines=1, max_lines=1, elem_classes="hint-box")
    if use_hint:  # fill in the blank
        return gr.Textbox(it["prefix"] + "__" + it["suffix"], show_label=False, container=False, key=key,
                          lines=1, max_lines=1, elem_classes="plain-box")
    return gr.Textbox(show_label=False, container=False, placeholder="Orð", key=key,
                      lines=1, max_lines=1, elem_classes="plain-box")