#!/usr/bin/env python3
"""
Create visual HTML map showing:
- Text broken into chunks
- Timestamps for each chunk
- Missing words highlighted in red
- Pattern analysis
"""

import os
import sys
import re
from openai import OpenAI

OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
if not OPENAI_API_KEY:
    print("❌ OPENAI_API_KEY not set!")
    sys.exit(1)

client = OpenAI(api_key=OPENAI_API_KEY)

# Reference text
REFERENCE_TEXT = """<excited> In the heart of an enchanted kingdom, where magic flowed like rivers and mythical beasts roamed freely, there lived a legendary warrior princess named Aisha. She was known throughout the land for her incredible courage, her mastery of ancient combat arts, and her ability to communicate with all living creatures.

<whisper> But few knew of the prophecy that surrounded her birth. The ancient seers had foretold that she would one day face the greatest darkness the world had ever seen, and only she possessed the power to defeat it.

Aisha spent her childhood training in the sacred mountains with the wise monks. <curious> They taught her not just the physical arts of combat, but also the spiritual disciplines that would strengthen her mind and soul. Every dawn, she would practice her sword techniques as the sun rose over the peaks. Every evening, she would meditate under the stars, learning to harness the cosmic energies that flowed through all things.

<angry> One fateful day, dark clouds gathered over the kingdom! The evil demon lord Ravaksh had awakened from his thousand-year slumber. His army of shadow warriors began to sweep across the land, corrupting everything they touched. Villages were destroyed, forests withered, and despair spread like wildfire.

<sigh> The kingdom's armies fought valiantly, but they were no match for Ravaksh's dark magic. One by one, the greatest warriors fell. The king's advisors urged him to flee, to abandon the kingdom and save what few lives they could.

<excited> But Aisha refused to give up! She knew this was the moment the prophecy had spoken of. She gathered the remaining defenders and devised a daring plan. They would launch a surprise attack on Ravaksh's fortress while he was still consolidating his power.

एक बार की बात है, ऐशा ने अपने सबसे भरोसेमंद योद्धाओं को इकट्ठा किया। उन्होंने एक खतरनाक योजना बनाई। वे रावक्ष के किले पर हमला करेंगे। यह असंभव लग रहा था, लेकिन उनके पास कोई विकल्प नहीं था। राज्य का भविष्य उनके हाथों में था।

<whisper> Under the cover of darkness, they approached the fortress. The shadow guards patrolled the walls, their red eyes glowing in the night. <laugh> But Aisha had learned from the wind spirits how to move unseen! She led her team through secret passages known only to the ancient monks.

The infiltration was successful. They reached the inner sanctum where Ravaksh was performing a dark ritual. <angry> He was trying to summon an even greater evil from the void! There was no time to waste.

<excited> Aisha drew her legendary sword, which began to glow with divine light! The blade had been forged by the celestial smiths themselves, imbued with the power to vanquish any darkness. As she charged forward, her companions engaged the shadow warriors, giving her a clear path to Ravaksh.

The battle that followed was epic. <curious> Ravaksh wielded powers that defied comprehension - he could bend reality itself, create illusions that seemed more real than reality, and summon storms of pure destruction. But Aisha had trained for this her entire life.

She dodged his dark bolts, deflected his curse spells, and pressed forward relentlessly. <whisper> She remembered the words of her master: "True strength comes not from power, but from purpose." Her purpose was clear - protect her people, save her kingdom, restore the light.

With a mighty cry, she leaped high into the air, her sword blazing like a star. <excited> The blade struck true, piercing through Ravaksh's dark armor and into his corrupted heart! The demon lord let out a terrible scream as the divine light consumed him.

<laugh> The shadow army dissolved like smoke in the wind! The dark clouds parted, and sunlight flooded the land once more. The corrupted forests began to heal, flowers bloomed, and the rivers ran clear again.

<excited> Aisha returned to her kingdom as a hero beyond measure! The people celebrated for seven days and seven nights. Songs were written about her bravery, statues were erected in her honor, and her legend would be told for a thousand generations.

<whisper> But Aisha remained humble. She knew that true victory was not in defeating enemies, but in protecting those she loved. <sigh> She had seen too much loss, too much suffering. From that day forward, she dedicated herself to teaching the next generation, ensuring they would be ready when darkness threatened again.

<excited> And so the kingdom prospered under her watchful protection, and peace reigned for many years to come!"""

AUDIO_FILE = "correct_format_test.wav"

print("="*80)
print("🎨 Creating Visual Map")
print("="*80)
print()

# Simple chunking (split into ~500 char chunks at sentence boundaries)
def simple_chunk(text, max_length=500):
    """Simple sentence-aware chunking."""
    sentences = re.split(r'([.!?]+\s+)', text)
    chunks = []
    current_chunk = ""
    
    for i in range(0, len(sentences), 2):
        sent = sentences[i]
        delim = sentences[i+1] if i+1 < len(sentences) else ""
        full_sent = sent + delim
        
        if len(current_chunk) + len(full_sent) > max_length and current_chunk:
            chunks.append(current_chunk.strip())
            current_chunk = full_sent
        else:
            current_chunk += full_sent
    
    if current_chunk.strip():
        chunks.append(current_chunk.strip())
    
    return chunks

chunks = simple_chunk(REFERENCE_TEXT, max_length=500)

print(f"📝 Text chunked into {len(chunks)} chunks")
print()

# Transcribe
print("🎧 Transcribing...")
with open(AUDIO_FILE, 'rb') as f:
    transcription = client.audio.transcriptions.create(
        model="whisper-1",
        file=f,
        response_format="verbose_json"
    )

transcript_words = []
if hasattr(transcription, 'words') and transcription.words:
    transcript_words = transcription.words
transcript_text = transcription.text.lower()
trans_word_set = set(re.findall(r'\b\w+\b', transcript_text))

if transcript_words:
    print(f"✅ Done! {len(transcript_words)} words with timestamps")
else:
    print(f"✅ Done! Transcribed (no word-level timestamps available)")
print()

# Analyze each chunk
print("🔍 Analyzing chunks...")

def clean_text(text):
    return re.sub(r'<[^>]+>', '', text).strip()

def get_words_in_chunk(chunk):
    clean = clean_text(chunk)
    words = []
    for match in re.finditer(r'\b[\w]+\b', clean.lower()):
        words.append(match.group())
    return words

# Estimate timestamps for chunks (based on audio duration)
audio_duration = (os.path.getsize(AUDIO_FILE) - 44) / (24000 * 2)
total_chars = sum(len(c) for c in chunks)
current_time = 0

chunk_data = []
for i, chunk in enumerate(chunks):
    chunk_words = get_words_in_chunk(chunk)
    
    # Estimate chunk duration
    chunk_duration = (len(chunk) / total_chars) * audio_duration
    start_time = current_time
    end_time = current_time + chunk_duration
    current_time = end_time
    
    # Find missing words
    missing_words = [w for w in chunk_words if w not in trans_word_set]
    found_words = [w for w in chunk_words if w in trans_word_set]
    
    chunk_data.append({
        'index': i + 1,
        'text': chunk,
        'clean': clean_text(chunk),
        'words': chunk_words,
        'missing': missing_words,
        'found': found_words,
        'start_time': start_time,
        'end_time': end_time,
        'duration': chunk_duration,
        'accuracy': len(found_words) / len(chunk_words) * 100 if chunk_words else 0
    })

print(f"✅ Analyzed {len(chunk_data)} chunks")
print()

# Generate HTML
print("🎨 Generating HTML visualization...")

html = """<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>TTS Chunking Visual Map</title>
    <style>
        body {
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            margin: 0;
            padding: 20px;
            background: #f5f5f5;
        }
        .container {
            max-width: 1400px;
            margin: 0 auto;
            background: white;
            padding: 30px;
            border-radius: 10px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        }
        h1 {
            color: #333;
            border-bottom: 3px solid #4CAF50;
            padding-bottom: 10px;
        }
        .summary {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
            gap: 20px;
            margin: 20px 0;
        }
        .stat-box {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 20px;
            border-radius: 8px;
            text-align: center;
        }
        .stat-box h3 {
            margin: 0 0 10px 0;
            font-size: 14px;
            opacity: 0.9;
        }
        .stat-box .value {
            font-size: 32px;
            font-weight: bold;
        }
        .chunk {
            margin: 30px 0;
            border: 2px solid #ddd;
            border-radius: 8px;
            overflow: hidden;
        }
        .chunk-header {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 15px 20px;
            display: flex;
            justify-content: space-between;
            align-items: center;
        }
        .chunk-title {
            font-weight: bold;
            font-size: 18px;
        }
        .chunk-stats {
            display: flex;
            gap: 20px;
            font-size: 14px;
        }
        .chunk-body {
            padding: 20px;
            background: #fafafa;
        }
        .chunk-text {
            line-height: 1.8;
            font-size: 16px;
            margin: 10px 0;
        }
        .emotion-tag {
            background: #e3f2fd;
            color: #1976d2;
            padding: 2px 8px;
            border-radius: 4px;
            font-weight: 600;
            font-size: 14px;
        }
        .word {
            display: inline;
        }
        .word.found {
            color: #2e7d32;
        }
        .word.missing {
            background: #ffebee;
            color: #c62828;
            padding: 2px 4px;
            border-radius: 3px;
            font-weight: 600;
            text-decoration: underline wavy #c62828;
        }
        .accuracy-bar {
            height: 30px;
            background: #e0e0e0;
            border-radius: 15px;
            overflow: hidden;
            margin: 15px 0;
        }
        .accuracy-fill {
            height: 100%;
            background: linear-gradient(90deg, #4CAF50 0%, #8BC34A 100%);
            display: flex;
            align-items: center;
            justify-content: center;
            color: white;
            font-weight: bold;
            transition: width 0.5s;
        }
        .missing-words-box {
            background: #fff3e0;
            border-left: 4px solid #ff9800;
            padding: 15px;
            margin: 15px 0;
            border-radius: 4px;
        }
        .missing-words-box h4 {
            margin: 0 0 10px 0;
            color: #e65100;
        }
        .missing-word-tag {
            display: inline-block;
            background: #ffebee;
            color: #c62828;
            padding: 5px 10px;
            margin: 3px;
            border-radius: 4px;
            font-weight: 600;
        }
        .timeline {
            background: #f5f5f5;
            padding: 10px;
            border-radius: 4px;
            font-family: monospace;
            color: #555;
        }
        .legend {
            display: flex;
            gap: 20px;
            margin: 20px 0;
            padding: 15px;
            background: #f5f5f5;
            border-radius: 8px;
        }
        .legend-item {
            display: flex;
            align-items: center;
            gap: 8px;
        }
        .legend-color {
            width: 20px;
            height: 20px;
            border-radius: 4px;
        }
        .pattern-analysis {
            background: #e8f5e9;
            border: 2px solid #4CAF50;
            padding: 20px;
            border-radius: 8px;
            margin: 20px 0;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>🎨 TTS Chunking Visual Map - Word-Level Analysis</h1>
        
        <div class="summary">
            <div class="stat-box">
                <h3>Total Chunks</h3>
                <div class="value">""" + str(len(chunks)) + """</div>
            </div>
            <div class="stat-box">
                <h3>Audio Duration</h3>
                <div class="value">""" + f"{audio_duration:.1f}s" + """</div>
            </div>
            <div class="stat-box">
                <h3>Total Words</h3>
                <div class="value">""" + str(sum(len(c['words']) for c in chunk_data)) + """</div>
            </div>
            <div class="stat-box">
                <h3>Missing Words</h3>
                <div class="value">""" + str(sum(len(c['missing']) for c in chunk_data)) + """</div>
            </div>
        </div>
        
        <div class="legend">
            <div class="legend-item">
                <div class="legend-color" style="background: #2e7d32;"></div>
                <span>Found in transcript</span>
            </div>
            <div class="legend-item">
                <div class="legend-color" style="background: #ffebee; border: 2px solid #c62828;"></div>
                <span>Missing from transcript</span>
            </div>
            <div class="legend-item">
                <div class="legend-color" style="background: #e3f2fd;"></div>
                <span>Emotion tag</span>
            </div>
        </div>
"""

# Add each chunk
for chunk_info in chunk_data:
    # Highlight missing words in text
    clean = chunk_info['clean']
    words = chunk_info['words']
    missing = set(chunk_info['missing'])
    
    # Create highlighted version
    highlighted = clean
    for word in words:
        if word in missing:
            # Use word boundary to avoid partial matches
            highlighted = re.sub(
                rf'\b{re.escape(word)}\b',
                f'<span class="word missing">{word}</span>',
                highlighted,
                count=1
            )
        else:
            highlighted = re.sub(
                rf'\b{re.escape(word)}\b',
                f'<span class="word found">{word}</span>',
                highlighted,
                count=1
            )
    
    # Show emotion tags
    text_with_tags = chunk_info['text']
    for match in re.finditer(r'<([^>]+)>', text_with_tags):
        emotion = match.group(1)
        text_with_tags = text_with_tags.replace(
            f'<{emotion}>',
            f'<span class="emotion-tag">{emotion}</span>'
        )
    
    accuracy = chunk_info['accuracy']
    missing_count = len(chunk_info['missing'])
    
    html += f"""
        <div class="chunk">
            <div class="chunk-header">
                <div class="chunk-title">Chunk {chunk_info['index']} / {len(chunks)}</div>
                <div class="chunk-stats">
                    <span>⏱️ {chunk_info['start_time']:.1f}s - {chunk_info['end_time']:.1f}s</span>
                    <span>📝 {len(chunk_info['words'])} words</span>
                    <span>❌ {missing_count} missing</span>
                </div>
            </div>
            <div class="chunk-body">
                <div class="accuracy-bar">
                    <div class="accuracy-fill" style="width: {accuracy}%">
                        {accuracy:.1f}% Accuracy
                    </div>
                </div>
                
                <div style="margin: 15px 0;">
                    <strong>With emotion tags:</strong><br>
                    {text_with_tags}
                </div>
                
                <div style="margin: 15px 0;">
                    <strong>Word-level analysis:</strong><br>
                    <div class="chunk-text">{highlighted}</div>
                </div>
    """
    
    if missing_count > 0:
        html += f"""
                <div class="missing-words-box">
                    <h4>❌ Missing Words ({missing_count}):</h4>
                    {''.join(f'<span class="missing-word-tag">{w}</span>' for w in chunk_info['missing'])}
                </div>
        """
    
    html += """
            </div>
        </div>
    """

# Add pattern analysis
best_chunks = sorted(chunk_data, key=lambda x: x['accuracy'], reverse=True)[:3]
worst_chunks = sorted(chunk_data, key=lambda x: x['accuracy'])[:3]

html += """
        <div class="pattern-analysis">
            <h2>📊 Pattern Analysis</h2>
            
            <h3>✅ Best Performing Chunks (Highest Accuracy):</h3>
            <ul>
"""

for chunk in best_chunks:
    html += f"""
                <li>Chunk {chunk['index']}: {chunk['accuracy']:.1f}% accuracy 
                    ({chunk['start_time']:.1f}s - {chunk['end_time']:.1f}s)</li>
    """

html += """
            </ul>
            
            <h3>⚠️ Chunks with Most Missing Words:</h3>
            <ul>
"""

for chunk in worst_chunks:
    html += f"""
                <li>Chunk {chunk['index']}: {chunk['accuracy']:.1f}% accuracy, 
                    {len(chunk['missing'])} missing words 
                    ({chunk['start_time']:.1f}s - {chunk['end_time']:.1f}s)</li>
    """

html += """
            </ul>
        </div>
    </div>
</body>
</html>
"""

# Save HTML
output_file = "visual_chunking_map.html"
with open(output_file, 'w', encoding='utf-8') as f:
    f.write(html)

print(f"✅ Visual map created: {output_file}")
print()
print(f"📊 Summary:")
print(f"   Total chunks: {len(chunks)}")
print(f"   Total words: {sum(len(c['words']) for c in chunk_data)}")
print(f"   Missing words: {sum(len(c['missing']) for c in chunk_data)}")
print(f"   Average accuracy: {sum(c['accuracy'] for c in chunk_data)/len(chunk_data):.1f}%")
print()
print(f"🌐 Open in browser: file://{os.path.abspath(output_file)}")
print()

