main.py

"""Main orchestration script for HabitHub calendar extraction.

This script coordinates the entire pipeline:
Phase 1: Split screenshot into month sections (OpenCV)
Phase 2: Extract data from each section (LLM)
Phase 3: Consolidate and validate data
"""

import argparse
import os
import sys

from dotenv import load_dotenv

import config
from split_screenshot import CalendarSplitter
from extract_with_llm import CalendarExtractor
from consolidate_data import DataConsolidator
from utils import find_first_image

# Load environment variables from .env file
load_dotenv()


def main():
    """Main entry point for the HabitHub export pipeline."""
    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Extract habit tracking data from HabitHub calendar screenshots'
    )
    parser.add_argument(
        'image_path',
        nargs='?',
        default='input/sample-screenshot.jpg',
        help='Path to the calendar screenshot image (default: input/sample-screenshot.jpg)'
    )
    args = parser.parse_args()

    print("HabitHub Calendar Export Tool")
    print("=" * 60)

    # Use the provided image path
    screenshot_path = args.image_path

    # Check if the file exists
    if not os.path.isfile(screenshot_path):
        print(f"\n✗ Error: Image file not found: {screenshot_path}")
        supported = ', '.join(config.SUPPORTED_EXTENSIONS)
        print(f"\n  Supported formats: {supported}")
        print("\n  Usage: python main.py [path/to/screenshot.jpg]")
        print(f"  Default: python main.py  (uses input/sample-screenshot.jpg)")
        sys.exit(1)

    # Validate file extension
    if not screenshot_path.lower().endswith(config.SUPPORTED_EXTENSIONS):
        print(f"\n✗ Error: Unsupported file format: {screenshot_path}")
        supported = ', '.join(config.SUPPORTED_EXTENSIONS)
        print(f"\n  Supported formats: {supported}")
        sys.exit(1)

    print(f"\n✓ Using screenshot: {screenshot_path}")

    # Phase 1: Split screenshot
    print("\nRunning Phase 1: Image Splitting...")
    splitter = CalendarSplitter(screenshot_path)
    if not splitter.run():
        print("\n✗ Phase 1 failed!")
        sys.exit(1)

    # Phase 2: Extract data with LLM
    print("\nRunning Phase 2: LLM-Based Data Extraction...")

    # Check for API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("\n✗ Error: OPENAI_API_KEY environment variable not set")
        print("\nPlease set your OpenAI API key:")
        print("  export OPENAI_API_KEY='your-api-key-here'")
        print("\nSkipping Phase 2 and Phase 3...")
        phase2_complete = False
        phase3_complete = False
    else:
        extractor = CalendarExtractor(api_key)
        metadata_path = os.path.join(config.OUTPUT_DIR, "split_metadata.json")

        all_month_data = extractor.process_all_months(metadata_path)

        if not all_month_data:
            print("\n✗ Phase 2 failed - no data extracted!")
            phase2_complete = False
            phase3_complete = False
        else:
            # Save extracted data
            extracted_data_path = os.path.join(config.OUTPUT_DIR,
                                               "extracted_data.json")
            extractor.save_extracted_data(all_month_data, extracted_data_path)
            phase2_complete = True

            # Phase 3: Consolidate and validate data
            print("\nRunning Phase 3: Data Consolidation...")

            consolidator = DataConsolidator(extracted_data_path)
            phase3_complete = consolidator.run(
                config.OUTPUT_CSV,
                config.OUTPUT_SUMMARY
            )

    # Summary
    print("\n" + "=" * 60)
    print("Pipeline Status:")
    print("=" * 60)
    print("✓ Phase 1: Image Splitting - COMPLETE")

    if phase2_complete:
        print("✓ Phase 2: LLM Extraction - COMPLETE")
    else:
        print("✗ Phase 2: LLM Extraction - FAILED or SKIPPED")

    if phase2_complete:
        if phase3_complete:
            print("✓ Phase 3: Data Consolidation - COMPLETE")
        else:
            print("✗ Phase 3: Data Consolidation - FAILED")
    else:
        print("⊘ Phase 3: Data Consolidation - SKIPPED")

    # Next steps
    if phase3_complete:
        print("\n🎉 All phases complete!")
        print("\nOutput files:")
        print(f"  - CSV data: {config.OUTPUT_CSV}")
        print(f"  - Summary: {config.OUTPUT_SUMMARY}")
        print(f"  - Raw data: {extracted_data_path}")
    elif phase2_complete:
        print("\nNext steps:")
        print("1. Review extracted data in output/extracted_data.json")
        print("2. Re-run Phase 3 to fix any issues")
    else:
        print("\nNext steps:")
        print("1. Review the split month images in output/split_months/")
        print("2. Set OPENAI_API_KEY environment variable")
        print("3. Re-run to complete Phase 2 and Phase 3")


if __name__ == "__main__":
    main()