ir/sheets_storage.py

"""
Google Sheets storage for UAE KB feedback data.

This module provides persistent storage for demo feedback using Google Sheets,
which allows easy viewing and analysis of user feedback data.
"""

import os
import json
from datetime import datetime
from typing import Optional, Dict, Any, List
import logging

# Import gspread at module level
try:
    import gspread
    from google.oauth2.service_account import Credentials
    GSPREAD_AVAILABLE = True
except ImportError:
    GSPREAD_AVAILABLE = False
    gspread = None
    Credentials = None

logger = logging.getLogger(__name__)

# Google Sheets configuration
SHEET_ID = "1wwnpQRx32nbmApNK1LyvzUDIVx-tW4mpwgRPLvjm4gw"
RATINGS_WORKSHEET = "ratings"
NOTES_WORKSHEET = "notes"

# Lazy-loaded client
_sheets_client = None
_spreadsheet = None


def _get_credentials():
    """Get Google credentials from environment variable."""
    if not GSPREAD_AVAILABLE:
        logger.warning("gspread not installed")
        return None

    creds_json = os.environ.get("GOOGLE_SHEETS_CREDENTIALS")
    if not creds_json:
        logger.warning("GOOGLE_SHEETS_CREDENTIALS not set")
        return None

    try:
        creds_dict = json.loads(creds_json)
        scopes = [
            "https://www.googleapis.com/auth/spreadsheets",
            "https://www.googleapis.com/auth/drive"
        ]
        return Credentials.from_service_account_info(creds_dict, scopes=scopes)
    except Exception as e:
        logger.error(f"Failed to load credentials: {e}")
        return None


def _get_spreadsheet():
    """Get or create the Google Sheets client."""
    global _spreadsheet

    if _spreadsheet is not None:
        return _spreadsheet

    if not GSPREAD_AVAILABLE:
        return None

    try:
        creds = _get_credentials()
        if creds is None:
            return None

        client = gspread.authorize(creds)
        _spreadsheet = client.open_by_key(SHEET_ID)
        logger.info(f"Connected to Google Sheet: {SHEET_ID}")
        return _spreadsheet
    except Exception as e:
        logger.error(f"Failed to connect to Google Sheets: {e}")
        return None


def _get_or_create_worksheet(name: str, headers: List[str]):
    """Get or create a worksheet with headers."""
    spreadsheet = _get_spreadsheet()
    if spreadsheet is None:
        return None

    try:
        # Try to get existing worksheet
        try:
            worksheet = spreadsheet.worksheet(name)
            # Check if headers exist (first row)
            first_row = worksheet.row_values(1)
            if not first_row or first_row[0] != headers[0]:
                # Insert headers at row 1
                worksheet.insert_row(headers, 1)
                logger.info(f"Added headers to worksheet: {name}")
        except gspread.WorksheetNotFound:
            # Create new worksheet
            worksheet = spreadsheet.add_worksheet(title=name, rows=1000, cols=len(headers))
            worksheet.append_row(headers)
            logger.info(f"Created worksheet: {name}")

        return worksheet
    except Exception as e:
        logger.error(f"Failed to get/create worksheet {name}: {e}")
        return None


def save_rating_to_sheets(
    query: str,
    category: str,
    entity_id: str,
    entity_name: str,
    rank: int,
    score: float,
    rating: str,
    page: int,
    client_ip: str = "unknown",
    comment: str = "",
    query_id: str = ""
) -> bool:
    """
    Save an entity rating to Google Sheets.

    Returns True if successful, False otherwise.
    """
    headers = [
        "timestamp", "query_id", "client_ip", "query", "category",
        "entity_id", "entity_name", "rank", "score", "rating", "page", "comment"
    ]

    worksheet = _get_or_create_worksheet(RATINGS_WORKSHEET, headers)
    if worksheet is None:
        return False

    try:
        row = [
            datetime.now().isoformat(),
            query_id,
            client_ip,
            query,
            category,
            entity_id,
            entity_name,
            rank,
            round(score, 4) if isinstance(score, float) else score,
            rating,
            page,
            comment
        ]
        worksheet.append_row(row)
        logger.info(f"Saved rating: {entity_name} - {rating}")
        return True
    except Exception as e:
        logger.error(f"Failed to save rating: {e}")
        return False


def save_notes_to_sheets(
    query: str,
    category: str,
    notes: str,
    client_ip: str = "unknown"
) -> bool:
    """
    Save notes/feedback to Google Sheets.

    Returns True if successful, False otherwise.
    """
    headers = ["timestamp", "client_ip", "query", "category", "notes"]

    worksheet = _get_or_create_worksheet(NOTES_WORKSHEET, headers)
    if worksheet is None:
        return False

    try:
        row = [
            datetime.now().isoformat(),
            client_ip,
            query,
            category,
            notes
        ]
        worksheet.append_row(row)
        logger.info(f"Saved notes for query: {query[:50]}...")
        return True
    except Exception as e:
        logger.error(f"Failed to save notes: {e}")
        return False


def is_sheets_enabled() -> bool:
    """Check if Google Sheets storage is configured and working."""
    return GSPREAD_AVAILABLE and os.environ.get("GOOGLE_SHEETS_CREDENTIALS") is not None


def test_connection() -> Dict[str, Any]:
    """Test the Google Sheets connection."""
    result = {
        "configured": is_sheets_enabled(),
        "connected": False,
        "sheet_id": SHEET_ID,
        "error": None
    }

    if not result["configured"]:
        result["error"] = "GOOGLE_SHEETS_CREDENTIALS not set"
        return result

    try:
        spreadsheet = _get_spreadsheet()
        if spreadsheet:
            result["connected"] = True
            result["title"] = spreadsheet.title
            result["worksheets"] = [ws.title for ws in spreadsheet.worksheets()]
    except Exception as e:
        result["error"] = str(e)

    return result