Py.Cafe

import solara
import pycafe
import pandas as pd
import altair as alt
import json
from datetime import datetime, timedelta
from cryptography.fernet import Fernet
#import solara.website
from pathlib import Path
import PIL.Image

# Global reactive states


continuous_update = solara.reactive(True)
file_path_state = solara.reactive("")
transcript_data_state = solara.reactive(None)
error_message = solara.reactive("")
selected_datum = solara.reactive(None)
solara.lab.theme.themes.light.primary = "#AEDEE9"
solara.lab.theme.themes.light.secondary = "#B0DFE9"


image_path = "./RushXplorer.png"


solara.lab.theme.themes.dark.primary ="#E17E35"#"#CF884F"### #"#CF884F" #

RUSHXPLORER_KEY = pycafe.get_secret(
        "RUSHXPLORER_KEY",
        """We need the RUSHXPLORER key to make the app work.""",
    )
def _init_fernet():
    return Fernet(RUSHXPLORER_KEY.encode())
_f = _init_fernet()

# Utility functions
def time_to_seconds(time_str: str) -> float:
    t = datetime.strptime(time_str, "%H:%M:%S")
    return timedelta(hours=t.hour, minutes=t.minute, seconds=t.second).total_seconds()

def seconds_to_time(seconds: int) -> str:
    return str(timedelta(seconds=int(seconds)))

def load_json(filename: str) -> dict:
    try:
        with open(filename, 'r') as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading {filename}: {e}")
        return {}

# Categorical options for relevance
relevance_options = ["low", "medium", "high"]
    
image = PIL.Image.open(image_path)

@solara.component
def Home():

    solara.lab.ThemeToggle()
    solara.Image(image,width="25%")
    solara.Markdown(
        """
Welcome to the RushXplorer Demo Dashboard! This proof-of-concept showcases AI-powered transcript processing and a future interactive dashboard for rushes editing:

- **Transcript**: Your raw footage transcripts—of any format—are **AI-preprocessed** into a structured JSON model, ready for exploration.
- **Characters**: AI-drafted character lists and metadata are fully editable, with future support for **prompt-driven LLM chats** to refine roles and relevance.
- **Storyline**: Narrative arcs and sub-arcs are **AI-generated**, editable on-the-fly, and can be reorganized via conversational prompts.
- **Highlights**: Key moments are initially selected by AI for emotional and narrative impact; adjust or add highlights through direct editing or LLM guidance.
- **Episode & Teaser**: Episode cuts and teasers are assembled by AI into EDLs (Edit Decision Lists) of **any format**, and can be re-formatted or re-ordered dynamically.
- **Editorial**: Manage review questions and tasks with AI-suggested priorities, all editable and extendable via chat.

Each tab combines interactive Altair visualizations with JSON viewers for transparency. Soon, you’ll be able to talk directly to the embedded LLM to adjust data, regenerate edits, and output custom EDL formats—all in real time.
        """
    )

@solara.component
def FileInput():
    solara.Button(
        "Insert Demo Rushes",
        on_click=lambda: file_path_state.set("./transcript_test.enc")
    )
#    def handle_file_load():
#        path = file_path_state.value.strip()
#        try:
#            with open(path, 'r') as f:
#                transcript_data_state.value = json.load(f)
#                error_message.value = ''
#        except Exception as e:
#            transcript_data_state.value = None
#            error_message.value = f"Error: {e}"

    def handle_file_load():
        path = file_path_state.value.strip()
        try:
            # Open encrypted file in binary mode
            with open(path, 'rb') as f_enc:
                encrypted = f_enc.read()
            # Decrypt and decode bytes to string
            decrypted = _f.decrypt(encrypted).decode('utf-8')
            # Parse JSON
            transcript_data_state.value = json.loads(decrypted)
            error_message.value = ''
        except Exception as e:
            transcript_data_state.value = None
            error_message.value = f"Error loading encrypted JSON: {e}"


    with solara.Column():
        solara.InputText(label="Transcript JSON file:", value=file_path_state)
        solara.Button("Load Transcript", on_click=handle_file_load)
        if error_message.value:
            solara.Error(error_message.value)

@solara.component
def TranscriptProcessor():
    data = transcript_data_state.value
    if data is None:
        solara.Markdown("No transcript loaded.")
        return
    transcripts = data.get('transcripts', [])
    df = pd.DataFrame(transcripts).dropna(subset=['timestamp', 'dialogue'])
    df['ts_sec'] = df['timestamp'].apply(time_to_seconds)
    df.sort_values('ts_sec', inplace=True)
    window_size = 120
    total = int(df['ts_sec'].max() if not df.empty else 0)
    speakers = sorted(df['role'].unique())

    windows = []
    for start in range(0, total, window_size):
        end = start + window_size
        w = df[(df['ts_sec'] >= start) & (df['ts_sec'] < end)]
        summary = {'start_time': seconds_to_time(start), 'end_time': seconds_to_time(end)}
        for s in speakers:
            wc = int(w[w['role']==s]['dialogue'].str.split().str.len().sum()) if not w.empty else 0
            summary[f"{s}_word_count"] = wc
        windows.append(summary)
    result_df = pd.DataFrame(windows)

    solara.DataFrame(result_df)

    # Heatmap of word counts
    melt_vars = [col for col in result_df.columns if col.endswith('_word_count')]
    melted = result_df.melt(
        id_vars=['start_time', 'end_time'],
        value_vars=melt_vars,
        var_name='speaker',
        value_name='word_count'
    )
    melted['speaker'] = melted['speaker'].str.replace('_word_count', '')

    heatmap = (
        alt.Chart(melted, title='Speaker Word Count Over Time')
        .mark_rect()
        .encode(
            alt.X('start_time:N', title='Start Time', sort=None),
            alt.Y('speaker:N', title='Speaker'),
            alt.Color('word_count:Q', title='Words Spoken'),
            tooltip=[
                alt.Tooltip('start_time', title='Start'),
                alt.Tooltip('end_time', title='End'),
                alt.Tooltip('speaker', title='Speaker'),
                alt.Tooltip('word_count', title='Words')
            ],
        )
        .configure_view(step=13, strokeWidth=0)
        .configure_axis(domain=False)
        .properties(width='container')
    )
    with solara.Card('Speaker Word Count Heatmap'):
        solara.AltairChart(heatmap, on_click=lambda d: selected_datum.set(d))

@solara.component
def Transcript():
    solara.Markdown("Complete, timestamped dialogue entries capturing every role’s spoken lines and moments in the raw rushes.")
    FileInput()
    TranscriptProcessor()

@solara.component
def ViewerJSON(filename: str, id_field: str = None):
    import json
    import pandas as pd

    relevance_options = ["low", "medium", "high"]
    raw = load_json(filename)
    if not raw:
        solara.Markdown(f"Failed to load {filename}")
        return

    # --- Special branch for editorial.json ---
    if isinstance(raw, dict) and "questions" in raw and "tasks" in raw:
        # Render questions
        show_dialog = solara.use_reactive(False)
        solara.Markdown("## Questions")
        for q in raw["questions"]:
            solara.Markdown(f"- {q}")
            with solara.Row():
                solara.Button("Answer", on_click=lambda: show_dialog.set(True))

        # Render tasks
        solara.Markdown("## Tasks")
        for task in raw["tasks"]:
            solara.Markdown("—")
            for k, v in task.items():
                with solara.Columns([1, 3]):
                    solara.Markdown(f"**{k.replace('_', ' ').title()}**")
                    # Priority as select
                    if k == "priority" and isinstance(v, str):
                        sel, set_sel = solara.use_state(v)
                        solara.Select(label="", value=sel, values=relevance_options, on_value=set_sel)
                    # related_highlights list
                    elif isinstance(v, list):
                        solara.Markdown(", ".join(v) or "None")
                    else:
                        solara.Markdown(str(v))
            with solara.Row():
                solara.Button("Execute task", on_click=lambda: show_dialog.set(True))
                #solara.Button("Show your answer", on_click=lambda: show_conditional.set(not show_conditional.value))
                #with solara.lab.ConfirmationDialog(open=show_dialog):
                #    solara.InputText("Type your additional task instructions here", autofocus=True)
                
        return

    # --- General JSON viewer below ---
    # Normalize raw into a list of entries
    if isinstance(raw, dict) and id_field and raw.get(id_field) is not None:
        items = [raw]
    elif isinstance(raw, dict):
        lists = [v for v in raw.values() if isinstance(v, list)]
        if len(lists) == 1:
            items = lists[0]
        elif all(isinstance(v, dict) for v in raw.values()):
            items = list(raw.values())
        else:
            solara.Markdown("Unsupported JSON dict structure")
            return
    elif isinstance(raw, list):
        items = raw
    else:
        solara.Markdown("Unsupported JSON structure")
        return

    if not items:
        solara.Markdown("No entries to display")
        return

    # Auto-detect id_field if not provided
    if id_field is None:
        possible = [k for k in items[0].keys() if k.endswith("_id")]
        id_field = possible[0] if possible else list(items[0].keys())[0]

    # Dropdown for selecting entry
    ids = [str(item.get(id_field, i)) for i, item in enumerate(items)]
    selected, set_selected = solara.use_state(ids[0])
    solara.Select(
        label=id_field.replace("_", " ").title(),
        value=selected,
        values=ids,
        on_value=set_selected,
    )

    # Render selected entry in two columns
    idx = ids.index(selected)
    entry = items[idx]
    for k, v in entry.items():
        with solara.Columns([1, 3]):
            solara.Markdown(f"**{k.replace('_', ' ').title()}**")
            # categorical select for relevance-like fields
            if isinstance(v, str) and v.lower() in relevance_options:
                sel, set_sel = solara.use_state(v.lower())
                solara.Select(label="", value=sel, values=relevance_options, on_value=set_sel)
            # primitive types
            elif isinstance(v, (str, int, float, bool)):
                val = solara.reactive(str(v))
                solara.InputText(label="", value=val, continuous_update=continuous_update.value)
            # list of dicts -> DataFrame
            elif isinstance(v, list) and v and isinstance(v[0], dict):
                df = pd.DataFrame(v)
                solara.DataFrame(df)
            # list of primitives
            elif isinstance(v, list):
                solara.Markdown(", ".join(map(str, v)))
            # nested dict -> code block
            elif isinstance(v, dict):
                code = json.dumps(v, indent=2)
                solara.Markdown(f"```json\n{code}\n```")
            else:
                solara.Markdown(str(v))


@solara.component
def Characters():
    solara.Markdown("A roster of all persons on scene, with metadata on their importance and which story arcs and highlights they appear in.")
    # Summary plot: number of arcs per character
    chars = load_json("CHARACTERS.json").get("characters", []) if isinstance(load_json("CHARACTERS.json"), dict) else load_json("CHARACTERS.json")
    df_chars = pd.DataFrame(chars)
    df_chars['arc_count'] = df_chars['relevant_arcs'].apply(len)
    arc_chart = (
        alt.Chart(df_chars, title='Relevant Arcs per Character')
        .mark_bar()
        .encode(
            alt.X('name:N', sort=None),
            alt.Y('arc_count:Q', title='Number of Arcs'),
            tooltip=['name', 'relative_importance']
        )
        .properties(width='container')
    )
    solara.AltairChart(arc_chart)
    ViewerJSON("CHARACTERS.json", id_field="character_id")

@solara.component
def Highlights():
    solara.Markdown("Selected key moments with IDs, arc associations, precise timestamps, summaries, emotional impact scores, and involved characters.")
    # Load storyline and highlights JSON
    sl = load_json("STORYLINE.json")
    storyline = sl if isinstance(sl, dict) else next(
        (el for el in sl if isinstance(el, dict) and 'story_arcs' in el), {}
    )
    hl = load_json("HIGHLIGHTS.json")
    highlights = hl if isinstance(hl, dict) else next(
        (el for el in hl if isinstance(el, dict) and 'highlights' in el), {}
    )

    # Helpers to convert between HH:MM:SS and seconds
    def to_seconds(ts: str) -> int:
        h, m, s = map(int, ts.split(':'))
        return h * 3600 + m * 60 + s

    def to_time(sec: int) -> str:
        h = sec // 3600
        m = (sec % 3600) // 60
        s = sec % 60
        return f"{h:02d}:{m:02d}:{s:02d}"

    # Define time windows across full timeline
    start_t = to_seconds(storyline.get('start_time', '00:00:00'))
    end_t = to_seconds(storyline.get('end_time', '00:00:00'))
    window_size = 60  # seconds per bin
    windows = []
    for ws in range(start_t, end_t, window_size):
        we = min(ws + window_size, end_t)
        windows.append({'start_s': ws, 'end_s': we,
                        'start_time': to_time(ws), 'end_time': to_time(we)})

    # Precompute highlight midpoints
    hl_items = []
    for item in highlights.get('highlights', []):
        s = to_seconds(item['start_timestamp'])
        e = to_seconds(item['end_timestamp'])
        hl_items.append({'arc_id': item['arc_id'], 'mid': (s + e)/2})

    # Build records: one row per window per arc if arc covers window
    records = []
    for w in windows:
        for arc in storyline.get('story_arcs', []):
            a_start = to_seconds(arc['start_timestamp'])
            a_end = to_seconds(arc['end_timestamp'])
            if not (w['end_s'] <= a_start or w['start_s'] >= a_end):
                cnt = sum(1 for h in hl_items if h['arc_id'] == arc['arc_id']
                          and w['start_s'] <= h['mid'] < w['end_s'])
                records.append({
                    'start_time': w['start_time'],
                    'arc_title': arc['title'],
                    'highlight_count': cnt
                })
    df = pd.DataFrame(records)
    if df.empty:
        solara.Markdown("No data available for highlight density plot.")
        return

    # State for click selection
    selected = solara.use_state(None)

    # Heatmap: window start vs arc, intensity = highlight count
    heatmap = (
        alt.Chart(df, title='Highlights Density per Story Arc Over Time')
        .mark_rect()
        .encode(
            x=alt.X('start_time:N', title='Window Start', sort=None),
            y=alt.Y('arc_title:N', title='Story Arc', sort=[arc['title'] for arc in storyline.get('story_arcs', [])]),
            color=alt.Color('highlight_count:Q', title='Highlights', scale=alt.Scale(scheme='orangered')),
            tooltip=[
                alt.Tooltip('start_time', title='Window Start'),
                alt.Tooltip('arc_title', title='Arc'),
                alt.Tooltip('highlight_count', title='Highlights')
            ]
        )
        .properties(width='container', height=400)
    )

    # Display chart with click handler
    with solara.Card('Story arcs & highlights over time'):
        solara.AltairChart(heatmap)#, on_click=lambda datum: selected.set(datum))
    # Show selected detail
    #if selected.value:
    #    solara.Markdown(f"**Selected**: Window `{selected.value['start_time']}`, Arc `{selected.value['arc_title']}`, Highlights `{selected.value['highlight_count']}`")

    ViewerJSON("HIGHLIGHTS.json", id_field="highlight_id")

@solara.component
def Storyline():
    solara.Markdown("Definitions of each narrative arc—IDs, titles, summaries, start/end times, and nested sub-arcs that structure the full timeline.")
    sl = load_json("STORYLINE.json")
    storyline = sl if isinstance(sl, dict) else next(
        (el for el in sl if isinstance(el, dict) and 'story_arcs' in el), {}
    )
    hl = load_json("HIGHLIGHTS.json")
    highlights = hl if isinstance(hl, dict) else next(
        (el for el in hl if isinstance(el, dict) and 'highlights' in el), {}
    )

    # Helpers to convert between HH:MM:SS and seconds
    def to_seconds(ts: str) -> int:
        h, m, s = map(int, ts.split(':'))
        return h * 3600 + m * 60 + s

    def to_time(sec: int) -> str:
        h = sec // 3600
        m = (sec % 3600) // 60
        s = sec % 60
        return f"{h:02d}:{m:02d}:{s:02d}"

    # Define time windows across full timeline
    start_t = to_seconds(storyline.get('start_time', '00:00:00'))
    end_t = to_seconds(storyline.get('end_time', '00:00:00'))
    window_size = 60  # seconds per bin
    windows = []
    for ws in range(start_t, end_t, window_size):
        we = min(ws + window_size, end_t)
        windows.append({'start_s': ws, 'end_s': we,
                        'start_time': to_time(ws), 'end_time': to_time(we)})

    # Precompute highlight midpoints
    hl_items = []
    for item in highlights.get('highlights', []):
        s = to_seconds(item['start_timestamp'])
        e = to_seconds(item['end_timestamp'])
        hl_items.append({'arc_id': item['arc_id'], 'mid': (s + e)/2})

    # Build records: one row per window per arc if arc covers window
    records = []
    for w in windows:
        for arc in storyline.get('story_arcs', []):
            a_start = to_seconds(arc['start_timestamp'])
            a_end = to_seconds(arc['end_timestamp'])
            if not (w['end_s'] <= a_start or w['start_s'] >= a_end):
                cnt = sum(1 for h in hl_items if h['arc_id'] == arc['arc_id']
                          and w['start_s'] <= h['mid'] < w['end_s'])
                records.append({
                    'start_time': w['start_time'],
                    'arc_title': arc['title'],
                    'highlight_count': cnt
                })
    df = pd.DataFrame(records)
    if df.empty:
        solara.Markdown("No data available for highlight density plot.")
        return

    # State for click selection
    selected = solara.use_state(None)

    # Heatmap: window start vs arc, intensity = highlight count
    heatmap = (
        alt.Chart(df, title='Highlights Density per Story Arc Over Time')
        .mark_rect()
        .encode(
            x=alt.X('start_time:N', title='Window Start', sort=None),
            y=alt.Y('arc_title:N', title='Story Arc', sort=[arc['title'] for arc in storyline.get('story_arcs', [])]),
            color=alt.Color('highlight_count:Q', title='Highlights', scale=alt.Scale(scheme='orangered')),
            tooltip=[
                alt.Tooltip('start_time', title='Window Start'),
                alt.Tooltip('arc_title', title='Arc'),
                alt.Tooltip('highlight_count', title='Highlights')
            ]
        )
        .properties(width='container', height=400)
    )

    # Display chart with click handler
    with solara.Card('Story arcs & highlights over time'):
        solara.AltairChart(heatmap)
    ViewerJSON("STORYLINE.json", id_field="arc_id")

@solara.component
def Episode():
    solara.Markdown("The assembled episode cut: chosen arcs, highlights, characters, and dialogue snippets mapped to a target run time.")
    ViewerJSON("EPISODE.json", id_field="episode_id")

@solara.component
def Teaser():
    solara.Markdown("A concise preview edit: top highlights and key moments distilled into a short teaser with associated timestamps and descriptions.")
    ViewerJSON("TEASER.json", id_field="teaser_id")

@solara.component
def Editorial():
    solara.Markdown("A collection of review questions and pending tasks—each with status, priority, and related highlights—guiding editorial decisions and improvements.")
    # Freeform text viewer
    ViewerJSON("EDITORIAL.json")

routes = [
    solara.Route(path='/', component=Home, label='RushXplorer'),
    solara.Route(path='transcript', component=Transcript, label='Transcript'),
    solara.Route(path='characters', component=Characters, label='Characters'),
    solara.Route(path='storyline', component=Storyline, label='Storyline'),
    solara.Route(path='highlights', component=Highlights, label='Highlights'),
    solara.Route(path='episode', component=Episode, label='Episode'),
    solara.Route(path='teaser', component=Teaser, label='Teaser'),
    solara.Route(path='editorial', component=Editorial, label='Editorial'),
]