cleanup_tool.py

#!/usr/bin/env python3
"""Main script for BibTeX cleanup tool."""

import argparse
import sys
import os
from pathlib import Path
from typing import List, Dict, Optional
import tempfile
import subprocess

from bibtex_parser import BibTeXParser, BibEntry
from citation_scanner import CitationScanner
from bibtex_formatter import BibTeXFormatter
from interactive_cli import InteractiveCLI


class BibTeXCleanupTool:
    """Main class for the BibTeX cleanup tool."""
    
    def __init__(self, tex_file: str, bib_file: str, output_file: str = None, non_interactive: bool = False):
        self.tex_file = tex_file
        self.bib_file = bib_file
        self.output_file = output_file or bib_file.replace('.bib', '_cleaned.bib')
        self.non_interactive = non_interactive
        
        self.parser = BibTeXParser()
        self.scanner = CitationScanner()
        self.formatter = BibTeXFormatter()
        self.cli = InteractiveCLI()
        
        self.processed_entries = []
        self.errors = []
    
    def run(self):
        """Run the cleanup tool."""
        try:
            # Display welcome
            self.cli.display_welcome()
            self.cli.display_file_info(self.tex_file, self.bib_file, self.output_file)
            
            # Step 1: Scan TeX file for citations
            print(f"\nScanning {self.tex_file} for citations...")
            citations = self.scanner.scan_tex_file(self.tex_file)
            
            # Step 2: Parse BibTeX file
            print(f"Parsing {self.bib_file}...")
            entries = self.parser.parse_file(self.bib_file)
            
            # Step 3: Check for undefined citations
            undefined = self.scanner.find_undefined_citations(
                self.tex_file, 
                self.parser.entry_map
            )
            
            self.cli.display_citations_found(len(citations), undefined)
            
            # Step 4: Process only cited entries
            cited_entries = []
            for citation in citations:
                if citation in self.parser.entry_map:
                    cited_entries.append(self.parser.entry_map[citation])
            
            if not cited_entries:
                print("\nNo cited entries found to process.")
                return
            
            # Ask user if they want to proceed (skip in non-interactive mode)
            if not self.non_interactive:
                if not self.cli.prompt_yes_no(f"\nProcess {len(cited_entries)} cited entries?"):
                    print("Cleanup cancelled.")
                    return
            else:
                print(f"\nProcessing {len(cited_entries)} cited entries in non-interactive mode...")
            
            # Step 5: Process each entry (interactively or automatically)
            if self.non_interactive:
                self.process_entries_automatically(cited_entries)
            else:
                self.process_entries_interactively(cited_entries)
            
            # Step 6: Save results
            if self.processed_entries:
                self.save_results()
            
            # Display summary
            self.cli.display_summary(
                len(cited_entries),
                len(self.processed_entries),
                self.errors
            )
            
        except KeyboardInterrupt:
            print("\n\nCleanup interrupted by user.")
            if self.processed_entries:
                if self.non_interactive or self.cli.prompt_yes_no("Save progress so far?"):
                    self.save_results()
        except Exception as e:
            print(f"\nError: {e}")
            self.errors.append(str(e))
    
    def process_entries_interactively(self, entries: List[BibEntry]):
        """Process entries one by one with user interaction."""
        total = len(entries)
        
        for index, entry in enumerate(entries, 1):
            # Clear screen for better readability (optional)
            # self.cli.clear_screen()
            
            # Display header
            self.cli.display_entry_header(entry.key, index, total)
            
            # Format the entry
            formatted_entry, suggestions = self.formatter.format_entry(entry)
            
            # Check if there are any changes
            original_text = entry.to_bibtex()
            formatted_text = formatted_entry.to_bibtex()
            
            if original_text == formatted_text:
                print("No changes needed for this entry.")
                self.processed_entries.append(entry)
                continue
            
            # Display suggestions
            self.cli.display_suggestions(suggestions)
            
            # Check if URL is missing and try to find it
            if 'url' not in formatted_entry.fields and not self.non_interactive:
                # Check if "Consider adding URL" is in suggestions
                if any("Consider adding URL" in s for s in suggestions):
                    print("\n🔍 This entry is missing a URL. Searching...")
                    
                    # Try multiple methods to find URL
                    url_found = False
                    
                    # Method 1: Check for DOI
                    doi = entry.fields.get('doi', '')
                    if doi:
                        doi = doi.replace('\\url{', '').replace('}', '').strip()
                        if not doi.startswith('http'):
                            url = f"https://doi.org/{doi}"
                        else:
                            url = doi
                        print(f"   Found DOI URL: {url[:80]}...")
                        if self.cli.prompt_yes_no("   Add this URL?"):
                            formatted_entry.fields['url'] = url
                            suggestions.append(f"Added URL from DOI: {url}")
                            formatted_text = formatted_entry.to_bibtex()
                            url_found = True
                    
                    # Method 2: Check for arXiv
                    if not url_found:
                        eprint = entry.fields.get('eprint', '')
                        journal = entry.fields.get('journal', '')
                        
                        arxiv_id = None
                        if eprint:
                            arxiv_id = eprint.replace('arXiv:', '').strip()
                        elif 'arxiv' in journal.lower():
                            import re
                            match = re.search(r'(\d{4}\.\d{4,5})', journal)
                            if match:
                                arxiv_id = match.group(1)
                        
                        if arxiv_id:
                            url = f"https://arxiv.org/abs/{arxiv_id}"
                            print(f"   Found arXiv URL: {url}")
                            if self.cli.prompt_yes_no("   Add this URL?"):
                                formatted_entry.fields['url'] = url
                                suggestions.append(f"Added URL from arXiv: {url}")
                                formatted_text = formatted_entry.to_bibtex()
                                url_found = True
                    
                    # Method 3: Try Claude API if available
                    if not url_found:
                        try:
                            from claude_url_finder import ClaudeURLFinder
                            claude_finder = ClaudeURLFinder()
                            if claude_finder.initialized:
                                result = claude_finder.find_paper_url(entry)
                                if result:
                                    url, confidence, source = result
                                    if confidence >= 0.7:
                                        print(f"   Found {source} URL (confidence: {confidence:.0%})")
                                        if self.cli.prompt_yes_no(f"   Add URL: {url[:80]}...?"):
                                            formatted_entry.fields['url'] = url
                                            suggestions.append(f"Added URL from {source}: {url}")
                                            formatted_text = formatted_entry.to_bibtex()
                                            url_found = True
                        except ImportError:
                            pass
                    
                    # Method 4: Manual search suggestion
                    if not url_found:
                        print("   Could not find URL automatically.")
                        print("   Suggested search query:")
                        title = entry.fields.get('title', '')
                        import re
                        clean_title = re.sub(r'\{([^}]+)\}', r'\1', title)
                        clean_title = re.sub(r'\\[a-zA-Z]+', '', clean_title).strip()
                        print(f'   "{clean_title}" {entry.fields.get("year", "")}')
            
            # Display diff
            self.cli.display_diff(original_text, formatted_text)
            
            # Process user choice
            while True:
                action = self.cli.prompt_user_action()
                
                if action == 'a':  # Accept
                    self.processed_entries.append(formatted_entry)
                    self.cli.changes_made += 1
                    print("Changes accepted.")
                    break
                    
                elif action == 's':  # Skip
                    self.processed_entries.append(entry)
                    self.cli.changes_skipped += 1
                    print("Entry skipped.")
                    break
                    
                elif action == 'e':  # Edit manually
                    edited = self.edit_entry_in_editor(formatted_text)
                    if edited:
                        # Parse the edited entry
                        temp_entries = self.parser.parse_string(edited)
                        if temp_entries:
                            self.processed_entries.append(temp_entries[0])
                            self.cli.changes_made += 1
                            print("Manual edits accepted.")
                            break
                    else:
                        print("Edit cancelled.")
                        
                elif action == 'd':  # Show detailed diff
                    self.cli.display_diff(original_text, formatted_text)
                    
                elif action == 'v':  # View side-by-side
                    self.cli.display_side_by_side(original_text, formatted_text)
                    
                elif action == 'q':  # Quit and save
                    if self.cli.prompt_yes_no("Save progress and quit?"):
                        return
                    
                elif action == 'x':  # Exit without saving
                    if self.cli.prompt_yes_no("Exit without saving?"):
                        sys.exit(0)
            
            # Display progress
            self.cli.display_progress(index, total)
    
    def process_entries_automatically(self, entries: List[BibEntry]):
        """Process entries automatically without user interaction."""
        total = len(entries)
        
        for index, entry in enumerate(entries, 1):
            print(f"\nProcessing entry {index}/{total}: {entry.key}")
            
            # Format the entry
            formatted_entry, suggestions = self.formatter.format_entry(entry)
            
            # Try to find missing URLs in non-interactive mode
            if 'url' not in formatted_entry.fields:
                import re
                
                # Check for DOI
                doi = entry.fields.get('doi', '')
                if doi:
                    doi = doi.replace('\\url{', '').replace('}', '').strip()
                    if not doi.startswith('http'):
                        url = f"https://doi.org/{doi}"
                    else:
                        url = doi
                    print(f"  Found DOI URL: {url[:60]}...")
                    formatted_entry.fields['url'] = url
                    suggestions.append("Added URL from DOI")
                
                # Check for arXiv
                elif entry.fields.get('eprint') or 'arxiv' in entry.fields.get('journal', '').lower():
                    eprint = entry.fields.get('eprint', '')
                    journal = entry.fields.get('journal', '')
                    
                    arxiv_id = None
                    if eprint:
                        arxiv_id = eprint.replace('arXiv:', '').strip()
                    elif 'arxiv' in journal.lower():
                        match = re.search(r'(\d{4}\.\d{4,5})', journal)
                        if match:
                            arxiv_id = match.group(1)
                    
                    if arxiv_id:
                        url = f"https://arxiv.org/abs/{arxiv_id}"
                        print(f"  Found arXiv URL: {url}")
                        formatted_entry.fields['url'] = url
                        suggestions.append("Added URL from arXiv")
                
                # Try Claude Code WebSearch if available
                else:
                    try:
                        from claude_code_integration import find_url_with_claude_code, IN_CLAUDE_CODE
                        if IN_CLAUDE_CODE:
                            result = find_url_with_claude_code(entry)
                            if result:
                                url, confidence, source = result
                                if confidence >= 0.8:
                                    print(f"  Found {source} URL via Claude Code: {url[:60]}...")
                                    formatted_entry.fields['url'] = url
                                    suggestions.append(f"Added URL from {source}")
                    except ImportError:
                        pass
                    
                    # Fallback to Claude API if available
                    if 'url' not in formatted_entry.fields:
                        try:
                            from claude_url_finder import ClaudeURLFinder
                            url_finder = ClaudeURLFinder()
                            if url_finder.initialized:
                                result = url_finder.find_paper_url(entry)
                                if result:
                                    url, confidence, source = result
                                    if confidence >= 0.8:  # Higher threshold for automatic mode
                                        print(f"  Found {source} URL (confidence: {confidence:.0%})")
                                        formatted_entry.fields['url'] = url
                                        suggestions.append(f"Added URL from {source}")
                        except ImportError:
                            pass
            
            # Check if there are any changes
            original_text = entry.to_bibtex()
            formatted_text = formatted_entry.to_bibtex()
            
            if original_text == formatted_text:
                print(f"  No changes needed for {entry.key}")
                self.processed_entries.append(entry)
            else:
                # Automatically accept all changes
                print(f"  Applying formatting changes to {entry.key}")
                if suggestions:
                    for suggestion in suggestions:
                        print(f"    - {suggestion}")
                self.processed_entries.append(formatted_entry)
                self.cli.changes_made += 1
            
            # Display progress
            if index % 10 == 0 or index == total:
                print(f"Progress: {index}/{total} entries processed")
    
    def edit_entry_in_editor(self, entry_text: str) -> Optional[str]:
        """Open entry in text editor for manual editing."""
        editor = os.environ.get('EDITOR', 'nano')  # Default to nano
        
        try:
            with tempfile.NamedTemporaryFile(mode='w', suffix='.bib', delete=False) as f:
                f.write(entry_text)
                temp_path = f.name
            
            # Open in editor
            subprocess.call([editor, temp_path])
            
            # Read back the edited content
            with open(temp_path, 'r') as f:
                edited = f.read()
            
            # Clean up
            os.unlink(temp_path)
            
            # Validate the edited entry
            valid, errors = self.formatter.validate_bibtex(edited)
            if not valid:
                print(f"Invalid BibTeX format: {', '.join(errors)}")
                if self.cli.prompt_yes_no("Try editing again?"):
                    return self.edit_entry_in_editor(edited)
                return None
            
            return edited
            
        except Exception as e:
            print(f"Error editing entry: {e}")
            return None
    
    def save_results(self):
        """Save processed entries to output file."""
        try:
            with open(self.output_file, 'w', encoding='utf-8') as f:
                for i, entry in enumerate(self.processed_entries):
                    if i > 0:
                        f.write('\n\n')
                    f.write(entry.to_bibtex())
            
            print(f"\nResults saved to: {self.output_file}")
            
            # Validate the output file
            test_parser = BibTeXParser()
            test_entries = test_parser.parse_file(self.output_file)
            print(f"Output file validated: {len(test_entries)} entries parsed successfully")
            
        except Exception as e:
            print(f"Error saving results: {e}")
            self.errors.append(f"Save error: {e}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='Interactive BibTeX cleanup tool',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s paper.tex references.bib
  %(prog)s paper.tex references.bib -o cleaned_refs.bib
  %(prog)s --tex main.tex --bib refs.bib --output refs_clean.bib
        """
    )
    
    parser.add_argument('tex_file', nargs='?', help='TeX file containing citations')
    parser.add_argument('bib_file', nargs='?', help='BibTeX file to clean')
    parser.add_argument('-o', '--output', dest='output_file',
                       help='Output file for cleaned entries (default: <input>_cleaned.bib)')
    parser.add_argument('--tex', dest='tex_alt',
                       help='Alternative way to specify TeX file')
    parser.add_argument('--bib', dest='bib_alt',
                       help='Alternative way to specify BibTeX file')
    parser.add_argument('--non-interactive', action='store_true',
                       help='Run in non-interactive mode (automatically accept all changes)')
    
    args = parser.parse_args()
    
    # Determine input files
    tex_file = args.tex_file or args.tex_alt
    bib_file = args.bib_file or args.bib_alt
    
    if not tex_file or not bib_file:
        parser.print_help()
        sys.exit(1)
    
    # Check files exist
    if not os.path.exists(tex_file):
        print(f"Error: TeX file '{tex_file}' not found")
        sys.exit(1)
    
    if not os.path.exists(bib_file):
        print(f"Error: BibTeX file '{bib_file}' not found")
        sys.exit(1)
    
    # Run the tool
    tool = BibTeXCleanupTool(tex_file, bib_file, args.output_file, args.non_interactive)
    tool.run()


if __name__ == '__main__':
    main()