import os
import spaces  # Added: Import spaces for ZeroGPU
import pandas as pd
import numpy as np
import logging
from typing import Tuple, List, Dict, Any
from app.storage import get_file_info, get_job, update_job_progress, complete_job, fail_job
from app.model import model_instance

logger = logging.getLogger("aku_sentiment_utils")

def parse_dataset_file(file_path: str) -> Tuple[pd.DataFrame, List[str], int, List[Dict[str, Any]]]:
    """
    Parses a CSV or XLSX file, handles encoding fallback, extracts column names and sample rows.
    """
    ext = os.path.splitext(file_path)[1].lower()

    if ext == ".csv":
        # Try UTF-8 first, fallback to latin-1
        try:
            df = pd.read_csv(file_path, encoding="utf-8")
        except UnicodeDecodeError:
            df = pd.read_csv(file_path, encoding="latin-1")
    elif ext in [".xlsx", ".xls"]:
        df = pd.read_excel(file_path)
    else:
        raise ValueError(f"Unsupported file format: {ext}. Only .csv and .xlsx are allowed.")

    columns = [str(col).strip() for col in df.columns]
    df.columns = columns

    total_rows = len(df)

    # Generate sample records (up to first 5 rows) for previewing column selection
    sample_df = df.head(5).replace({np.nan: None})
    sample_records = sample_df.to_dict(orient="records")

    return df, columns, total_rows, sample_records


@spaces.GPU  # Added: GPU decorator to allow inference on ZeroGPU
def run_background_sentiment_job(job_id: str):
    """
    Background worker task function executed in a separate thread/executor.
    Processes dataset in batches of size 32, updates job progress, and calculates metrics.
    """
    try:
        job = get_job(job_id)
        if not job:
            logger.error(f"Job ID {job_id} not found in storage.")
            return

        file_info = get_file_info(job["file_id"])
        if not file_info:
            fail_job(job_id, "Uploaded file reference not found.")
            return

        file_path = file_info["file_path"]
        text_column = job["text_column"]

        # Parse full dataframe
        df, columns, total_rows, _ = parse_dataset_file(file_path)

        if text_column not in df.columns:
            fail_job(job_id, f"Selected column '{text_column}' does not exist in dataset.")
            return

        text_series = df[text_column].astype(str).tolist()
        batch_size = 32

        sentiments = []
        scores = []

        # Process in batches to provide fine-grained progress updates
        for i in range(0, total_rows, batch_size):
            batch_texts = text_series[i:i + batch_size]
            batch_predictions = model_instance.predict_batch(batch_texts, batch_size=batch_size)

            for pred in batch_predictions:
                sentiments.append(pred["label"])
                scores.append(pred["score"])

            processed_count = min(i + batch_size, total_rows)
            update_job_progress(job_id, processed_count, total_rows)

        # Append new columns to original DataFrame
        output_df = df.copy()
        output_df["sentiment"] = sentiments
        output_df["sentiment_score"] = scores

        # Calculate summary statistics
        pos_count = sentiments.count("Positive")
        neu_count = sentiments.count("Neutral")
        neg_count = sentiments.count("Negative")

        pos_pct = round((pos_count / total_rows * 100), 2) if total_rows > 0 else 0.0
        neu_pct = round((neu_count / total_rows * 100), 2) if total_rows > 0 else 0.0
        neg_pct = round((neg_count / total_rows * 100), 2) if total_rows > 0 else 0.0

        summary = {
            "total_rows": total_rows,
            "positive": pos_count,
            "neutral": neu_count,
            "negative": neg_count,
            "positive_pct": pos_pct,
            "neutral_pct": neu_pct,
            "negative_pct": neg_pct
        }

        # Prepare a sample preview for the UI table (first 500 rows maximum)
        preview_df = output_df.head(500).replace({np.nan: None})
        preview_records = preview_df.to_dict(orient="records")

        # Save output and mark job complete
        complete_job(job_id, summary, preview_records, output_df)
        logger.info(f"Job {job_id} successfully completed. {total_rows} rows processed.")

    except Exception as e:
        logger.error(f"Error processing job {job_id}: {str(e)}", exc_info=True)
        fail_job(job_id, str(e))