main.py

"""PickLine Analyzer CLI.

Usage:
  python main.py harvest --window today
  python main.py score
  python main.py topk --window today --k 10
  python main.py screenshot path/to/img.png --platform Tinder
  python main.py stats
"""
from __future__ import annotations

import argparse
import base64
import json
import logging
import os
import re
import sys
import time
from pathlib import Path

import anthropic
from dotenv import load_dotenv

from arctic_shift import ArcticShiftClient
from harvest import harvest, time_window
from score import score_all
from store import Store
from topk import top_k

DB_PATH = Path("picklines.db")
CONFIG_DIR = Path("config")
LOG_DIR = Path("output")

VISION_MODEL = "claude-haiku-4-5-20251001"

SCREENSHOT_PROMPT = """\
This is a screenshot from {platform} showing a dating/messaging conversation.

Extract the pickup line / opening message and reply details. Return JSON:
{{
  "line_text": "the exact opening message or pickup line",
  "response": "what the other person replied, or null",
  "outcome_hint": "any outcome clue (date plan, number exchange, hookup, etc.), or null",
  "platform": "{platform}"
}}

User-provided context: {context}

Return ONLY the JSON object — no markdown fences, no explanation.
"""


def setup_logging(debug: bool = False) -> None:
    LOG_DIR.mkdir(exist_ok=True)
    logging.basicConfig(
        level=logging.DEBUG if debug else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        handlers=[
            logging.FileHandler(LOG_DIR / "run.log", encoding="utf-8"),
            logging.StreamHandler(sys.stdout),
        ],
    )


def load_subreddits() -> list[str]:
    path = CONFIG_DIR / "subreddits.txt"
    if not path.exists():
        raise SystemExit(f"Missing {path}")
    return [l.strip() for l in path.read_text().splitlines() if l.strip() and not l.startswith("#")]


# ── subcommands ────────────────────────────────────────────────────────────────

def cmd_harvest(args: argparse.Namespace, store: Store, ai: anthropic.Anthropic) -> None:
    arctic = ArcticShiftClient(delay_seconds=1.0)
    subs = [args.subreddit] if args.subreddit else load_subreddits()
    since_ts, until_ts = time_window(args.window)
    run_id = store.start_run("harvest", ",".join(subs), args.window, since_ts, until_ts)
    try:
        posts, lines = harvest(store, arctic, ai, subs, args.window, args.max_posts)
        store.finish_run(run_id, posts, lines)
        print(f"\nHarvest done: {posts} posts scanned, {lines} new lines found.")
    except Exception as e:
        store.finish_run(run_id, 0, 0, error=str(e))
        raise


def cmd_score(args: argparse.Namespace, store: Store, ai: anthropic.Anthropic) -> None:
    run_id = store.start_run("score", "all", None, None, None)
    try:
        count = score_all(store, ai, limit=args.limit)
        store.finish_run(run_id, 0, count)
        print(f"\nScored {count} lines.")
    except Exception as e:
        store.finish_run(run_id, 0, 0, error=str(e))
        raise


def cmd_topk(args: argparse.Namespace, store: Store) -> None:
    since_ts, until_ts = time_window(args.window)
    lines = top_k(store, since_ts, until_ts, k=args.k)
    sep = "─" * 62
    print(f"\n  Top {args.k} Pickup Lines — {args.window.upper()}")
    print(f"  {sep}")
    if not lines:
        print("  No scored lines found for this window.")
        return
    for i, ln in enumerate(lines, 1):
        comp = ln.get("composite_score") or 0
        cre = ln.get("creativity") or 0
        rq = ln.get("response_quality") or 0
        hp = ln.get("hookup_probability") or 0
        print(f"\n  #{i}  Score {comp:.1f}  [C:{cre:.1f} R:{rq:.1f} H:{hp:.1f}]  r/{ln.get('subreddit','?')}")
        print(f"  Line:    {ln['line_text']}")
        if ln.get("response"):
            print(f"  Reply:   {ln['response']}")
        if ln.get("outcome_hint"):
            print(f"  Outcome: {ln['outcome_hint']}")
        if ln.get("url"):
            print(f"  Source:  {ln['url']}")
    print(f"\n  {sep}")


def cmd_screenshot(args: argparse.Namespace, store: Store, ai: anthropic.Anthropic) -> None:
    img_path = Path(args.image)
    if not img_path.exists():
        raise SystemExit(f"File not found: {img_path}")

    media_map = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp"}
    media_type = media_map.get(img_path.suffix.lower(), "image/jpeg")
    img_b64 = base64.b64encode(img_path.read_bytes()).decode()

    platform = args.platform or "dating app"
    prompt = SCREENSHOT_PROMPT.format(platform=platform, context=args.context or "")

    try:
        msg = ai.messages.create(
            model=VISION_MODEL,
            max_tokens=512,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": img_b64}},
                    {"type": "text", "text": prompt},
                ],
            }],
        )
        raw = msg.content[0].text.strip()
        raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw, flags=re.MULTILINE).strip()
        item = json.loads(raw)
    except Exception as e:
        raise SystemExit(f"Screenshot analysis failed: {e}") from e

    line_text = (item.get("line_text") or "").strip()
    if not line_text:
        print("No pickup line detected in this screenshot.")
        return

    new = store.upsert_line({
        "post_id": f"screenshot_{img_path.stem}_{int(time.time())}",
        "subreddit": item.get("platform") or platform,
        "line_text": line_text,
        "response": item.get("response"),
        "context": f"{item.get('platform', platform)} screenshot",
        "outcome_hint": item.get("outcome_hint"),
        "upvotes": 0,
        "created_utc": int(time.time()),
        "url": str(img_path.resolve()),
        "source": "screenshot",
    })

    print(f"\nLine:    {line_text!r}")
    if item.get("response"):
        print(f"Reply:   {item['response']!r}")
    if item.get("outcome_hint"):
        print(f"Outcome: {item['outcome_hint']!r}")
    print(f"\n{'Stored new line.' if new else 'Already in database.'} Run `python main.py score` to AI-score it.")


def cmd_stats(args: argparse.Namespace, store: Store) -> None:
    s = store.stats()
    print(f"\nPickLine Analyzer")
    print(f"  Total lines : {s['total_lines']}")
    print(f"  Scored      : {s['scored']}")
    print(f"  Unscored    : {s['unscored']}")
    print(f"  Avg scores  : C={s['avg_creativity']} R={s['avg_response_quality']} H={s['avg_hookup_probability']} → {s['avg_composite']}")
    print(f"\n  By subreddit:")
    for row in s["by_subreddit"]:
        print(f"    r/{row['subreddit']:<30} {row['n']} lines")


# ── entry point ────────────────────────────────────────────────────────────────

def main() -> None:
    load_dotenv()
    parser = argparse.ArgumentParser(description="PickLine Analyzer — top-K pickup lines from Reddit")
    parser.add_argument("--debug", action="store_true")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_h = sub.add_parser("harvest", help="Pull lines from Reddit")
    p_h.add_argument("--window", default="today", choices=["today", "yesterday", "weekend"])
    p_h.add_argument("--subreddit", default=None, help="Override config with a single subreddit")
    p_h.add_argument("--max-posts", type=int, default=500, dest="max_posts")

    p_s = sub.add_parser("score", help="AI-score unscored lines")
    p_s.add_argument("--limit", type=int, default=200)

    p_t = sub.add_parser("topk", help="Print top-K lines")
    p_t.add_argument("--window", default="today", choices=["today", "yesterday", "weekend"])
    p_t.add_argument("--k", type=int, default=10)

    p_ss = sub.add_parser("screenshot", help="Analyze a screenshot")
    p_ss.add_argument("image", help="Path to image (PNG/JPG/WEBP)")
    p_ss.add_argument("--platform", default=None)
    p_ss.add_argument("--context", default="")

    sub.add_parser("stats", help="Database stats")

    args = parser.parse_args()
    setup_logging(args.debug)

    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key and args.cmd in ("harvest", "score", "screenshot"):
        raise SystemExit("Set ANTHROPIC_API_KEY in .env")

    ai = anthropic.Anthropic(api_key=api_key or "")
    Path("output").mkdir(exist_ok=True)
    store = Store(DB_PATH)

    try:
        match args.cmd:
            case "harvest":    cmd_harvest(args, store, ai)
            case "score":      cmd_score(args, store, ai)
            case "topk":       cmd_topk(args, store)
            case "screenshot": cmd_screenshot(args, store, ai)
            case "stats":      cmd_stats(args, store)
    finally:
        store.close()


if __name__ == "__main__":
    main()