#!/usr/bin/env python3
"""
Token-level transcript viewer for deception detection experiments.
Shows exactly which parts of conversations (system/user/assistant) are being evaluated.
"""

import json
from pathlib import Path
from typing import List, Any
from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.models import get_model_and_tokenizer


def create_token_viewer_html(experiment_path: str, max_samples: int = 10):
    """Create an HTML viewer showing token-level scores and message detection flags."""
    
    exp_path = Path(experiment_path)
    experiment = Experiment(ExperimentConfig.from_path(exp_path))
    
    # Load scores data
    with open(exp_path / 'scores.json') as f:
        scores_data = json.load(f)
    
    print(f"Available datasets in scores: {list(scores_data.keys())}")
    
    # Get the evaluated dataset
    dataset_name = 'upper_bound_combined__plain'
    if dataset_name not in scores_data:
        print(f"Dataset {dataset_name} not found in scores. Using first available.")
        dataset_name = list(scores_data.keys())[0]
    
    dataset_scores, dataset_labels = scores_data[dataset_name]
    print(f"Found {len(dataset_scores)} samples with scores")
    
    # Find the matching dataset in the experiment
    eval_dataset = None
    for dataset in experiment.eval_datasets:
        if dataset.name == dataset_name:
            eval_dataset = dataset
            break
    
    if eval_dataset is None:
        print(f"Could not find dataset {dataset_name} in experiment eval datasets")
        return
    
    # Create HTML template
    html_template = '''
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Token-Level Transcript Viewer</title>
        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css" rel="stylesheet">
        <style>
            .message-box {{ 
                border: 1px solid #ddd; 
                margin: 10px 0; 
                padding: 15px; 
                border-radius: 8px; 
            }}
            .system-msg {{ background-color: #f8f9fa; border-left: 4px solid #6c757d; }}
            .user-msg {{ background-color: #e3f2fd; border-left: 4px solid #2196f3; }}
            .assistant-msg {{ background-color: #f3e5f5; border-left: 4px solid #9c27b0; }}
            .detect-enabled {{ 
                background-color: #fff3cd !important; 
                border: 2px solid #856404 !important;
                box-shadow: 0 0 5px rgba(133,100,4,0.3);
            }}
            .detect-disabled {{ opacity: 0.6; }}
            .score-info {{ 
                font-size: 0.9em; 
                color: #666; 
                margin-top: 10px; 
                padding: 8px; 
                background-color: #f8f9fa; 
                border-radius: 4px;
            }}
            .sample-header {{
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                padding: 20px;
                margin: 20px 0;
                border-radius: 10px;
            }}
            .label-badge {{
                font-size: 1.1em;
                padding: 8px 16px;
            }}
            .honest {{ color: #28a745; }}
            .deceptive {{ color: #dc3545; }}
            .token-score {{ 
                display: inline-block; 
                padding: 2px 4px; 
                margin: 1px; 
                border-radius: 3px; 
                font-family: monospace; 
                font-size: 0.9em;
            }}
            .positive-score {{ background-color: #ffebee; color: #c62828; }}
            .negative-score {{ background-color: #e8f5e8; color: #2e7d32; }}
            .neutral-score {{ background-color: #f5f5f5; color: #666; }}
        </style>
    </head>
    <body>
        <div class="container-fluid mt-4">
            <h1 class="text-center mb-4">🔍 Token-Level Transcript Viewer</h1>
            
            <div class="alert alert-info">
                <h5>📋 Detection Configuration</h5>
                <p><strong>What This Shows:</strong> Each conversation sample with message-level detection flags and token scores.</p>
                <ul class="mb-0">
                    <li><strong>Yellow highlighted boxes</strong> = Messages being evaluated (detect=True)</li>
                    <li><strong>Faded boxes</strong> = Messages not being evaluated (detect=False)</li>
                    <li><strong>Token scores</strong> = Individual token contributions to deception detection</li>
                    <li><strong>Positive scores (red)</strong> = Tokens suggesting deception</li>
                    <li><strong>Negative scores (green)</strong> = Tokens suggesting honesty</li>
                </ul>
            </div>
            
            {samples_html}
        </div>
        
        <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/js/bootstrap.bundle.min.js"></script>
    </body>
    </html>
    '''
    
    # Get tokenizer for the model
    try:
        _, tokenizer = get_model_and_tokenizer(experiment.cfg.model_name, device="cpu")
    except Exception as e:
        print(f"Could not load tokenizer: {e}")
        tokenizer = None
    
    samples_html = ""
    
    # Process each sample (limit to max_samples or available samples)
    num_samples = min(max_samples, len(dataset_scores) if isinstance(dataset_scores, list) else 0, len(eval_dataset.dialogues))
    
    for i in range(num_samples):
        dialogue = eval_dataset.dialogues[i]
        label = eval_dataset.labels[i]
        sample_scores = dataset_scores[i] if isinstance(dataset_scores, list) and i < len(dataset_scores) else []
        
        label_class = 'honest' if label.value == 'honest' else 'deceptive'
        label_color = 'success' if label.value == 'honest' else 'danger'
        
        # Sample header
        samples_html += f'''
        <div class="sample-header">
            <div class="d-flex justify-content-between align-items-center">
                <h3>Sample {i+1}</h3>
                <span class="badge bg-{label_color} label-badge">{label.value.upper()}</span>
            </div>
            <div class="mt-2">
                <small>Overall Score: {sample_scores[0] if sample_scores and len(sample_scores) > 0 else 'N/A'}</small>
                <small class="ms-3">Total Tokens with Scores: {len(sample_scores) if isinstance(sample_scores, list) else 0}</small>
            </div>
        </div>
        '''
        
        # Process each message in the dialogue
        for msg_idx, message in enumerate(dialogue):
            role = message.role
            content = message.content
            detect_flag = message.detect
            
            # Determine message styling
            role_class = f"{role}-msg"
            detect_class = "detect-enabled" if detect_flag else "detect-disabled"
            
            # Create message content with token highlighting if scores available
            if detect_flag and isinstance(sample_scores, list) and len(sample_scores) > 0 and tokenizer:
                # Tokenize the message content
                try:
                    tokens = tokenizer.tokenize(content)
                    
                    # Create token-by-token display with scores
                    content_with_scores = ""
                    if len(tokens) <= len(sample_scores):
                        for j, token in enumerate(tokens):
                            if j < len(sample_scores):
                                score = sample_scores[j] if isinstance(sample_scores[j], (int, float)) else 0
                                score_class = "positive-score" if score > 0.1 else "negative-score" if score < -0.1 else "neutral-score"
                                content_with_scores += f'<span class="token-score {score_class}" title="Score: {score:.3f}">{token}</span> '
                            else:
                                content_with_scores += f'<span class="token-score neutral-score">{token}</span> '
                    else:
                        content_with_scores = content  # Fallback if token counts don't match
                except Exception:
                    content_with_scores = content  # Fallback if tokenization fails
            else:
                content_with_scores = content
            
            samples_html += f'''
            <div class="message-box {role_class} {detect_class}">
                <div class="d-flex justify-content-between align-items-start mb-2">
                    <strong>{role.upper()}</strong>
                    <div>
                        <span class="badge {'bg-warning text-dark' if detect_flag else 'bg-secondary'}">
                            {'EVALUATING' if detect_flag else 'NOT EVALUATING'}
                        </span>
                    </div>
                </div>
                <div class="content">
                    {content_with_scores}
                </div>
                <div class="score-info">
                    <strong>Detection Status:</strong> {'✅ This message IS being evaluated for deception' if detect_flag else '❌ This message is NOT being evaluated'}
                    {f'<br><strong>Message Length:</strong> ~{len(content.split())} words' if content else ''}
                </div>
            </div>
            '''
        
        samples_html += "<hr class='my-4'>"
    
    # Create final HTML
    final_html = html_template.format(samples_html=samples_html)
    
    # Save to file
    output_dir = Path('html-outputs')
    output_dir.mkdir(exist_ok=True)
    output_file = output_dir / 'token_transcript_viewer.html'
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(final_html)
    
    print(f"✅ Token transcript viewer created: {output_file.absolute()}")
    print(f"📊 Analyzed {num_samples} samples")
    print(f"🎯 Open the file in your browser to see token-level evaluation details")
    
    return output_file


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python token_transcript_viewer.py <experiment_path> [max_samples]")
        print("Example: python token_transcript_viewer.py results/train_on_all_XYZ_lr__repe_honesty__plain__mistral-small-3.1-24b__20250712_122604")
        sys.exit(1)
    
    experiment_path = sys.argv[1]
    max_samples = int(sys.argv[2]) if len(sys.argv) > 2 else 10
    
    create_token_viewer_html(experiment_path, max_samples) 