src/display/Display.tsx

import React, { useEffect, useRef, useState } from 'react';
import Grid from '@mui/material/Grid';
import Box from '@mui/material/Box';
import { Button, Typography } from '@mui/material';
import ScrollingSections, { ScrollingSectionsRef } from '../scrollingTranscript/components/ScrollingSections';
import { Conversation } from '../types/Conversation';
import { EntailmentDisplay } from './EntailmentDisplay';

interface DisplayDialogProps {
    source: { name: string; key: string; } | undefined;
}

export const DisplayDialog: React.FC<DisplayDialogProps> = ({ source }) => {
    const [conversation, setConversation] = useState<Conversation | undefined>(undefined);
    const [audioLoaded, setAudioLoaded] = useState(false);
    const [isPlaying, setIsPlaying] = useState(false);
    const [currentSectionIndex, setCurrentSectionIndex] = useState(0);
    const audioRef = useRef<HTMLAudioElement>(null);
    const audioSrc = `/sources/${source?.key}/output_audio.wav`;
    console.log('Audio source:', audioSrc);

    const scrollingSentencesRef = useRef<ScrollingSectionsRef>(null);
    const scrollingSourceRef = useRef<ScrollingSectionsRef>(null);

    const handleSectionChange = (index: number) => {
        console.log(`Section changed to index: ${index}`);
        setCurrentSectionIndex(index);
    };

    useEffect(() => {
        setAudioLoaded(false);
        setIsPlaying(false);
        setCurrentSectionIndex(0);
        scrollingSentencesRef.current?.pause();
        scrollingSourceRef.current?.pause();
        scrollingSentencesRef.current?.goToSection(0);
        scrollingSourceRef.current?.goToSection(0);
        if (source === undefined) {
            return;
        }
        const jsonSrc = `/sources/${source.key}/output_segments.json`;
        fetch(jsonSrc)
            .then((res) => res.json())
            .then((data) => setConversation(data))
            .catch((err) => console.error("Failed to load JSON", err));
    }, [source]);

    useEffect(() => {
        const playAudio = async () => {
            try {
                if (audioRef.current) {
                    await audioRef.current.play();
                    console.log('Playback started');
                }
            } catch (err) {
                console.warn('Playback failed:', err);
            }
        };

        const stopAudio = async () => {
            if (audioRef.current) {
                await audioRef.current.pause();
            }
        };

        if (isPlaying && audioLoaded) {
            playAudio();
        } else {
            stopAudio();
        }
    }, [isPlaying, audioLoaded]);


    useEffect(() => {
        if (scrollingSentencesRef.current && scrollingSourceRef.current) {
            if (isPlaying) {
                scrollingSentencesRef.current.play();
                scrollingSourceRef.current.play();
            }
            else {
                scrollingSentencesRef.current.pause();
                scrollingSourceRef.current.pause();
            }
        }
    }, [isPlaying])

    const getSectionStartTime = (index: number): number => {
        if (!conversation || index < 0 || index >= conversation.sections.length) {
            return 0;
        }
        // sum up the durations of all previous sections
        return conversation.sections.slice(0, index).reduce((acc, section) => acc + section.duration, 0);
    }

    if (source === undefined) {
        return (
            <Typography variant="body1" color="textSecondary" gutterBottom style={{ textAlign: 'center' }}>
                Please select a source from the dropdown above.
            </Typography>
        );
    }

    const formatMinutes = (seconds: number): string => {
        const minutes = Math.floor(seconds / 60);
        const remainingSeconds = Math.floor(seconds % 60);
        return `${minutes}:${remainingSeconds < 10 ? '0' : ''}${remainingSeconds}`;
    };

    return (
        <Box sx={{ height: '100vh', display: 'flex', flexDirection: 'column' }}>
            <Box sx={{ display: 'flex', justifyContent: 'center', p: 2 }}>
                <Typography variant="h6" gutterBottom>
                    Model: {conversation?.model || 'Unknown'} - Duration: {conversation ? formatMinutes(conversation.sections.reduce((acc, section) => acc + section.duration, 0)) : 'N/A'}
                </Typography>
            </Box>
            <Box sx={{ display: 'flex', justifyContent: 'center', p: 2 }}>
                <Button variant="outlined" color="secondary" disabled={!audioLoaded} onClick={() => {
                    setIsPlaying(!isPlaying);
                }}>
                    {isPlaying ? "Pause" : `Start playing ${source.name}`}
                </Button>
            </Box>

            <Box sx={{ display: 'flex', justifyContent: 'center', mb: 2 }}>
                <audio src={audioSrc} ref={audioRef} onLoadedData={() => {
                    setAudioLoaded(true);
                    console.log('Audio loaded');
                }} />
            </Box>
            <Box sx={{ display: 'flex', justifyContent: 'center', mb: 2 }}>
                <span>
                    Entailment between source and transcript:
                    {
                        conversation && conversation.sections[currentSectionIndex]?.entailment && (
                            <EntailmentDisplay
                                entailment={conversation.sections[currentSectionIndex].entailment}
                            />
                        )
                    }
                </span>
            </Box>
            <Box sx={{ display: 'flex', justifyContent: 'center' }}>
                <Box
                    sx={{
                        display: 'flex',
                        flexWrap: 'wrap',
                        justifyContent: 'center',
                        maxWidth: '80%',
                        width: '100%',
                        mb: 2,
                    }}
                >
                    {conversation && conversation.sections.map((section, index) => {
                        console.log('Current section index:', currentSectionIndex);
                        return (
                            <Button
                                key={index}
                                variant={currentSectionIndex === index ? "contained" : "outlined"}
                                color="primary"
                                onClick={() => {
                                    scrollingSentencesRef.current?.goToSection(index);
                                    scrollingSourceRef.current?.goToSection(index);
                                    audioRef.current!.currentTime = getSectionStartTime(index);
                                }}
                                sx={{ margin: '4px' }}
                                disabled={isPlaying}
                            >
                                {index + 1}
                            </Button>
                        );
                    })}
                </Box>
            </Box>
            <Grid container spacing={2} sx={{ flex: 1, height: '100%' }}>
                <Grid size={{ xs: 12, md: 4 }} sx={{ display: 'flex', flexDirection: 'column' }}>
                    <Box sx={{ p: 2, textAlign: 'center' }}>
                        <h3>Transcript</h3>
                    </Box>
                    <Box sx={{ p: 2, flex: 1 }}>
                        {conversation && (
                            <ScrollingSections
                                ref={scrollingSentencesRef}
                                sections={conversation.sections}
                                onSectionChange={handleSectionChange}
                                displayProp='sentence'
                                onComplete={() => {
                                    setIsPlaying(false);
                                    scrollingSentencesRef.current?.goToSection(0);
                                    scrollingSourceRef.current?.goToSection(0);
                                    audioRef.current!.currentTime = 0;
                                }}
                            />
                        )}
                    </Box>
                </Grid>
                <Grid size={{ xs: 12, md: 8 }}>
                    <Box sx={{ p: 2, textAlign: 'center' }}>
                        <h3>Source</h3>
                    </Box>
                    {conversation && (
                        <ScrollingSections
                            ref={scrollingSourceRef}
                            sections={conversation.sections}
                            onSectionChange={handleSectionChange}
                            displayProp='source'
                            showProgress={false}
                            displayBeforeAndAfter={false}
                            startAtTop={true}
                        />
                    )}
                </Grid>
            </Grid>
        </Box>
    );
};