import os
from googleapiclient.discovery import build
from datetime import datetime
import json
import sys
from dotenv import load_dotenv

load_dotenv()

YOUTUBE_API_KEY = os.getenv("YOUTUBE_API_KEY")
CHANNEL_ID = os.getenv("YOUTUBE_CHANNEL_ID")
OUTPUT_FILENAME = os.getenv("OUTPUT_FILENAME")


def validate_config():
    if not YOUTUBE_API_KEY:
        print("Error: YOUTUBE_API_KEY environment variable not set.", file=sys.stderr)
        sys.exit(1)
    if not CHANNEL_ID:
        print("Error: YOUTUBE_CHANNEL_ID environment variable not set.", file=sys.stderr)
        sys.exit(1)


def build_youtube_client():
    return build("youtube", "v3", developerKey=YOUTUBE_API_KEY)


def get_playlist_id(youtube, channel_id):
    try:
        request = youtube.channels().list(part="contentDetails", id=channel_id)
        response = request.execute()
        # Return the playlist ID containing all uploaded videos
        return response['items'][0]['contentDetails']['relatedPlaylists']['uploads'] if response['items'] else None
    except Exception as e:
        print(f"Error fetching channel details: {e}", file=sys.stderr)
        return None


def get_video_ids(youtube, playlist_id):
    video_ids, next_page_token = [], None
    while True:
        try:
            # Fetch playlist items in batches of 50 (max allowed)
            request = youtube.playlistItems().list(
                part="contentDetails", 
                playlistId=playlist_id, 
                maxResults=50, 
                pageToken=next_page_token
            )
            response = request.execute()
            # Extract video IDs from the current page
            video_ids.extend([item['contentDetails']['videoId'] for item in response['items']])
            # Check if there is a next page, else exit loop
            next_page_token = response.get('nextPageToken')
            if not next_page_token:
                break
        except Exception as e:
            print(f"Error fetching playlist items: {e}", file=sys.stderr)
            break
    return video_ids


def get_video_details(youtube, video_ids):
    all_video_data = []
    # Process video IDs in chunks of 50 due to API limits
    for i in range(0, len(video_ids), 50):
        video_chunk = ",".join(video_ids[i:i+50])
        try:
            request = youtube.videos().list(part="snippet,statistics", id=video_chunk)
            response = request.execute()
            # Timestamp for when data was extracted
            timestamp = datetime.utcnow().isoformat() + 'Z'
            # Extract relevant video info and stats
            for item in response['items']:
                video_data = {
                    'videoId': item['id'],
                    'title': item['snippet']['title'],
                    'publishedAt': item['snippet']['publishedAt'],
                    'viewCount': item['statistics'].get('viewCount'),
                    'likeCount': item['statistics'].get('likeCount'),
                    'commentCount': item['statistics'].get('commentCount'),
                    'extractionTimestamp': timestamp
                }
                all_video_data.append(video_data)
        except Exception as e:
            print(f"Error fetching video details: {e}", file=sys.stderr)
    return all_video_data


if __name__ == "__main__":
    print("Starting YouTube data extraction...")

    # Step 1: Check if environment variables are set
    validate_config()

    # Step 2: Create YouTube API client
    youtube = build_youtube_client()

    # Step 3: Get playlist ID for channel uploads
    playlist_id = get_playlist_id(youtube, CHANNEL_ID)
    if not playlist_id:
        print(f"Could not retrieve playlist ID for channel {CHANNEL_ID}.", file=sys.stderr)
        sys.exit(1)

    # Step 4: Retrieve all video IDs from the uploads playlist
    video_ids = get_video_ids(youtube, playlist_id)
    if not video_ids:
        print(f"No videos found for channel {CHANNEL_ID}.", file=sys.stderr)
        sys.exit(1)

    # Step 5: Get detailed video data for all video IDs
    video_data = get_video_details(youtube, video_ids)
    if not video_data:
        print("No video data fetched.", file=sys.stderr)
        sys.exit(1)

    # Step 6: Write the collected video data to the output JSON file
    try:
        with open(OUTPUT_FILENAME, 'w') as f:
            json.dump(video_data, f, indent=4)
        print(f"Data saved to {OUTPUT_FILENAME}")
    except IOError as e:
        print(f"Error writing data to {OUTPUT_FILENAME}: {e}", file=sys.stderr)
        sys.exit(1)

    print("Extraction process finished.")
    sys.exit(0)

import os
import psycopg2
from pyspark.sql import SparkSession
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

POSTGRES_USER = os.getenv("POSTGRES_USER")
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD")
POSTGRES_HOST = os.getenv("POSTGRES_HOST")
POSTGRES_PORT = os.getenv("POSTGRES_PORT")
POSTGRES_DB = os.getenv("POSTGRES_DB")
POSTGRES_SSLMODE = os.getenv("POSTGRES_SSLMODE")

# Initialize Spark Session
spark = SparkSession.builder \
    .appName("YouTube Data Transformation") \
    .config("spark.jars", "/home/main/downloads/postgresql-42.7.3.jar") \
    .getOrCreate()

# Load JSON data
df_raw = spark.read.option("multiline", "true").json("/home/main/youtube/ambuso.json")
df_raw.createOrReplaceTempView("youtube_videos_raw")

# Transform raw data
sql_query = """
    SELECT
        videoId,
        title,
        publishedAt,
        CAST(viewCount AS INT) AS viewCount,
        CAST(likeCount AS INT) AS likeCount,
        CAST(commentCount AS INT) AS commentCount,
        extractionTimestamp
    FROM youtube_videos_raw
    WHERE viewCount IS NOT NULL
    ORDER BY viewCount DESC
"""
df_transformed = spark.sql(sql_query)
df_transformed.createOrReplaceTempView("youtube_transformed")

# Enrich data
enriched_query = """
    SELECT
        videoId,
        title,
        publishedAt,
        viewCount,
        likeCount,
        commentCount,
        extractionTimestamp,
        YEAR(TO_TIMESTAMP(publishedAt)) AS year,
        MONTH(TO_TIMESTAMP(publishedAt)) AS month,
        DAY(TO_TIMESTAMP(publishedAt)) AS day_of_month,
        DATE_FORMAT(TO_TIMESTAMP(publishedAt), 'EEEE') AS day_of_week,
        HOUR(TO_TIMESTAMP(publishedAt)) AS hour,
        WEEKOFYEAR(TO_TIMESTAMP(publishedAt)) AS week,
        CASE
            WHEN viewCount > 100000 THEN 'viral'
            WHEN viewCount > 5000 THEN 'high'
            ELSE 'normal'
        END AS performance_class
    FROM youtube_transformed
"""
df_enriched = spark.sql(enriched_query)
df_enriched.show()

# Create schema and table if needed
try:
    conn = psycopg2.connect(
        dbname=POSTGRES_DB,
        user=POSTGRES_USER,
        password=POSTGRES_PASSWORD,
        host=POSTGRES_HOST,
        port=POSTGRES_PORT,
        sslmode=POSTGRES_SSLMODE
    )
    conn.autocommit = True
    cursor = conn.cursor()

    cursor.execute("CREATE SCHEMA IF NOT EXISTS dataengineering;")
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS dataengineering.youtube_videos_enriched (
            videoId TEXT,
            title TEXT,
            publishedAt TIMESTAMP,
            viewCount INT,
            likeCount INT,
            commentCount INT,
            extractionTimestamp TIMESTAMP,
            year INT,
            month INT,
            day_of_month INT,
            day_of_week TEXT,
            hour INT,
            week INT,
            performance_class TEXT
        );
    """)
    cursor.close()
    conn.close()
    print("Schema and table checked or created.")
except Exception as e:
    print("Failed to create schema or table:", e)

# Write to PostgreSQL
jdbc_url = f"jdbc:postgresql://{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}"
connection_properties = {
    "user": POSTGRES_USER,
    "password": POSTGRES_PASSWORD,
    "driver": "org.postgresql.Driver"
}

df_enriched.write \
    .jdbc(
        url=jdbc_url,
        table="dataengineering.youtube_videos_enriched",
        mode="append",
        properties=connection_properties
    )

print("Enriched data successfully appended to PostgreSQL.")

from airflow import DAG
from airflow.operators.bash import BashOperator
from datetime import datetime

default_args = {
    'start_date': datetime(2025, 6, 7),
    'retries': 1,
}

with DAG(
    'youtube_data_pipeline',
    default_args=default_args,
    description='A pipeline to extract and transform YouTube data',
    schedule='@daily',  # corrected argument
    catchup=False,
) as dag:

    extract_task = BashOperator(
        task_id='extract_youtube_data',
        bash_command='/home/main/youtube/env/bin/python /home/main/youtube/extract.py',
    )

    transform_task = BashOperator(
        task_id='transform_and_load',
        bash_command='/home/main/youtube/env/bin/python /home/main/youtube/transform.py',
    )

    extract_task >> transform_task

SELECT *
FROM (
    SELECT DISTINCT ON ("videoId")
        "videoId",
        "title",
        "viewCount",
        "publishedAt"
    FROM dataengineering.youtube_videos_enriched
    ORDER BY "videoId", "publishedAt" DESC
) sub
ORDER BY "viewCount" DESC
LIMIT 10;