# app.py

import torch
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    RobertaTokenizer
)
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from googleapiclient.discovery import build
import config  # Ensure you have a config.py with YOUTUBE_API_KEY
import logging
import warnings
import re
import emoji

# ---------------------------
# Logging Configuration
# ---------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# ---------------------------
# Suppress specific warnings (Optional)
# ---------------------------
warnings.filterwarnings('ignore', message='Some weights of the model checkpoint')
warnings.filterwarnings('ignore', message='Please make sure the config includes `forced_bos_token_id=0`')
warnings.filterwarnings('ignore', message='The tokenizer class you load from this checkpoint is not the same type')

# ---------------------------
# FastAPI app initialization
# ---------------------------
app = FastAPI(
    title="YouTube Comments Emotion & Sarcasm Analyzer",
    description="Analyzes YouTube comments for emotions and sarcasm.",
    version="1.0.0"
)

# ---------------------------
# Pydantic model to handle requests
# ---------------------------
class EmotionRequest(BaseModel):
    video_id: str
    max_comments: int = 100  # Default to 100 if not specified
    page: int = 1
    page_size: int = 10

# ---------------------------
# Utility Functions
# ---------------------------
def demojize_text(text):
    """
    Convert emojis in text to textual descriptions.
    """
    return emoji.demojize(text, delimiters=(" ", " "))

def normalize_text(text):
    """
    Normalize slang and abbreviations in the text.
    """
    text = re.sub(r'\br\b', ' are ', text)
    text = re.sub(r'\bu\b', ' you ', text)
    text = re.sub(r'\btho\b', ' though ', text)
    # Add more normalization rules as needed
    return text

def clean_text(text):
    """
    Clean the input text by removing URLs, mentions, special characters, and converting emojis.
    """
    text = re.sub(r'http\S+', '', text)          # Remove URLs
    text = re.sub(r'@\w+', '', text)             # Remove mentions
    text = re.sub(r'[^A-Za-z0-9\s]', '', text)   # Remove special characters
    text = demojize_text(text)                   # Convert emojis to text
    text = normalize_text(text)                  # Normalize slang and abbreviations
    text = text.lower()                           # Convert to lowercase
    return text

# ---------------------------
# Analyzer class for emotion and sarcasm
# ---------------------------
class EmotionSarcasmAnalyzer:
    def __init__(self, device: torch.device):
        self.device = device

        # ---------------------------
        # Emotion Model and Tokenizer Setup
        # ---------------------------
        self.emotion_model = AutoModelForSequenceClassification.from_pretrained(
            'nateraw/bert-base-uncased-emotion',
            ignore_mismatched_sizes=True  # To ignore unused weights
        ).to(self.device)
        self.emotion_tokenizer = AutoTokenizer.from_pretrained(
            'nateraw/bert-base-uncased-emotion'
        )
        self.emotion_labels = ['anger', 'joy', 'optimism', 'sadness']
        logger.info("Emotion model initialized.")

        # ---------------------------
        # Sarcasm Model and Tokenizer Setup
        # ---------------------------
        self.sarcasm_model = AutoModelForSequenceClassification.from_pretrained(
            'cardiffnlp/twitter-roberta-base-sentiment-latest',
            ignore_mismatched_sizes=True  # To ignore unused weights
        ).to(self.device)
        self.sarcasm_tokenizer = RobertaTokenizer.from_pretrained(
            'cardiffnlp/twitter-roberta-base-sentiment-latest'
        )
        self.sarcasm_labels = ['non-sarcastic', 'sarcastic']
        logger.info("Sarcasm model initialized.")

    def preprocess_text(self, text: str) -> str:
        """
        Preprocess the text for model input.
        """
        return clean_text(text)

    def analyze_emotion(self, text: str) -> str:
        """
        Analyze emotion from the text.
        """
        try:
            inputs = self.emotion_tokenizer(
                text,
                return_tensors="pt",
                padding=True,
                truncation=True
            ).to(self.device)
            with torch.no_grad():
                outputs = self.emotion_model(**inputs)
            logits = outputs.logits
            predicted_class_id = torch.argmax(logits, dim=1).item()
            emotion = self.emotion_labels[predicted_class_id]
            logger.debug(f"Emotion prediction: {emotion} for text: {text}")
            return emotion
        except Exception as e:
            logger.error(f"Error in emotion analysis: {e}")
            return "error"

    def analyze_sarcasm(self, text: str) -> str:
        """
        Analyze sarcasm from the text.
        """
        try:
            inputs = self.sarcasm_tokenizer(
                text,
                return_tensors="pt",
                padding=True,
                truncation=True
            ).to(self.device)
            with torch.no_grad():
                outputs = self.sarcasm_model(**inputs)
            logits = outputs.logits
            predicted_class_id = torch.argmax(logits, dim=1).item()
            sarcasm = self.sarcasm_labels[predicted_class_id]
            logger.debug(f"Sarcasm prediction: {sarcasm} for text: {text}")
            return sarcasm
        except Exception as e:
            logger.error(f"Error in sarcasm analysis: {e}")
            return "error"

    def analyze_comment(self, comment: str) -> dict:
        """
        Analyze a single comment for emotion and sarcasm.
        """
        try:
            preprocessed = self.preprocess_text(comment)
            emotion = self.analyze_emotion(preprocessed)
            sarcasm = self.analyze_sarcasm(preprocessed)
            return {
                "comment": comment,
                "emotion": emotion,
                "sarcasm": sarcasm
            }
        except Exception as e:
            logger.error(f"Error analyzing comment: {e}")
            return {
                "comment": comment,
                "emotion": "error",
                "sarcasm": "error"
            }

# ---------------------------
# Function to fetch comments from YouTube
# ---------------------------
def fetch_comments(video_id: str, max_results: int) -> List[str]:
    """
    Fetch comments from a YouTube video using YouTube Data API.
    """
    youtube = build("youtube", "v3", developerKey=config.YOUTUBE_API_KEY)
    comments = []
    try:
        request = youtube.commentThreads().list(
            part="snippet",
            videoId=video_id,
            textFormat="plainText",
            maxResults=min(max_results, 100)  # YouTube API max is 100 per request
        )
        response = request.execute()
        while request is not None and len(comments) < max_results:
            for item in response.get("items", []):
                comment = item["snippet"]["topLevelComment"]["snippet"]["textDisplay"]
                comments.append(comment)
                if len(comments) >= max_results:
                    break
            request = youtube.commentThreads().list_next(request, response)
    except Exception as e:
        logger.error(f"Error fetching comments: {e}")
    logger.info(f"Fetched {len(comments)} comments from video ID: {video_id}")
    return comments[:max_results]

# ---------------------------
# Device Configuration
# ---------------------------
def get_device() -> torch.device:
    """
    Returns the available device (GPU if available, else CPU).
    """
    return torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

device = get_device()
logger.info(f"Using device: {device}")

# ---------------------------
# Instantiate the Analyzer
# ---------------------------
analyzer = EmotionSarcasmAnalyzer(device=device)

# ---------------------------
# FastAPI Endpoint: Analyze Emotion and Sarcasm
# ---------------------------
@app.post("/analyze_emotion_sarcasm/", summary="Analyze Emotions and Sarcasm in YouTube Comments")
async def analyze_emotion_sarcasm(request: EmotionRequest):
    """
    Endpoint to analyze emotions and sarcasm in YouTube comments.
    """
    # Validate pagination parameters
    if request.page < 1 or request.page_size < 1:
        logger.warning("Invalid pagination parameters received.")
        raise HTTPException(status_code=400, detail="Page and page_size must be positive integers.")

    # Fetch comments from YouTube
    comments = fetch_comments(request.video_id, request.max_comments)
    if not comments:
        logger.warning(f"No comments found for video ID: {request.video_id}")
        raise HTTPException(status_code=404, detail="No comments found or unable to fetch comments.")

    logger.info(f"Analyzing {len(comments)} comments.")

    # Analyze each comment
    results = []
    for comment in comments:
        analysis_result = analyzer.analyze_comment(comment)
        results.append(analysis_result)

    # Pagination
    start_idx = (request.page - 1) * request.page_size
    end_idx = start_idx + request.page_size
    paginated_results = results[start_idx:end_idx]

    logger.info(f"Returning comments {start_idx + 1} to {min(end_idx, len(results))}.")

    return {
        "video_id": request.video_id,
        "max_comments": request.max_comments, 
        "page": request.page,
        "page_size": request.page_size,
        "total_comments_analyzed": len(results),
        "results": paginated_results
    }  

# ---------------------------
# Root Endpoint for Health Check
# ---------------------------
@app.get("/", summary="Health Check")
async def root():
    return {"message": "YouTube Comments Emotion & Sarcasm Analyzer is runn
I think this is my current script, somehow the output its good for sarcasm sclassification, but for emotion detection the model is not good enough,  also it extracts emojis but doesn't label them... or involved in the analysis