hiffi-artist-intelligence/collectors/youtube_collector.py

"""
collectors/youtube_collector.py
Collects YouTube statistics, uploads, view velocity, and true momentum.
Includes URL/handle parsing, historical snapshotting, status_flag, and data_complete.
"""

from datetime import datetime, timedelta, timezone
import os
import re
import sqlite3
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlparse
from googleapiclient.discovery import build


def parse_youtube_identifier(raw_input: str) -> Tuple[str, str]:
    """Parses raw input (URL, handle, or ID) into a query type and clean value."""
    val = raw_input.strip()

    if "youtube.com" in val or "youtu.be" in val:
        parsed = urlparse(val if val.startswith("http") else f"https://{val}")
        path_parts = [p for p in parsed.path.split("/") if p]
        if not path_parts:
            raise ValueError(f"Invalid YouTube URL: {val}")

        if path_parts[0].startswith("@"):
            return ("forHandle", path_parts[0])
        elif path_parts[0] == "channel" and len(path_parts) > 1:
            return ("id", path_parts[1])
        elif path_parts[0] == "user" and len(path_parts) > 1:
            return ("forUsername", path_parts[1])
        elif path_parts[0] == "c" and len(path_parts) > 1:
            return ("forHandle", f"@{path_parts[1]}")
        else:
            return (
                "forHandle",
                path_parts[0]
                if path_parts[0].startswith("@")
                else f"@{path_parts[0]}",
            )

    if val.startswith("UC") and len(val) == 24:
        return ("id", val)

    if val.startswith("@"):
        return ("forHandle", val)

    return ("forHandle", f"@{val}")


class YouTubeCollector:

    def __init__(
        self,
        api_key: Optional[str] = None,
        cache_db_path: str = ".cache/youtube_metrics.sqlite",
    ):
        self.api_key = api_key or os.getenv("YOUTUBE_API_KEY")
        if not self.api_key:
            raise ValueError("YOUTUBE_API_KEY not found in environment.")
        self.youtube = build("youtube", "v3", developerKey=self.api_key)
        self.cache_db_path = cache_db_path
        self._init_db()

    def _init_db(self):
        os.makedirs(os.path.dirname(self.cache_db_path), exist_ok=True)
        with sqlite3.connect(self.cache_db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS channel_snapshots (
                    channel_id TEXT,
                    snapshot_date TEXT,
                    lifetime_views INTEGER,
                    subscriber_count INTEGER,
                    video_count INTEGER,
                    PRIMARY KEY (channel_id, snapshot_date)
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS channel_metrics (
                    identifier TEXT PRIMARY KEY,
                    channel_id TEXT,
                    channel_title TEXT,
                    subscriber_count INTEGER,
                    lifetime_views INTEGER,
                    video_count INTEGER,
                    recent_avg_views REAL,
                    daily_view_velocity REAL,
                    momentum_30d REAL,
                    status_flag TEXT,
                    data_complete INTEGER,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

    def _resolve_channel(self, identifier: str) -> Optional[Dict[str, Any]]:
        param_type, param_value = parse_youtube_identifier(identifier)
        req_params = {
            "part": "snippet,statistics,contentDetails",
            param_type: param_value,
        }

        try:
            resp = self.youtube.channels().list(**req_params).execute()
            items = resp.get("items", [])
            return items[0] if items else None
        except Exception as e:
            print(f"API Resolution error for '{identifier}': {e}")
            return None

    def _record_snapshot(
        self, channel_id: str, views: int, subs: int, videos: int
    ):
        today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        with sqlite3.connect(self.cache_db_path) as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO channel_snapshots 
                (channel_id, snapshot_date, lifetime_views, subscriber_count, video_count)
                VALUES (?, ?, ?, ?, ?)
            """,
                (channel_id, today, views, subs, videos),
            )

    def _calculate_historical_channel_momentum(
        self, channel_id: str
    ) -> Optional[float]:
        with sqlite3.connect(self.cache_db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT snapshot_date, lifetime_views 
                FROM channel_snapshots 
                WHERE channel_id = ? 
                ORDER BY snapshot_date DESC
            """,
                (channel_id,),
            )
            rows = cursor.fetchall()

        if len(rows) < 2:
            return None

        latest_date = datetime.strptime(rows[0]["snapshot_date"], "%Y-%m-%d")
        latest_views = rows[0]["lifetime_views"]

        for r in rows[1:]:
            past_date = datetime.strptime(r["snapshot_date"], "%Y-%m-%d")
            diff_days = (latest_date - past_date).days
            if 7 <= diff_days <= 45:
                past_views = r["lifetime_views"]
                delta_views = latest_views - past_views
                return float(delta_views / max(diff_days, 1))

        return None

    def _get_recent_video_metrics(
        self, uploads_playlist_id: str, max_results: int = 25
    ) -> Dict[str, Optional[float]]:
        now = datetime.now(timezone.utc)
        t_30d = now - timedelta(days=30)
        t_60d = now - timedelta(days=60)

        try:
            pl_resp = (
                self.youtube.playlistItems()
                .list(
                    part="contentDetails,snippet",
                    playlistId=uploads_playlist_id,
                    maxResults=max_results,
                )
                .execute()
            )
        except Exception:
            return {
                "recent_avg_views": None,
                "daily_view_velocity": None,
                "momentum_30d": None,
            }

        video_ids = [
            item["contentDetails"]["videoId"]
            for item in pl_resp.get("items", [])
        ]
        if not video_ids:
            return {
                "recent_avg_views": None,
                "daily_view_velocity": None,
                "momentum_30d": None,
            }

        v_resp = (
            self.youtube.videos()
            .list(part="snippet,statistics", id=",".join(video_ids))
            .execute()
        )

        views_0_to_30: List[int] = []
        views_30_to_60: List[int] = []
        daily_velocities: List[float] = []

        for item in v_resp.get("items", []):
            stats = item.get("statistics", {})
            v_count = int(stats.get("viewCount", 0))
            pub_str = item.get("snippet", {}).get("publishedAt")
            if not pub_str:
                continue

            pub_date = datetime.fromisoformat(pub_str.replace("Z", "+00:00"))
            days_active = max((now - pub_date).total_seconds() / 86400.0, 0.25)

            daily_velocities.append(v_count / days_active)

            if pub_date >= t_30d:
                views_0_to_30.append(v_count)
            elif t_60d <= pub_date < t_30d:
                views_30_to_60.append(v_count)

        avg_views = (
            float(sum(views_0_to_30) / len(views_0_to_30))
            if views_0_to_30
            else None
        )
        daily_vel = (
            float(sum(daily_velocities) / len(daily_velocities))
            if daily_velocities
            else None
        )

        if views_30_to_60 and views_0_to_30:
            prev_sum = sum(views_30_to_60)
            curr_sum = sum(views_0_to_30)
            momentum = (
                ((curr_sum - prev_sum) / float(max(prev_sum, 1000))) * 100.0
                if prev_sum > 0
                else 100.0
            )
        elif views_0_to_30 and not views_30_to_60:
            momentum = 100.0
        else:
            momentum = None

        return {
            "recent_avg_views": avg_views,
            "daily_view_velocity": daily_vel,
            "momentum_30d": momentum,
        }

    def get_channel_metrics(
        self,
        identifier: str,
        fetch_recent_videos: bool = True,
        force_refresh: bool = False,
    ) -> Dict[str, Any]:
        if not force_refresh:
            with sqlite3.connect(self.cache_db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT * FROM channel_metrics WHERE identifier = ?",
                    (identifier,),
                )
                row = cursor.fetchone()
                if row and row["status_flag"] != "FAILED_LOOKUP":
                    return dict(row)

        ch_data = self._resolve_channel(identifier)
        if not ch_data:
            return {
                "identifier": identifier,
                "channel_id": None,
                "channel_title": None,
                "subscriber_count": None,
                "lifetime_views": None,
                "video_count": None,
                "recent_avg_views": None,
                "daily_view_velocity": None,
                "momentum_30d": None,
                "status_flag": "FAILED_LOOKUP",
                "data_complete": 0,
            }

        ch_id = ch_data["id"]
        title = ch_data["snippet"]["title"]
        stats = ch_data["statistics"]
        uploads_pl = ch_data["contentDetails"]["relatedPlaylists"]["uploads"]

        subs = int(stats.get("subscriberCount", 0))
        views = int(stats.get("viewCount", 0))
        vids = int(stats.get("videoCount", 0))

        self._record_snapshot(ch_id, views, subs, vids)

        recent_metrics = {
            "recent_avg_views": None,
            "daily_view_velocity": None,
            "momentum_30d": None,
        }
        if fetch_recent_videos:
            recent_metrics = self._get_recent_video_metrics(uploads_pl)

        momentum = recent_metrics["momentum_30d"]
        if momentum is None:
            momentum = self._calculate_historical_channel_momentum(ch_id)

        # Reconcile status_flag and data_complete
        if momentum is not None and recent_metrics["daily_view_velocity"] is not None:
            status_flag = "COMPLETE"
            data_complete = 1
        elif recent_metrics["daily_view_velocity"] is not None:
            status_flag = "CATALOG_ACTIVE_NO_RECENT_UPLOAD"
            momentum = 0.0
            data_complete = 1
        else:
            status_flag = "INACTIVE_OR_DORMANT"
            momentum = None
            data_complete = 0

        result = {
            "identifier": identifier,
            "channel_id": ch_id,
            "channel_title": title,
            "subscriber_count": subs,
            "lifetime_views": views,
            "video_count": vids,
            "recent_avg_views": recent_metrics["recent_avg_views"],
            "daily_view_velocity": recent_metrics["daily_view_velocity"],
            "momentum_30d": momentum,
            "status_flag": status_flag,
            "data_complete": data_complete,
        }

        with sqlite3.connect(self.cache_db_path) as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO channel_metrics 
                (identifier, channel_id, channel_title, subscriber_count, lifetime_views, 
                 video_count, recent_avg_views, daily_view_velocity, momentum_30d, status_flag, data_complete)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
                (
                    result["identifier"],
                    result["channel_id"],
                    result["channel_title"],
                    result["subscriber_count"],
                    result["lifetime_views"],
                    result["video_count"],
                    result["recent_avg_views"],
                    result["daily_view_velocity"],
                    result["momentum_30d"],
                    result["status_flag"],
                    result["data_complete"],
                ),
            )

        return result