#!/usr/bin/env python3
"""
Google Sheets tracker for dataset processing status.

This module provides functionality to track dataset processing status
in Google Sheets. It's designed to be optional and non-intrusive.
"""

from __future__ import annotations

import logging
import os
from datetime import datetime
from pathlib import Path
from typing import Any

from dotenv import load_dotenv
from google.oauth2 import service_account
from googleapiclient.discovery import build

env_path = Path(__file__).resolve().parent.parent / ".env"
if env_path.exists():
    load_dotenv(dotenv_path=env_path, override=False)

logger = logging.getLogger(__name__)


class GoogleSheetsTracker:
    """Track dataset processing status in Google Sheets."""

    def __init__(
        self,
        spreadsheet_id: str,
        worksheet_name: str = "Sheet1",
        credentials_path: str | None = None,
    ):
        """Initialize Google Sheets tracker.

        Args:
            spreadsheet_id: Google Sheets spreadsheet ID
            worksheet_name: Name of the worksheet to use
            credentials_path: Path to service account JSON credentials
        """
        self.spreadsheet_id = spreadsheet_id
        self.worksheet_name = worksheet_name

        credentials_path_raw = credentials_path or os.getenv(
            "GOOGLE_SHEETS_CREDENTIALS_PATH"
        )

        if credentials_path_raw:
            credentials_path_obj = Path(credentials_path_raw)
            if not credentials_path_obj.is_absolute():
                base_dir = Path(__file__).resolve().parent.parent
                credentials_path_obj = base_dir / credentials_path_raw
            self.credentials_path = str(credentials_path_obj.resolve())
        else:
            self.credentials_path = None

        self.service = None
        self.sheet = None
        self.authenticated = False
        self.dataset_name_col = None
        self.status_col = None
        self.error_col = None
        self.last_shard_col = None
        self.total_rows_col = None
        self.skip_descriptions_col = None
        self.last_updated_col = None
        self.headers = []

    def authenticate(self) -> bool:
        """Authenticate with Google Sheets API.

        Returns:
            True if authentication successful, False otherwise
        """
        if not self.credentials_path:
            return False

        if not os.path.exists(self.credentials_path):
            return False

        try:
            credentials = service_account.Credentials.from_service_account_file(
                self.credentials_path,
                scopes=["https://www.googleapis.com/auth/spreadsheets"],
            )
            self.service = build("sheets", "v4", credentials=credentials)
            self.sheet = self.service.spreadsheets()
            self.authenticated = True
            self._load_headers()
            return True
        except Exception:
            return False

    def _load_headers(self) -> None:
        """Load column headers from the sheet."""
        if self.sheet is None:
            return

        try:
            result = (
                self.sheet.values()
                .get(
                    spreadsheetId=self.spreadsheet_id,
                    range=f"{self.worksheet_name}!1:1",
                )
                .execute()
            )
            values = result.get("values", [])
            if values:
                self.headers = [str(h).strip().casefold() for h in values[0]]
                self.dataset_name_col = self._find_column_index(
                    ["dataset name", "dataset", "dataset_id", "id", "name"]
                )
                self.status_col = self._find_column_index(["status", "state"])
                self.error_col = self._find_column_index(
                    ["error", "error message"]
                )
                self.last_shard_col = self._find_column_index(
                    ["last completed shard", "last_completed_shard",
                     "last shard", "last_shard", "shard"]
                )
                self.total_rows_col = self._find_column_index(
                    ["total rows", "rows", "total_rows"]
                )
                self.skip_descriptions_col = self._find_column_index(
                    ["skip descriptions", "skip_descriptions", "descriptions"]
                )
                self.last_updated_col = self._find_column_index(
                    ["last updated", "last_updated", "updated", "timestamp"]
                )

                if self.dataset_name_col is None:
                    logger.warning("Could not find dataset name column in sheet")
                if self.status_col is None:
                    logger.warning("Could not find status column in sheet")
                if self.error_col is None:
                    logger.warning("Could not find error column in sheet")
        except Exception:
            self.headers = []
            self.dataset_name_col = 0
            self.status_col = 1
            self.error_col = 2

    def _find_column_index(self, possible_names: list[str]) -> int | None:
        """Find column index by name.

        Args:
            possible_names: List of possible column names (will be casefolded)

        Returns:
            Column index or None if not found
        """
        for name in possible_names:
            try:
                return self.headers.index(name.casefold())
            except ValueError:
                continue
        return None

    def get_dataset_status(self, dataset_id: str) -> dict[str, Any]:
        """Get current status of a dataset from the sheet.

        Args:
            dataset_id: Dataset identifier

        Returns:
            Dictionary with status information
        """
        if not self.authenticated or self.sheet is None:
            return {"status": None, "error": None, "row": None}

        try:
            # Calculate maximum column index needed for status data
            max_col = -1
            for col in [
                self.dataset_name_col,
                self.status_col,
                self.error_col,
            ]:
                if col is not None and col > max_col:
                    max_col = col

            # Use dynamic range: A to max_column (or Z if no columns found)
            if max_col >= 0:
                # +1 because _col_letter is 1-indexed
                end_col = self._col_letter(max_col + 1)
                range_str = f"{self.worksheet_name}!A:{end_col}"
            else:
                # Fallback to Z if no status columns found
                range_str = f"{self.worksheet_name}!A:Z"

            result = (
                self.sheet.values()
                .get(
                    spreadsheetId=self.spreadsheet_id,
                    range=range_str,
                )
                .execute()
            )
            values = result.get("values", [])

            if len(values) < 2:
                return {"status": None, "error": None, "row": None}

            for row_idx, row in enumerate(values[1:], start=2):
                if (
                    self.dataset_name_col is not None
                    and len(row) > self.dataset_name_col
                    and str(row[self.dataset_name_col]).strip().casefold()
                    == dataset_id.casefold()
                ):
                    status = (
                        row[self.status_col].strip()
                        if self.status_col is not None and len(row) > self.status_col
                        else None
                    )
                    error = (
                        row[self.error_col].strip()
                        if self.error_col is not None and len(row) > self.error_col
                        else None
                    )
                    return {"status": status, "error": error, "row": row_idx}

            return {"status": None, "error": None, "row": None}
        except Exception:
            return {"status": None, "error": None, "row": None}

    def update_status(
        self,
        dataset_id: str,
        status: str,
        error_message: str | None = None,
    ) -> bool:
        """Update dataset status in the sheet.

        Args:
            dataset_id: Dataset identifier
            status: Status value (e.g., "IN PROGRESS", "DONE", "NOT STARTED")
            error_message: Error message to update in ERROR column

        Returns:
            True if update successful, False otherwise
        """
        if not self.authenticated or self.sheet is None:
            return False

        try:
            dataset_status = self.get_dataset_status(dataset_id)
            row_num = dataset_status.get("row")

            if row_num is None:
                return False

            updates = []

            if self.status_col is not None:
                status_range = (
                    f"{self.worksheet_name}!"
                    f"{self._col_letter(self.status_col + 1)}{row_num}"
                )
                updates.append({"range": status_range, "values": [[status]]})

            if self.error_col is not None:
                if error_message:
                    error_text = str(error_message)[:500]
                    error_range = (
                        f"{self.worksheet_name}!"
                        f"{self._col_letter(self.error_col + 1)}{row_num}"
                    )
                    updates.append({"range": error_range, "values": [[error_text]]})
                else:
                    error_range = (
                        f"{self.worksheet_name}!"
                        f"{self._col_letter(self.error_col + 1)}{row_num}"
                    )
                    updates.append({"range": error_range, "values": [[""]]})

            if updates:
                body = {"valueInputOption": "RAW", "data": updates}
                (
                    self.sheet.values()
                    .batchUpdate(
                        spreadsheetId=self.spreadsheet_id,
                        body=body,
                    )
                    .execute()
                )
                return True

            # No columns to update (both status_col and error_col are None)
            return False
        except Exception:
            return False

    def _col_letter(self, col_num: int) -> str:
        """Convert column number to letter (1 -> A, 2 -> B, etc.).

        Args:
            col_num: Column number (1-indexed)

        Returns:
            Column letter
        """
        result = ""
        while col_num > 0:
            col_num -= 1
            result = chr(65 + (col_num % 26)) + result
            col_num //= 26
        return result

    def check_should_process(
        self, dataset_id: str, force: bool = False
    ) -> tuple[bool, str]:
        """Check if dataset should be processed based on current status.

        Args:
            dataset_id: Dataset identifier
            force: Force processing even if status is "Done" or "In Progress"

        Returns:
            Tuple of (should_process, reason)
        """
        if force:
            return (True, "Force processing enabled")

        if not self.authenticated:
            return (True, "Sheet tracking not available")

        dataset_status = self.get_dataset_status(dataset_id)
        status = dataset_status.get("status", "")

        if not status:
            return (True, "No status found in sheet")

        status_upper = status.upper().strip()

        if status_upper == "DONE":
            return (False, "Dataset already marked as DONE")

        if status_upper == "IN PROGRESS":
            return (False, "Dataset already IN PROGRESS")

        if status_upper == "NOT STARTED":
            return (True, "Status is NOT STARTED, processing")

        return (
            False,
            f"Unknown status '{status_upper}': "
            f"Not processing until status is 'NOT STARTED'",
        )

    def mark_in_progress(self, dataset_id: str) -> bool:
        """Mark dataset as in progress.

        Args:
            dataset_id: Dataset identifier

        Returns:
            True if update successful
        """
        return self.update_status(dataset_id, "IN PROGRESS")

    def mark_completed(self, dataset_id: str) -> bool:
        """Mark dataset as completed.

        Args:
            dataset_id: Dataset identifier

        Returns:
            True if update successful
        """
        return self.update_status(dataset_id, "DONE")

    def mark_error(self, dataset_id: str, error_message: str) -> bool:
        """Mark dataset as error with error message.

        Args:
            dataset_id: Dataset identifier
            error_message: Error message to record in ERROR column

        Returns:
            True if update successful
        """
        return self.update_status(
            dataset_id, "NOT STARTED", error_message=error_message
        )

    def reset_to_not_started(
        self, dataset_id: str, error_message: str | None = None
    ) -> bool:
        """Reset dataset status to NOT STARTED.

        Used when script is interrupted or fails.

        Args:
            dataset_id: Dataset identifier
            error_message: Optional error message to record in ERROR column

        Returns:
            True if update successful
        """
        return self.update_status(
            dataset_id, "NOT STARTED", error_message=error_message
        )

    def get_checkpoint(self, dataset_id: str) -> dict[str, Any] | None:
        """Get checkpoint data for a dataset.

        Args:
            dataset_id: Dataset identifier

        Returns:
            Dictionary with checkpoint data or None if not found
        """
        if not self.authenticated or self.sheet is None:
            return None

        try:
            dataset_status = self.get_dataset_status(dataset_id)
            row_num = dataset_status.get("row")

            if row_num is None:
                return None

            # Calculate maximum column index needed for checkpoint data
            max_col = -1
            for col in [
                self.last_shard_col,
                self.total_rows_col,
                self.skip_descriptions_col,
                self.last_updated_col,
            ]:
                if col is not None and col > max_col:
                    max_col = col

            # Use dynamic range: A to max_column (or Z if no columns found)
            if max_col >= 0:
                # +1 because _col_letter is 1-indexed
                end_col = self._col_letter(max_col + 1)
                range_str = f"{self.worksheet_name}!A{row_num}:{end_col}{row_num}"
            else:
                # Fallback to Z if no checkpoint columns found
                range_str = f"{self.worksheet_name}!A{row_num}:Z{row_num}"

            result = (
                self.sheet.values()
                .get(
                    spreadsheetId=self.spreadsheet_id,
                    range=range_str,
                )
                .execute()
            )
            values = result.get("values", [])

            if not values or not values[0]:
                return None

            row = values[0]
            checkpoint = {}

            if self.last_shard_col is not None and len(row) > self.last_shard_col:
                val = row[self.last_shard_col].strip()
                checkpoint["last_completed_shard"] = int(val) if val else 0

            if self.total_rows_col is not None and len(row) > self.total_rows_col:
                val = row[self.total_rows_col].strip()
                checkpoint["total_rows"] = int(val) if val else 0

            if (
                self.skip_descriptions_col is not None
                and len(row) > self.skip_descriptions_col
            ):
                val = row[self.skip_descriptions_col].strip()
                checkpoint["skip_descriptions"] = int(val) if val else 0

            if self.last_updated_col is not None and len(row) > self.last_updated_col:
                checkpoint["timestamp"] = row[self.last_updated_col].strip()

            return checkpoint if checkpoint else None

        except Exception:
            return None

    def update_checkpoint(
        self,
        dataset_id: str,
        last_shard: int,
        total_rows: int,
        skip_descriptions: int,
    ) -> bool:
        """Update checkpoint data for a dataset.

        Args:
            dataset_id: Dataset identifier
            last_shard: Last completed shard index
            total_rows: Total rows processed
            skip_descriptions: Total descriptions processed

        Returns:
            True if update successful
        """
        if not self.authenticated or self.sheet is None:
            return False

        try:
            dataset_status = self.get_dataset_status(dataset_id)
            row_num = dataset_status.get("row")

            if row_num is None:
                return False

            updates = []

            if self.last_shard_col is not None:
                shard_range = (
                    f"{self.worksheet_name}!"
                    f"{self._col_letter(self.last_shard_col + 1)}{row_num}"
                )
                updates.append({"range": shard_range, "values": [[last_shard]]})

            if self.total_rows_col is not None:
                rows_range = (
                    f"{self.worksheet_name}!"
                    f"{self._col_letter(self.total_rows_col + 1)}{row_num}"
                )
                updates.append({"range": rows_range, "values": [[total_rows]]})

            if self.skip_descriptions_col is not None:
                desc_range = (
                    f"{self.worksheet_name}!"
                    f"{self._col_letter(self.skip_descriptions_col + 1)}{row_num}"
                )
                updates.append(
                    {"range": desc_range, "values": [[skip_descriptions]]}
                )

            if self.last_updated_col is not None:
                timestamp = datetime.now().isoformat()
                updated_range = (
                    f"{self.worksheet_name}!"
                    f"{self._col_letter(self.last_updated_col + 1)}{row_num}"
                )
                updates.append({"range": updated_range, "values": [[timestamp]]})

            if updates:
                body = {"valueInputOption": "RAW", "data": updates}
                (
                    self.sheet.values()
                    .batchUpdate(
                        spreadsheetId=self.spreadsheet_id,
                        body=body,
                    )
                    .execute()
                )

            return True
        except Exception:
            return False


def create_tracker(
    spreadsheet_id: str | None = None,
    worksheet_name: str | None = None,
    credentials_path: str | None = None,
) -> GoogleSheetsTracker | None:
    """Create and initialize Google Sheets tracker.

    Args:
        spreadsheet_id: Google Sheets spreadsheet ID
        worksheet_name: Name of the worksheet
        credentials_path: Path to service account credentials

    Returns:
        GoogleSheetsTracker instance or None if initialization fails
    """
    if not spreadsheet_id:
        spreadsheet_id = os.getenv("GOOGLE_SHEETS_ID")

    if not spreadsheet_id:
        return None

    if not worksheet_name:
        worksheet_name = os.getenv("GOOGLE_SHEETS_WORKSHEET_NAME", "Sheet1")

    try:
        tracker = GoogleSheetsTracker(
            spreadsheet_id=spreadsheet_id,
            worksheet_name=worksheet_name,
            credentials_path=credentials_path,
        )
        if tracker.authenticate():
            return tracker
    except Exception:
        pass

    return None