UmbraAudio/scripts/test_parser.py

#!/usr/bin/env python3
"""
Test script for document parser and entity extraction.
Tests the parsing of DealHunterShip manuscript.
"""

import json
import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from umbra_audio.agents.entity_extractor import (
    extract_characters_from_text,
    merge_characters,
)
from umbra_audio.core.document_parser import (
    extract_dialogue_lines,
    get_document_stats,
    get_section_content_for_audio,
    parse_manuscript,
    validate_manuscript,
)


def test_document_parser():
    """Test the document parser on DealHunterShip."""
    print("=" * 80)
    print("TESTING DOCUMENT PARSER")
    print("=" * 80)

    # Read the manuscript
    manuscript_path = Path(__file__).parent.parent / "DealHunterShip.txt"
    if not manuscript_path.exists():
        print(f"ERROR: Manuscript not found at {manuscript_path}")
        return False

    with open(manuscript_path, encoding='utf-8') as f:
        text = f.read()

    print(f"\nManuscript loaded: {len(text)} characters")

    # Validate manuscript
    print("\n1. Validating manuscript structure...")
    validation = validate_manuscript(text)
    print(f"   Valid: {validation['valid']}")
    print(f"   Errors: {validation['errors']}")
    print(f"   Warnings: {validation['warnings']}")
    print(f"   Stats: {json.dumps(validation['stats'], indent=2)}")

    # Parse manuscript
    print("\n2. Parsing manuscript...")
    manuscript = parse_manuscript(text)
    print(f"   Title: {manuscript.title}")
    print(f"   Author: {manuscript.author}")
    print(f"   Sections found: {len(manuscript.sections)}")

    # Analyze sections
    print("\n3. Section analysis:")
    chapter_count = 0
    scene_count = 0
    for section in manuscript.sections:
        if section.section_type == 'chapter':
            chapter_count += 1
            print(f"   Chapter {section.chapter_number}: {section.chapter_title[:50]}...")
            print(f"      Content length: {len(section.content)} chars")
            print(f"      Scenes: {section.scene_index}")
        elif section.section_type == 'scene':
            scene_count += 1

    print(f"\n   Total chapters: {chapter_count}")
    print(f"   Total scenes: {scene_count}")

    # Test audio content extraction
    print("\n4. Testing audio content extraction...")
    if manuscript.sections:
        first_section = manuscript.sections[0]
        audio_content = get_section_content_for_audio(first_section)
        print(f"   First section audio content length: {len(audio_content)} chars")
        print(f"   Preview: {audio_content[:200]}...")

    # Test dialogue extraction
    print("\n5. Testing dialogue extraction...")
    all_dialogues = []
    for section in manuscript.sections:
        dialogues = extract_dialogue_lines(section)
        all_dialogues.extend(dialogues)

    print(f"   Total dialogue lines extracted: {len(all_dialogues)}")

    # Show sample dialogues
    if all_dialogues:
        print("\n   Sample dialogues:")
        for d in all_dialogues[:5]:
            print(f"      Speaker: {d['speaker']}")
            print(f"      Dialogue: {d['dialogue'][:80]}...")
            print()

    # Get document stats
    print("\n6. Document statistics:")
    stats = get_document_stats(text)
    print(f"   Total characters: {stats['total_chars']}")
    print(f"   Total lines: {stats['total_lines']}")
    print(f"   Word count: {stats['word_count']}")
    print(f"   Chapters: {stats['chapter_count']}")
    print(f"   Scenes: {stats['scene_count']}")
    print(f"   Dialogue lines: {stats['dialogue_count']}")

    print("\n   Top character mentions:")
    for name, count in list(stats['character_mentions'].items())[:10]:
        print(f"      {name}: {count}")

    return validation['valid']


def test_entity_extraction():
    """Test entity extraction on DealHunterShip."""
    print("\n" + "=" * 80)
    print("TESTING ENTITY EXTRACTION")
    print("=" * 80)

    # Read the manuscript
    manuscript_path = Path(__file__).parent.parent / "DealHunterShip.txt"
    if not manuscript_path.exists():
        print(f"ERROR: Manuscript not found at {manuscript_path}")
        return False

    with open(manuscript_path, encoding='utf-8') as f:
        text = f.read()

    print(f"\nManuscript loaded: {len(text)} characters")

    # Extract characters
    print("\n1. Extracting characters...")
    characters = extract_characters_from_text(text)
    print(f"   Characters found: {len(characters)}")

    # Show top characters
    print("\n2. Top characters by dialogue count:")
    for char in characters[:15]:
        if char.dialogue_count > 0:
            print(f"\n   {char.canonical_name}:")
            print(f"      Dialogue lines: {char.dialogue_count}")
            if char.aliases:
                print(f"      Aliases: {', '.join(char.aliases[:3])}")
            if char.voice_clues:
                print(f"      Voice clues: {', '.join(char.voice_clues[:2])}")
            if char.sample_dialogues:
                print(f"      Sample: {char.sample_dialogues[0][:60]}...")

    # Merge characters
    print("\n3. Merging characters with aliases...")
    merged = merge_characters(characters)
    print(f"   Characters after merge: {len(merged)}")

    # Save results
    output_path = Path(__file__).parent.parent / "DealHunterShip_characters.json"
    output_data = {
        "book_title": "Deal Hunter",
        "total_characters_found": len(merged),
        "characters": []
    }

    for char in merged:
        output_data["characters"].append({
            "canonical_name": char.canonical_name,
            "aliases": char.aliases,
            "is_speaker": char.is_speaker,
            "dialogue_count": char.dialogue_count,
            "voice_clues": char.voice_clues,
            "sample_dialogues": char.sample_dialogues,
            "qwen_voice_prompt": ""  # To be filled by voice agent
        })

    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)

    print(f"\n   Results saved to: {output_path}")

    return True


def test_combined_workflow():
    """Test the combined parser + extraction workflow."""
    print("\n" + "=" * 80)
    print("TESTING COMBINED WORKFLOW")
    print("=" * 80)

    # Read the manuscript
    manuscript_path = Path(__file__).parent.parent / "DealHunterShip.txt"
    if not manuscript_path.exists():
        print(f"ERROR: Manuscript not found at {manuscript_path}")
        return False

    with open(manuscript_path, encoding='utf-8') as f:
        text = f.read()

    # Parse manuscript
    print("\n1. Parsing manuscript...")
    manuscript = parse_manuscript(text)
    print(f"   Sections: {len(manuscript.sections)}")

    # Extract characters
    print("\n2. Extracting characters...")
    characters = extract_characters_from_text(text)
    print(f"   Characters: {len(characters)}")

    # Match characters to sections
    print("\n3. Matching characters to sections...")
    character_sections = {}

    for char in characters[:10]:  # Top 10 characters
        char_sections = []
        for section in manuscript.sections:
            if char.canonical_name in section.content or \
               any(alias in section.content for alias in char.aliases):
                char_sections.append({
                    'chapter': section.chapter_number,
                    'scene': section.scene_index,
                    'title': section.chapter_title
                })
        character_sections[char.canonical_name] = char_sections

    print("\n   Character appearances by chapter:")
    for char_name, sections in character_sections.items():
        if sections:
            chapter_nums = sorted(set(s['chapter'] for s in sections))
            print(f"      {char_name}: Chapters {chapter_nums}")

    # Generate voice prompts
    print("\n4. Generating voice prompts...")
    voice_prompts = {}

    for char in characters[:10]:
        if char.is_speaker and char.voice_clues:
            # Simple voice prompt generation based on clues
            gender = "female" if any(word in char.canonical_name.lower() for word in ['princess', 'kainda', 'vicky', 'nanny', 'willamina']) else "male"
            age = "young" if any(word in char.voice_clues for word in ['young', 'youth']) else "mature"

            voice_prompts[char.canonical_name] = {
                "gender": gender,
                "age": age,
                "clues": char.voice_clues[:3],
                "prompt": f"A {age} {gender} voice with {', '.join(char.voice_clues[:2])} characteristics."
            }

    print(f"   Voice prompts generated: {len(voice_prompts)}")

    # Save combined results
    output_path = Path(__file__).parent.parent / "DealHunterShip_analysis.json"
    output_data = {
        "manuscript": {
            "title": manuscript.title,
            "author": manuscript.author,
            "sections": len(manuscript.sections),
            "chapters": len([s for s in manuscript.sections if s.section_type == 'chapter']),
        },
        "characters": [
            {
                "name": char.canonical_name,
                "aliases": char.aliases,
                "is_speaker": char.is_speaker,
                "dialogue_count": char.dialogue_count,
                "voice_clues": char.voice_clues,
                "voice_prompt": voice_prompts.get(char.canonical_name, {})
            }
            for char in characters[:20]
        ],
        "character_sections": character_sections
    }

    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)

    print(f"\n   Analysis saved to: {output_path}")

    return True


def main():
    """Run all tests."""
    print("UmbraAudio Document Parser & Entity Extraction Tests")
    print("=" * 80)

    results = {}

    try:
        results['document_parser'] = test_document_parser()
    except Exception as e:
        print(f"ERROR in document parser test: {e}")
        results['document_parser'] = False

    try:
        results['entity_extraction'] = test_entity_extraction()
    except Exception as e:
        print(f"ERROR in entity extraction test: {e}")
        results['entity_extraction'] = False

    try:
        results['combined_workflow'] = test_combined_workflow()
    except Exception as e:
        print(f"ERROR in combined workflow test: {e}")
        results['combined_workflow'] = False

    # Summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80)

    all_passed = True
    for test_name, passed in results.items():
        status = "PASS" if passed else "FAIL"
        print(f"   {test_name}: {status}")
        if not passed:
            all_passed = False

    print("\n" + "=" * 80)
    if all_passed:
        print("ALL TESTS PASSED!")
    else:
        print("SOME TESTS FAILED!")
    print("=" * 80)

    return 0 if all_passed else 1


if __name__ == "__main__":
    sys.exit(main())