#!/usr/bin/env python3
"""
Supadata.ai Quota Tracker

Tracks API usage for Supadata.ai YouTube transcript service.
Uses Elestio PostgreSQL for persistent storage.

STORY-005: Enhance Supadata integration with quota tracking.

Author: Genesis System
Version: 1.0.0
"""

import logging
import sys
from datetime import datetime, timezone
from typing import Dict, Any, Optional, Tuple

# Add genesis-memory to path for elestio_config
sys.path.insert(0, '/mnt/e/genesis-system/data/genesis-memory')

try:
    import psycopg2
    from psycopg2 import sql
    from psycopg2.extras import RealDictCursor
except ImportError:
    raise ImportError("psycopg2 required. Install with: pip install psycopg2-binary")

from elestio_config import PostgresConfig

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SupadataQuotaTracker:
    """
    Tracks Supadata.ai API quota usage with PostgreSQL persistence.

    Free tier limit: 100 calls/month
    Resets on the 1st of each month.
    """

    FREE_TIER_LIMIT = 100
    TABLE_NAME = "supadata_transcript_usage"

    def __init__(self, connection_params: Optional[Dict] = None):
        """
        Initialize the quota tracker.

        Args:
            connection_params: PostgreSQL connection parameters.
                              Defaults to Elestio config if not provided.
        """
        self.connection_params = connection_params or PostgresConfig.get_connection_params()
        self._ensure_table_exists()

    def _get_connection(self):
        """Create a new database connection."""
        try:
            conn = psycopg2.connect(**self.connection_params)
            return conn
        except psycopg2.Error as e:
            logger.error(f"Failed to connect to PostgreSQL: {e}")
            raise ConnectionError(f"PostgreSQL connection failed: {e}")

    def _ensure_table_exists(self):
        """Create the usage tracking table if it doesn't exist."""
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS supadata_transcript_usage (
            id SERIAL PRIMARY KEY,
            video_id VARCHAR(20) NOT NULL,
            month_key VARCHAR(7) NOT NULL,  -- Format: YYYY-MM
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            status VARCHAR(20) DEFAULT 'success',
            metadata JSONB DEFAULT '{}'
        );

        CREATE INDEX IF NOT EXISTS idx_supadata_month_key
            ON supadata_transcript_usage(month_key);
        CREATE INDEX IF NOT EXISTS idx_supadata_video_id
            ON supadata_transcript_usage(video_id);
        """

        try:
            conn = self._get_connection()
            cursor = conn.cursor()
            cursor.execute(create_table_sql)
            conn.commit()
            cursor.close()
            conn.close()
            logger.info("Supadata usage table verified/created")
        except psycopg2.Error as e:
            logger.error(f"Failed to create table: {e}")
            raise

    def _get_current_month_key(self) -> str:
        """Get the current month key in YYYY-MM format."""
        return datetime.now(timezone.utc).strftime("%Y-%m")

    def check_quota(self) -> Dict[str, Any]:
        """
        Check the remaining quota for the current month.

        Returns:
            Dict with quota information:
            - limit: Total monthly limit
            - used: Number of calls used this month
            - remaining: Number of calls remaining
            - month: Current month key
            - reset_date: Next reset date (1st of next month)
        """
        month_key = self._get_current_month_key()

        try:
            conn = self._get_connection()
            cursor = conn.cursor()

            cursor.execute(
                "SELECT COUNT(*) FROM supadata_transcript_usage WHERE month_key = %s",
                (month_key,)
            )
            result = cursor.fetchone()
            used = result[0] if result else 0

            cursor.close()
            conn.close()

            # Calculate next reset date
            now = datetime.now(timezone.utc)
            if now.month == 12:
                reset_date = datetime(now.year + 1, 1, 1)
            else:
                reset_date = datetime(now.year, now.month + 1, 1)

            return {
                "limit": self.FREE_TIER_LIMIT,
                "used": used,
                "remaining": max(0, self.FREE_TIER_LIMIT - used),
                "month": month_key,
                "reset_date": reset_date.strftime("%Y-%m-%d"),
                "percentage_used": round((used / self.FREE_TIER_LIMIT) * 100, 1)
            }

        except psycopg2.Error as e:
            logger.error(f"Failed to check quota: {e}")
            raise

    def record_usage(
        self,
        video_id: str,
        status: str = "success",
        metadata: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """
        Record an API call usage.

        Args:
            video_id: The YouTube video ID that was fetched
            status: Status of the call ('success', 'failed', 'cached')
            metadata: Optional additional metadata about the call

        Returns:
            Dict with the recorded usage and updated quota info
        """
        month_key = self._get_current_month_key()
        metadata = metadata or {}

        try:
            conn = self._get_connection()
            cursor = conn.cursor()

            import json
            cursor.execute(
                """
                INSERT INTO supadata_transcript_usage (video_id, month_key, status, metadata)
                VALUES (%s, %s, %s, %s)
                RETURNING id, created_at
                """,
                (video_id, month_key, status, json.dumps(metadata))
            )

            result = cursor.fetchone()
            record_id = result[0]
            created_at = result[1]

            conn.commit()
            cursor.close()
            conn.close()

            # Get updated quota info
            quota_info = self.check_quota()

            logger.info(
                f"Recorded Supadata usage: video={video_id}, "
                f"remaining={quota_info['remaining']}/{quota_info['limit']}"
            )

            return {
                "record_id": record_id,
                "video_id": video_id,
                "created_at": created_at.isoformat() if created_at else None,
                "status": status,
                **quota_info
            }

        except psycopg2.Error as e:
            logger.error(f"Failed to record usage: {e}")
            raise

    def get_usage_stats(self, months: int = 3) -> Dict[str, Any]:
        """
        Get usage statistics for recent months.

        Args:
            months: Number of months to include in statistics

        Returns:
            Dict with detailed usage statistics
        """
        try:
            conn = self._get_connection()
            cursor = conn.cursor(cursor_factory=RealDictCursor)

            # Get monthly breakdown
            cursor.execute(
                """
                SELECT
                    month_key,
                    COUNT(*) as total_calls,
                    COUNT(CASE WHEN status = 'success' THEN 1 END) as successful,
                    COUNT(CASE WHEN status = 'failed' THEN 1 END) as failed,
                    COUNT(DISTINCT video_id) as unique_videos,
                    MIN(created_at) as first_call,
                    MAX(created_at) as last_call
                FROM supadata_transcript_usage
                WHERE month_key >= TO_CHAR(
                    CURRENT_DATE - INTERVAL '%s months', 'YYYY-MM'
                )
                GROUP BY month_key
                ORDER BY month_key DESC
                """,
                (months,)
            )

            monthly_stats = cursor.fetchall()

            # Get all-time stats
            cursor.execute(
                """
                SELECT
                    COUNT(*) as total_calls,
                    COUNT(DISTINCT video_id) as unique_videos,
                    COUNT(DISTINCT month_key) as active_months,
                    MIN(created_at) as first_ever_call
                FROM supadata_transcript_usage
                """
            )
            all_time = cursor.fetchone()

            # Get recent videos
            cursor.execute(
                """
                SELECT video_id, created_at, status
                FROM supadata_transcript_usage
                ORDER BY created_at DESC
                LIMIT 10
                """
            )
            recent_videos = cursor.fetchall()

            cursor.close()
            conn.close()

            # Current quota
            current_quota = self.check_quota()

            return {
                "current_month": current_quota,
                "monthly_breakdown": [dict(row) for row in monthly_stats],
                "all_time": dict(all_time) if all_time else {},
                "recent_videos": [dict(row) for row in recent_videos],
                "generated_at": datetime.now(timezone.utc).isoformat() + "Z"
            }

        except psycopg2.Error as e:
            logger.error(f"Failed to get usage stats: {e}")
            raise

    def is_quota_available(self) -> bool:
        """
        Check if quota is available for another API call.

        Returns:
            True if quota remaining > 0, False otherwise
        """
        quota = self.check_quota()
        return quota["remaining"] > 0

    def alert_if_low(self, threshold: int = 10) -> Optional[Dict[str, Any]]:
        """
        Check if quota is running low and return a warning if so.

        Args:
            threshold: Number of remaining calls that triggers the warning

        Returns:
            Warning dict if quota < threshold, None otherwise
        """
        quota = self.check_quota()

        if quota["remaining"] <= threshold:
            warning = {
                "level": "CRITICAL" if quota["remaining"] <= 5 else "WARNING",
                "message": f"Supadata quota running low: {quota['remaining']} calls remaining",
                "remaining": quota["remaining"],
                "limit": quota["limit"],
                "reset_date": quota["reset_date"],
                "percentage_used": quota["percentage_used"]
            }

            logger.warning(warning["message"])
            return warning

        return None

    def check_video_already_fetched(self, video_id: str) -> Optional[Dict[str, Any]]:
        """
        Check if a video has already been fetched (for caching/deduplication).

        Args:
            video_id: The YouTube video ID to check

        Returns:
            Dict with fetch info if already fetched, None otherwise
        """
        try:
            conn = self._get_connection()
            cursor = conn.cursor(cursor_factory=RealDictCursor)

            cursor.execute(
                """
                SELECT video_id, created_at, status, metadata
                FROM supadata_transcript_usage
                WHERE video_id = %s AND status = 'success'
                ORDER BY created_at DESC
                LIMIT 1
                """,
                (video_id,)
            )

            result = cursor.fetchone()
            cursor.close()
            conn.close()

            if result:
                return dict(result)
            return None

        except psycopg2.Error as e:
            logger.error(f"Failed to check video status: {e}")
            raise


class QuotaExhaustedError(Exception):
    """Raised when Supadata API quota is exhausted."""
    pass


def fetch_transcript_with_quota(
    client,
    url: str,
    tracker: Optional[SupadataQuotaTracker] = None,
    skip_if_fetched: bool = True,
    **kwargs
) -> Dict[str, Any]:
    """
    Wrapper function that checks quota before fetching a transcript.

    Args:
        client: SupadataClient instance
        url: YouTube video URL or ID
        tracker: SupadataQuotaTracker instance (creates one if not provided)
        skip_if_fetched: Skip if video already fetched successfully
        **kwargs: Additional arguments to pass to get_transcript

    Returns:
        Transcript data or error dict

    Raises:
        QuotaExhaustedError: If no quota remaining
    """
    if tracker is None:
        tracker = SupadataQuotaTracker()

    # Extract video ID
    video_id = client.extract_video_id(url)

    # Check if already fetched
    if skip_if_fetched:
        existing = tracker.check_video_already_fetched(video_id)
        if existing:
            logger.info(f"Video {video_id} already fetched at {existing['created_at']}")
            return {
                "video_id": video_id,
                "status": "cached",
                "message": "Video transcript already fetched",
                "original_fetch": existing
            }

    # Check quota
    if not tracker.is_quota_available():
        quota = tracker.check_quota()
        raise QuotaExhaustedError(
            f"Supadata quota exhausted. Used {quota['used']}/{quota['limit']} calls. "
            f"Resets on {quota['reset_date']}."
        )

    # Check for low quota warning
    warning = tracker.alert_if_low(threshold=10)

    # Attempt the fetch
    try:
        result = client.get_transcript(url, **kwargs)

        # Record successful usage
        tracker.record_usage(
            video_id=video_id,
            status="success",
            metadata={"lang": kwargs.get("lang"), "mode": kwargs.get("mode", "auto")}
        )

        if warning:
            result["quota_warning"] = warning

        return result

    except Exception as e:
        # Record failed attempt (still counts against quota based on API behavior)
        tracker.record_usage(
            video_id=video_id,
            status="failed",
            metadata={"error": str(e)}
        )
        raise


# CLI for testing
if __name__ == "__main__":
    import argparse
    import json

    parser = argparse.ArgumentParser(description="Supadata Quota Tracker CLI")
    parser.add_argument("--check", action="store_true", help="Check current quota")
    parser.add_argument("--stats", action="store_true", help="Show usage statistics")
    parser.add_argument("--record", help="Record usage for a video ID")
    parser.add_argument("--alert", type=int, default=10, help="Alert threshold")

    args = parser.parse_args()

    try:
        tracker = SupadataQuotaTracker()

        if args.check:
            quota = tracker.check_quota()
            print("\n=== Supadata Quota Status ===")
            print(f"Month:     {quota['month']}")
            print(f"Used:      {quota['used']}/{quota['limit']}")
            print(f"Remaining: {quota['remaining']}")
            print(f"Usage:     {quota['percentage_used']}%")
            print(f"Resets:    {quota['reset_date']}")

            warning = tracker.alert_if_low(args.alert)
            if warning:
                print(f"\n[{warning['level']}] {warning['message']}")

        elif args.stats:
            stats = tracker.get_usage_stats()
            print("\n=== Supadata Usage Statistics ===")
            print(json.dumps(stats, indent=2, default=str))

        elif args.record:
            result = tracker.record_usage(args.record)
            print(f"\nRecorded usage for video: {args.record}")
            print(f"Remaining quota: {result['remaining']}/{result['limit']}")

        else:
            # Default: show quota
            quota = tracker.check_quota()
            print(f"Quota: {quota['remaining']}/{quota['limit']} remaining ({quota['month']})")

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)