#!/usr/bin/env python3

import pickle
from pathlib import Path
from tqdm import tqdm
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.token_html import get_token_scores_html
from deception_detection.results.utils import get_scores_from_name

def get_toggled_html_all_tokens(masked_html: str, unmasked_html: str) -> str:
    """Modified toggle to show unmasked (all tokens) by default"""
    return f"""
        <script>
        function toggleMask() {{
        var masked = document.getElementById('masked');
        var unmasked = document.getElementById('unmasked');
            if (document.getElementById('maskToggle').checked) {{
            masked.style.display = 'block';
            unmasked.style.display = 'none';
            }} else {{
            masked.style.display = 'none';
            unmasked.style.display = 'block';
            }}
        }}
        </script>
        <div class="form-check mb-3">
            <input class="form-check-input" type="checkbox" id="maskToggle" 
                onchange="toggleMask()">
            <label class="form-check-label" for="maskToggle">Show scores for evaluation
            tokens only (original behavior)</label>
        </div>
        <div id="masked" style="display: none">{masked_html}</div>
        <div id="unmasked">{unmasked_html}</div>"""

def get_full_html(title: str, content: str, subtitle: str = "") -> str:
    """Generate complete HTML page"""
    return f"""
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>{title}</title>
        <style>
            body {{
                max-width: 1200px;
                margin: 0 auto;
                padding: 20px;
                background: #f8f9fa;
                font-family: Arial, sans-serif;
                line-height: 1.6;
            }}
            .token {{
                display: inline-block;
                padding: 2px;
                margin: 1px;
            }}
            h1 {{ color: #333; }}
            h2 {{ color: #666; }}
        </style>
    </head>
    <body>
        <h1>{title}</h1>
        {f"<h2>{subtitle}</h2>" if subtitle else ""}
        {content}
    </body>
    </html>
    """

def export_all_tokens_html():
    """Export HTML with all tokens highlighted by default"""
    
    # Load experiment
    with open('experiment.pkl', 'rb') as inp:
        experiment = pickle.load(inp)
    
    # Find the upper_bound dataset
    upper_bound_dataset = None
    for dataset in experiment.eval_datasets:
        if "upper_bound_combined" in dataset.name:
            upper_bound_dataset = dataset
            break
    
    if not upper_bound_dataset:
        print("No upper_bound_combined dataset found!")
        return
    
    print(f"Found dataset: {upper_bound_dataset.name}")
    print(f"Number of samples: {len(upper_bound_dataset.dialogues)}")
    
    # Create output directory
    output_dir = Path("html_all_tokens")
    output_dir.mkdir(exist_ok=True)
    
    # Create scorer
    scorer = create_prompt_scorer("mean", experiment, [0.01], default_fpr=0.01, 
                                 control_dataset_name="alpaca_XYZ__plain")
    
    # Get tokenized dataset and scores
    tok_dataset = TokenizedDataset.from_dataset(upper_bound_dataset, experiment.tokenizer)
    all_token_scores = get_scores_from_name(experiment, upper_bound_dataset.name, all_tokens=True)
    normalized_scores_array = scorer.normalize_scores(all_token_scores)
    
    # Create index HTML
    index_html = "<h2>Upper Bound Combined Dataset - All Tokens Highlighted</h2>\n"
    index_html += "<ul>\n"
    
    # Export first 20 samples from each label
    honest_samples = []
    deceptive_samples = []
    
    for i, label in enumerate(all_token_scores.labels):
        if label.value == "honest" and len(honest_samples) < 20:
            honest_samples.append(i)
        elif label.value == "deceptive" and len(deceptive_samples) < 20:
            deceptive_samples.append(i)
    
    samples_to_export = honest_samples + deceptive_samples
    
    for i in tqdm(samples_to_export, desc="Exporting samples"):
        # Get scores and tokens
        normalized_scores = normalized_scores_array[i]
        str_tokens = tok_dataset.str_tokens[i]
        
        # Create masked scores (original behavior)
        if tok_dataset.detection_mask is not None:
            detection_mask = tok_dataset.detection_mask[i]
            masked_scores = normalized_scores * detection_mask[-len(normalized_scores):]
        else:
            masked_scores = normalized_scores
            
        # Generate HTML
        masked_html = get_token_scores_html(str_tokens, masked_scores, extent=6.0)
        unmasked_html = get_token_scores_html(str_tokens, normalized_scores, extent=6.0)
        
        # Create toggle HTML (unmasked by default)
        toggled_html = get_toggled_html_all_tokens(masked_html, unmasked_html)
        
        # Get label and score
        label = all_token_scores.labels[i].value
        score = scorer.reduce(all_token_scores.scores[i:i+1])[0]
        
        # Create full HTML
        subtitle = f"Score: {score:.2f}, Label: {label.title()}"
        full_html = get_full_html(
            title=f"upper_bound_combined__plain sample {i}",
            content=toggled_html,
            subtitle=subtitle
        )
        
        # Save file
        filename = f"sample_{i}_{label}.html"
        with open(output_dir / filename, 'w') as f:
            f.write(full_html)
        
        # Add to index
        index_html += f'<li><a href="{filename}">Sample {i} ({label.title()}, score: {score:.2f})</a></li>\n'
    
    index_html += "</ul>\n"
    
    # Save index
    index_full_html = get_full_html(
        title="All Tokens HTML Export - Upper Bound Combined",
        content=index_html
    )
    
    with open(output_dir / "index.html", 'w') as f:
        f.write(index_full_html)
    
    print(f"\n✅ HTML export completed!")
    print(f"📁 Files saved to: {output_dir.absolute()}")
    print(f"🌐 Open in browser: {(output_dir / 'index.html').absolute()}")

if __name__ == "__main__":
    export_all_tokens_html() 