main.py

#!/usr/bin/env python3
"""
instagram-places-mapper
───────────────────────
Extract places worth visiting from an Instagram saved_posts.json export,
geocode them, and produce a CSV + KML file ready to import into Google My Maps.

Usage
─────
  python main.py <path/to/saved_posts.json>                        # full pipeline
  python main.py <path/to/saved_posts.json> --no-extract           # re-geocode existing CSV
  python main.py <path/to/saved_posts.json> --no-geocode           # extract only (writes CSV)
  python main.py <path/to/saved_posts.json> --no-resume            # ignore checkpoint, start fresh
  python main.py <path/to/saved_posts.json> --transcribe           # add transcription fallback step
  python main.py <path/to/saved_posts.json> --no-extract --transcribe  # transcribe-only run

Output (written to --output-dir, default: same folder as the JSON file)
  places_full.csv   — one row per place with coordinates + reel URL
  places_map.kml    — Google My Maps–ready KML organised by country › city

Environment
───────────
  ANTHROPIC_API_KEY  required for extraction (set in .env or shell)

Resuming interrupted extraction
───────────────────────────────
  A checkpoint file (.checkpoint.json) is written next to the CSV after every
  10 Claude API calls. Re-running without --no-resume picks up where it left off.
"""

import argparse
import os
import sys
from pathlib import Path

from pipeline.extract import DEFAULT_MODEL, DEFAULT_OLLAMA_MODEL, MODELS, MODELS_OLLAMA
from pipeline.ollama import DEFAULT_BASE_URL as OLLAMA_DEFAULT_URL


# ── .env loader ───────────────────────────────────────────────────────────────

def _load_env(start: Path) -> None:
    """Walk up from start looking for a .env file and load it."""
    for directory in [start, *start.parents]:
        env_file = directory / ".env"
        if env_file.exists():
            with open(env_file) as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith("#") and "=" in line:
                        k, v = line.split("=", 1)
                        os.environ.setdefault(k.strip(), v.strip())
            return


# ── CLI ───────────────────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Extract places worth visiting from an Instagram export and pin them on Google My Maps.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("input", help="Path to saved_posts.json from Instagram export")
    p.add_argument("--output-dir", default=None,
                   help="Directory for output files (default: same folder as input JSON)")
    p.add_argument("--no-extract", action="store_true",
                   help="Skip Claude extraction; geocode and export the existing CSV")
    p.add_argument("--no-geocode", action="store_true",
                   help="Skip geocoding; only run extraction and export")
    p.add_argument("--no-resume", action="store_true",
                   help="Ignore checkpoint; re-extract all posts from scratch")
    p.add_argument("--transcribe", action="store_true",
                   help="After extraction, download and transcribe reels that had no caption info, "
                        "then re-run the active provider (Claude or Ollama) on the transcript. "
                        "With --provider ollama the whole step is $0 and stays on-device. "
                        "Requires openai-whisper + yt-dlp.")
    p.add_argument("--whisper-model", default="small",
                   choices=["tiny", "base", "small", "medium", "large"],
                   help="Whisper model size (default: small). Larger = more accurate but slower.")
    p.add_argument("--browser", default="chrome",
                   help="Browser to pull Instagram cookies from for yt-dlp auth "
                        "(default: chrome). Options: chrome, firefox, brave, edge. "
                        "Note: safari cookies are sandboxed on macOS and cannot be read. "
                        "Use 'none' to skip cookie auth.")
    p.add_argument("--provider", default="anthropic", choices=["anthropic", "ollama"],
                   help="Inference provider: 'anthropic' (default, requires API key) or "
                        "'ollama' (free, local — requires 'ollama serve' running).")
    p.add_argument("--model", default=DEFAULT_MODEL, choices=list(MODELS),
                   help=f"Anthropic model (default: {DEFAULT_MODEL}). Ignored when --provider ollama.")
    p.add_argument("--ollama-model", default=DEFAULT_OLLAMA_MODEL,
                   choices=list(MODELS_OLLAMA),
                   help=f"Ollama model to use (default: {DEFAULT_OLLAMA_MODEL}). "
                        "Only used when --provider ollama.")
    p.add_argument("--ollama-url", default=OLLAMA_DEFAULT_URL,
                   help=f"Ollama server URL (default: {OLLAMA_DEFAULT_URL}).")
    p.add_argument("--batch-size", type=int, default=10,
                   help="Number of captions sent per LLM call (default: 10). "
                        "Higher values reduce API round-trips and cost; "
                        "lower values give faster partial results.")
    p.add_argument("--batch-api", action="store_true",
                   help="Use the Anthropic Batch API (50%% cost discount, async). "
                        "Submits all requests at once, polls until complete "
                        "(typically a few minutes). Anthropic-only; ignored for Ollama.")
    p.add_argument("-y", "--yes", action="store_true",
                   help="Skip the cost-estimate confirmation prompt and run immediately.")
    return p


# ── Entry point ───────────────────────────────────────────────────────────────

def main() -> None:
    args = build_parser().parse_args()

    input_path = Path(args.input).resolve()
    if not input_path.exists():
        sys.exit(f"Error: {input_path} does not exist.")

    output_dir = Path(args.output_dir).resolve() if args.output_dir else input_path.parent
    output_dir.mkdir(parents=True, exist_ok=True)

    csv_path = output_dir / "places_full.csv"
    kml_path = output_dir / "places_map.kml"

    # Load .env starting from the input file's directory
    _load_env(input_path.parent)

    needs_llm = not args.no_extract or args.transcribe

    # Resolve the effective model ID based on provider
    active_model = args.ollama_model if args.provider == "ollama" else args.model

    # ── Cost estimate ─────────────────────────────────────────────────────────
    if needs_llm:
        if args.provider == "anthropic" and "ANTHROPIC_API_KEY" not in os.environ:
            sys.exit(
                "Error: ANTHROPIC_API_KEY is not set.\n"
                "Add it to a .env file next to your saved_posts.json, or export it in your shell.\n"
                "Or run with --provider ollama to use a free local model."
            )

        from pipeline import extract as extract_mod
        from pipeline.extract import estimate_cost, load_posts, _load_checkpoint, _warm_start_from_csv
        from pathlib import Path as _Path
        from pipeline.transcribe import _skipped_posts

        posts = load_posts(str(input_path))
        checkpoint_file = csv_path.with_suffix(".checkpoint.json")
        checkpoint = _load_checkpoint(checkpoint_file) if not args.no_resume else {}
        if not args.no_resume:
            _warm_start_from_csv(checkpoint, checkpoint_file, str(csv_path))

        new_extract_urls = {p["url"] for p in posts if p["url"] not in checkpoint} if not args.no_extract else set()

        transcribe_count = 0
        if args.transcribe:
            trans_checkpoint_path = _Path(str(csv_path.with_suffix("")) + ".transcribe_checkpoint.json")
            trans_checkpoint = {}
            if not args.no_resume and trans_checkpoint_path.exists():
                import json as _json
                with open(trans_checkpoint_path) as f:
                    trans_checkpoint = _json.load(f)
            skipped = _skipped_posts(str(input_path), str(csv_path))
            transcribe_count = sum(1 for p in skipped if p["url"] not in trans_checkpoint)

        provider_label = f"Ollama ({active_model})" if args.provider == "ollama" else "Claude"

        _sep = "─" * 62
        print(_sep)
        print(" Cost estimate")
        print(_sep)
        if not args.no_extract:
            print(f"  Extraction : {len(new_extract_urls)} new posts × {provider_label}")
        if args.transcribe:
            print(f"  Transcribe : {transcribe_count} reels × Whisper (free) + {provider_label}")
        print(f"  Geocoding  : free  (Nominatim / OpenStreetMap)")
        print(f"  KML export : free  (local)")
        print()

        if args.provider == "ollama":
            print(f"  Provider   : Ollama (local) — $0.00 total")
            print(f"  Model      : {active_model}  ({MODELS_OLLAMA.get(active_model, {}).get('label', '')})")
            print(f"  ✓ All processing is free and stays on your machine.")
        else:
            batch_api_flag = args.batch_api and args.provider == "anthropic"
            print(f"  {'Model':<22}  {'Extract':>9}  {'Transcribe':>11}  {'Total':>9}  {'Notes'}")
            print(f"  {'─'*22}  {'─'*9}  {'─'*11}  {'─'*9}  {'─'*20}")
            for mid, info in MODELS.items():
                est = estimate_cost(posts, new_extract_urls, mid, transcribe_count,
                                    batch_size=args.batch_size,
                                    use_batch_api=batch_api_flag)
                marker = "← selected" if mid == args.model else ""
                rec    = " (recommended)" if mid == "claude-haiku-4-5" else ""
                print(
                    f"  {mid:<22}  ${est['extract_cost']:>8.3f}  ${est['transcribe_cost']:>10.3f}"
                    f"  ${est['total_cost']:>8.3f}  {marker}{rec}"
                )
            if batch_api_flag:
                print(f"  50% Batch API discount applied to extraction costs.")
            print(f"  Prices are estimates. Actual cost may vary by ±20%.")
            print(f"  Use --provider ollama for free local inference.")
        print(_sep)

        if not args.yes:
            try:
                answer = input(f"\nProceed with {args.provider}/{active_model}? [Y/n]: ").strip().lower()
            except (EOFError, KeyboardInterrupt):
                print("\nAborted.")
                sys.exit(0)
            if answer and answer not in ("y", "yes"):
                print("Aborted. Use --model to pick a different model.")
                sys.exit(0)
        print()

    # ── Step 1: Extract ───────────────────────────────────────────────────────
    if args.no_extract:
        if not csv_path.exists():
            sys.exit(f"Error: --no-extract requires an existing {csv_path}")
        print(f"Skipping extraction — using {csv_path}\n")
    else:
        from pipeline import extract
        extract.run(
            json_path=str(input_path),
            output_csv=str(csv_path),
            resume=not args.no_resume,
            model=active_model,
            provider=args.provider,
            ollama_url=args.ollama_url,
            batch_size=args.batch_size,
            use_batch_api=args.batch_api and args.provider == "anthropic",
        )
        print()

    # ── Step 1b: Transcription fallback ──────────────────────────────────────
    if args.transcribe:
        from pipeline import transcribe
        transcribe.run(
            json_path=str(input_path),
            extracted_csv=str(csv_path),
            whisper_model_name=args.whisper_model,
            browser=None if args.browser == "none" else args.browser,
            resume=not args.no_resume,
            model=active_model,
            provider=args.provider,
            ollama_url=args.ollama_url,
        )
        print()

    # ── Step 2: Geocode ───────────────────────────────────────────────────────
    if not args.no_geocode:
        from pipeline import geocode
        geocode.run(input_csv=str(csv_path))
        print()
    else:
        print("Skipping geocoding.\n")

    # ── Step 2b: Manual overrides ─────────────────────────────────────────────
    import csv as _csv
    from pipeline import override as override_mod
    override_path = output_dir / "places_override.csv"
    with open(csv_path, encoding="utf-8") as f:
        csv_rows = list(_csv.DictReader(f))
    csv_rows, applied = override_mod.apply(csv_rows, override_path)
    if applied:
        # Write back the corrected coordinates before export
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            from pipeline import extract as _ext
            writer = _csv.DictWriter(f, fieldnames=_ext.FIELDNAMES)
            writer.writeheader()
            for row in csv_rows:
                writer.writerow({k: row.get(k, "") for k in _ext.FIELDNAMES})
    print()

    # ── Step 3: Export KML ────────────────────────────────────────────────────
    from pipeline import export
    export.run(input_csv=str(csv_path), output_kml=str(kml_path))


if __name__ == "__main__":
    main()