# -*- coding: utf-8 -*-
"""Common base for the vertical videos: 9:16 format, subtitles,
progress bar and a system of 'beats' synced with the audio."""
import json, os
from contextlib import contextmanager
from manim import *

# --- TikTok format 1080x1920 --------------------------------------------------
if os.environ.get("DRAFT"):
    config.pixel_width, config.pixel_height = 540, 960
else:
    config.pixel_width, config.pixel_height = 1080, 1920
config.frame_width  = 9.0
config.frame_height = 16.0
config.frame_rate   = 30
config.background_color = "#0B0F1A"

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

FONT    = "Roboto"
AMBER   = "#FFD166"
C_GREEN   = "#06D6A0"
C_PINK    = "#EF476F"
C_SKY = "#4CC9F0"
C_GRAY    = "#8D99AE"

GAP      = 0.30   # silence between segments
SUB_Y    = -4.25  # subtitles sit above the TikTok UI
TOP_Y    =  6.45


class _Beat:
    """Accumulates the time used inside a segment so the rest can be filled."""
    def __init__(self, scene, floor):
        self.s, self.floor, self.used = scene, floor, 0.0

    def play(self, *anims, run_time=1.0, **kw):
        self.s.play(*anims, run_time=run_time, **kw)
        self.used += run_time

    def wait(self, t):
        if t > 0.02:
            self.s.wait(t)
            self.used += t

    def add(self, *m):
        self.s.add(*m)

    def remove(self, *m):
        self.s.remove(*m)

    def close(self):
        rest = self.floor - self.used
        if rest > 0.02:
            self.s.wait(rest)
            self.used += rest
        if GAP > 0.02:
            self.s.wait(GAP)
            self.used += GAP
        return self.used


class TikTok(Scene):
    NAME = ""

    # -- infrastructure -------------------------------------------------------
    def prepare(self):
        with open(os.path.join(ROOT, "audio", self.NAME, "segments.json"), encoding="utf-8") as f:
            self.data = json.load(f)
        self.segs = self.data["segments"]
        self.timeline, self.clock = [], 0.0
        self._sub = None
        self._expected = sum(s["dur"] + GAP for s in self.segs)
        self.backdrop()
        self._progress_bar()

    def backdrop(self):
        dots = VGroup(*[
            Dot(radius=0.035, color="#243050", fill_opacity=0.55).move_to([x, y, 0])
            for x in np.arange(-4.2, 4.3, 1.05)
            for y in np.arange(-7.5, 7.6, 1.05)
        ])
        self.add(dots)
        return dots

    def _progress_bar(self):
        w = config.frame_width
        track = Rectangle(width=w, height=0.075, stroke_width=0,
                          fill_color="#1B2233", fill_opacity=1).move_to(UP * (TOP_Y + 1.15))
        bar = Rectangle(width=0.002, height=0.075, stroke_width=0,
                        fill_color=AMBER, fill_opacity=1).move_to(track.get_left(), LEFT)
        self.elapsed = 0.0
        total, left, y = self._expected, -w / 2, track.get_y()

        def upd(m, dt):
            self.elapsed += dt
            frac = min(self.elapsed / total, 1.0)
            nw = max(frac * w, 0.002)
            m.stretch_to_fit_width(nw)
            m.move_to([left + nw / 2, y, 0])

        bar.add_updater(upd)
        self.add(track, bar)

    def subtitle(self, text):
        """Changes the subtitle without spending animation time."""
        if self._sub:
            self.remove(self._sub)
        t = Text(text, font=FONT, font_size=44, weight=BOLD, color=WHITE,
                 line_spacing=0.8).move_to(UP * SUB_Y)
        if t.width > 8.0:
            t.scale(8.0 / t.width)
        box = RoundedRectangle(corner_radius=0.18, stroke_width=0,
                               fill_color="#000000", fill_opacity=0.55,
                               width=t.width + 0.55, height=t.height + 0.42).move_to(t)
        self._sub = VGroup(box, t)
        self.add(self._sub)

    @contextmanager
    def beat(self, i, sub=True):
        """Block synced with audio segment i."""
        if sub:
            self.subtitle(self.segs[i]["sub"])
        b = _Beat(self, self.segs[i]["dur"])
        yield b
        used = b.close()
        self.timeline.append({"i": i, "start": round(self.clock, 3),
                              "dur": round(used, 3), "audio": self.segs[i]["dur"]})
        self.clock += used

    def finish(self):
        if self._sub:
            self.play(FadeOut(self._sub), run_time=0.4)
            self.clock += 0.4
        self.wait(0.5)
        self.clock += 0.5
        os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
        with open(os.path.join(ROOT, "out", f"{self.NAME}_timeline.json"), "w") as f:
            json.dump({"name": self.NAME, "gap": GAP, "total": round(self.clock, 3),
                       "beats": self.timeline}, f, indent=2)

    # -- visual helpers -------------------------------------------------------
    def chip(self, text, color=AMBER):
        t = Text(text, font=FONT, font_size=34, weight=BOLD, color="#0B0F1A")
        box = RoundedRectangle(corner_radius=0.3, stroke_width=0, fill_color=color,
                               fill_opacity=1, width=t.width + 0.7, height=t.height + 0.45)
        g = VGroup(box, t).move_to(UP * TOP_Y)
        return g

    def fit(self, m, w=8.3):
        if m.width > w:
            m.scale(w / m.width)
        return m

    def big(self, text, size=76, color=WHITE):
        return Text(text, font=FONT, font_size=size, weight=BOLD, color=color)