import gradio as gr
from dataclasses import dataclass
from pathlib import Path
import json
import hashlib
import os
from typing import List, Tuple, Iterator
import assemblyai as aai
from google import generativeai
from pydub import AudioSegment
import asyncio
import io
from itertools import groupby
from datetime import datetime

prompt = '''
You are an expert transcript editor. Your task is to enhance this transcript for maximum readability while maintaining the core message.

IMPORTANT: Respond ONLY with the enhanced transcript. Do not include any explanations, headers, or phrases like "Here is the transcript."

Note: Below you'll find an auto-generated transcript that may help with speaker identification, but focus on creating your own high-quality transcript from the audio.

Think about your job as if you were transcribing an interview for a print book where the priority is the reading audience. It should just be a total pleasure to read this as a written artifact where all the flubs and repetitions and conversational artifacts and filler words and false starts are removed, where a bunch of helpful punctuation is added. It should basically read like somebody wrote it specifically for reading rather than just something somebody said extemporaneously. 

Please:
1. Fix speaker attribution errors, especially at segment boundaries. Watch for incomplete thoughts that were likely from the previous speaker.

2. Optimize AGGRESSIVELY for readability over verbatim accuracy:
   - Readability is the most important thing!!
   - Remove ALL conversational artifacts (yeah, so, I mean, etc.)
   - Remove ALL filler words (um, uh, like, you know)
   - Remove false starts and self-corrections completely
   - Remove redundant phrases and hesitations
   - Convert any indirect or rambling responses into direct statements
   - Break up run-on sentences into clear, concise statements
   - Maintain natural conversation flow while prioritizing clarity and directness

3. Format the output consistently:
   - Keep the "Speaker X 00:00:00" format (no brackets, no other formatting)
   - DO NOT change the timestamps. You're only seeing a chunk of the full transcript, which is why your 0:00:00 is not the true beginning. Keep the timestamps as they are.
   - Add TWO line breaks between speaker/timestamp and the text
   - Use proper punctuation and capitalization
   - Add paragraph breaks for topic changes
   - When you add paragraph breaks between the same speaker's remarks, no need to restate the speaker attribution
   - Don't go more than four sentences without adding a paragraph break. Be liberal with your paragraph breaks. 
   - Preserve distinct speaker turns

Example input:
Speaker A 00:01:15

Um, yeah, so like, I've been working on this new project at work, you know? And uh, what's really interesting is that, uh, we're seeing these amazing results with the new approach we're taking. Like, it's just, you know, it's really transforming how we do things.

And then, I mean, the thing is, uh, when we showed it to the client last week, they were just, you know, completely blown away by what we achieved. Like, they couldn't even believe it was the same system they had before.

Example output:
Speaker A 00:01:15

I've been working on this new project at work, and we're seeing amazing results with our new approach. It's really transforming how we do things.

When we showed it to the client last week, they were completely blown away by what we achieved. They couldn't believe it was the same system they had before.

Enhance the following transcript, starting directly with the speaker format: 
'''


@dataclass
class Utterance:
    """A single utterance from a speaker"""
    speaker: str
    text: str
    start: int
    end: int

    @property
    def timestamp(self) -> str:
        seconds = int(self.start // 1000)
        hours = seconds // 3600
        minutes = (seconds % 3600) // 60
        seconds = seconds % 60
        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"


class Transcriber:
    def __init__(self, api_key: str):
        aai.settings.api_key = api_key
        self.cache_dir = Path("transcripts/.cache")
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def get_transcript(self, audio_path: Path) -> List[Utterance]:
        cache_file = self.cache_dir / f"{audio_path.stem}.json"
        
        if cache_file.exists():
            with open(cache_file) as f:
                data = json.load(f)
                if data["hash"] == self._get_file_hash(audio_path):
                    return [
                        Utterance(
                            speaker=u["speaker"],
                            text=u["text"],
                            start=u["start"],
                            end=u["end"]
                        )
                        for u in data["utterances"]
                    ]

        config = aai.TranscriptionConfig(speaker_labels=True, language_code="en")
        transcript = aai.Transcriber().transcribe(str(audio_path), config=config)
        
        utterances = [
            Utterance(
                speaker=u.speaker,
                text=u.text,
                start=u.start,
                end=u.end
            )
            for u in transcript.utterances
        ]
        
        cache_data = {
            "hash": self._get_file_hash(audio_path),
            "utterances": [
                {
                    "speaker": u.speaker,
                    "text": u.text,
                    "start": u.start,
                    "end": u.end
                }
                for u in utterances
            ]
        }
        with open(cache_file, "w") as f:
            json.dump(cache_data, f, indent=2)
            
        return utterances

    def _get_file_hash(self, file_path: Path) -> str:
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()


class Enhancer:
    def __init__(self, api_key: str):
        generativeai.configure(api_key=api_key)
        self.model = generativeai.GenerativeModel("gemini-2.0-flash-lite-preview-02-05")
        self.prompt = prompt

    async def enhance_chunks(self, chunks: List[Tuple[str, io.BytesIO]]) -> List[str]:
        semaphore = asyncio.Semaphore(3)
        
        async def process_chunk(i: int, chunk: Tuple[str, io.BytesIO]) -> str:
            text, audio = chunk
            async with semaphore:
                audio.seek(0)
                response = await self.model.generate_content_async(
                    [self.prompt, text, {"mime_type": "audio/mp3", "data": audio.read()}]
                )
                return response.text

        tasks = [process_chunk(i, chunk) for i, chunk in enumerate(chunks)]
        results = await asyncio.gather(*tasks)
        return results


@dataclass
class SpeakerDialogue:
    speaker: str
    utterances: List[Utterance]
    
    @property
    def start(self) -> int:
        return self.utterances[0].start
    
    @property
    def end(self) -> int:
        return self.utterances[-1].end
    
    @property
    def timestamp(self) -> str:
        return self.utterances[0].timestamp
    
    def format(self, markdown: bool = False) -> str:
        texts = [u.text + "\n\n" for u in self.utterances]
        combined_text = ''.join(texts).rstrip()
        if markdown:
            return f"**Speaker {self.speaker}** *{self.timestamp}*\n\n{combined_text}"
        return f"Speaker {self.speaker} {self.timestamp}\n\n{combined_text}"


def group_utterances_by_speaker(utterances: List[Utterance]) -> Iterator[SpeakerDialogue]:
    for speaker, group in groupby(utterances, key=lambda u: u.speaker):
        yield SpeakerDialogue(speaker=speaker, utterances=list(group))


def estimate_tokens(text: str, chars_per_token: int = 4) -> int:
    return (len(text) + chars_per_token - 1) // chars_per_token


def chunk_dialogues(dialogues: Iterator[SpeakerDialogue], max_tokens: int = 2000, chars_per_token: int = 4) -> List[List[SpeakerDialogue]]:
    chunks = []
    current_chunk = []
    current_text = ""
    
    for dialogue in dialogues:
        formatted = dialogue.format()
        new_text = current_text + "\n\n" + formatted if current_text else formatted
        
        if current_chunk and estimate_tokens(new_text, chars_per_token) > max_tokens:
            chunks.append(current_chunk)
            current_chunk = [dialogue]
            current_text = formatted
        else:
            current_chunk.append(dialogue)
            current_text = new_text
    
    if current_chunk:
        chunks.append(current_chunk)
    
    return chunks


def format_chunk(dialogues: List[SpeakerDialogue], markdown: bool = False) -> str:
    return "\n\n".join(dialogue.format(markdown=markdown) for dialogue in dialogues)


def prepare_audio_chunks(audio_path: Path, utterances: List[Utterance]) -> List[Tuple[str, io.BytesIO]]:
    dialogues = group_utterances_by_speaker(utterances)
    chunks = chunk_dialogues(dialogues)
    
    audio = AudioSegment.from_file(audio_path)
    
    prepared = []
    for chunk in chunks:
        segment = audio[chunk[0].start:chunk[-1].end]
        buffer = io.BytesIO()
        segment.export(buffer, format="mp3", parameters=["-q:a", "9"])
        prepared.append((format_chunk(chunk, markdown=False), buffer))
    
    return prepared


def apply_markdown_formatting(text: str) -> str:
    import re
    pattern = r"(Speaker \w+) (\d{2}:\d{2}:\d{2})"
    return re.sub(pattern, r"**\1** *\2*", text)


def rename_speakers(text: str, speaker_map: dict) -> str:
    """Replace speaker labels using the provided mapping"""
    result = text
    for old_name, new_name in speaker_map.items():
        # Replace both markdown and plain text formats
        result = result.replace(f"**Speaker {old_name}**", f"**{new_name}**")
        result = result.replace(f"Speaker {old_name}", new_name)
    return result


def create_downloadable_file(content: str, prefix: str) -> str:
    """Create a temporary file with the content and return filepath"""
    temp_dir = Path("temp_downloads")
    temp_dir.mkdir(exist_ok=True)
    
    # Create a unique filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{prefix}_{timestamp}.md"
    filepath = temp_dir / filename
    
    # Write content to file
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(content)
        
    return str(filepath)


def process_audio(audio_file):
    try:
        temp_path = Path("temp_audio")
        temp_path.mkdir(exist_ok=True)
        temp_file = temp_path / "temp_audio.mp3"
        
        try:
            with open(temp_file, "wb") as f:
                f.write(audio_file)

            # Initial state - clear both transcripts
            yield (
                gr.update(value="", visible=True),  # original transcript
                gr.update(value="", visible=True),  # enhanced transcript
                None,  # original download
                None,  # enhanced download
            )
            
            # Get transcript
            transcriber = Transcriber(os.getenv("ASSEMBLYAI_API_KEY"))
            utterances = transcriber.get_transcript(temp_file)
            dialogues = list(group_utterances_by_speaker(utterances))
            original = format_chunk(dialogues, markdown=True)
            
            # Create downloadable file for original transcript
            original_file = create_downloadable_file(original, "original_transcript")
            
            # Show original transcript
            yield (
                gr.update(value=original, visible=True),
                gr.update(value="", visible=True),
                original_file,
                None,
            )
            
            try:
                enhancer = Enhancer(os.getenv("GOOGLE_API_KEY"))
                chunks = prepare_audio_chunks(temp_file, utterances)
                enhanced = asyncio.run(enhancer.enhance_chunks(chunks))
                merged = "\n\n".join(chunk.strip() for chunk in enhanced)
                merged = apply_markdown_formatting(merged)
                
                # Create downloadable file for enhanced transcript
                enhanced_file = create_downloadable_file(merged, "enhanced_transcript")
                
                # Show final result
                yield (
                    gr.update(value=original, visible=True),
                    gr.update(value=merged, visible=True),
                    original_file,
                    enhanced_file,
                )
                
            except Exception as e:
                yield (
                    gr.update(value=original, visible=True),
                    gr.update(value=f"Error: {str(e)}", visible=True),
                    original_file,
                    None,
                )
            
        finally:
            # Cleanup temp files
            if os.path.exists(temp_file):
                os.remove(temp_file)
        
    except Exception as e:
        if isinstance(e, gr.Error):
            raise
        raise gr.Error(f"Error processing audio: {str(e)}")


# Create the Gradio interface
with gr.Blocks(title="Transcript Enhancer") as demo:
    gr.Markdown("""
    # 🎙️ Audio Transcript Enhancer
    
    Upload an audio file to get both an automated transcript and an enhanced version using AI.
    
    1. The original transcript is generated using AssemblyAI with speaker detection
    2. The enhanced version uses Google's Gemini to improve clarity and readability
    """)
    
    with gr.Row():
        audio_input = gr.File(
            label="Upload Audio File", 
            type="binary", 
            file_count="single", 
            file_types=["audio"]
        )
    
    with gr.Row():
        transcribe_btn = gr.Button("📝 Transcribe & Enhance")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Original Transcript")
            original_download = gr.File(
                label="Download as Markdown",
                file_count="single",
                visible=True,
                interactive=False,
            )
            original_output = gr.Markdown()
        
        with gr.Column():
            gr.Markdown("### Enhanced Transcript")
            enhanced_download = gr.File(
                label="Download as Markdown",
                file_count="single",
                visible=True,
                interactive=False,
            )
            enhanced_output = gr.Markdown()
    
    # Add some CSS to style the download buttons
    gr.Markdown("""
        <style>
        .download-button {
            margin-top: 10px;
        }
        </style>
    """)

    transcribe_btn.click(
        fn=process_audio,
        inputs=[audio_input],
        outputs=[
            original_output,
            enhanced_output,
            original_download,
            enhanced_download
        ]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(max_file_size=5 * gr.FileSize.GB)  # Backend limit