#!/usr/bin/env python
"""
Google Spreadsheet MCP Server
A Model Context Protocol (MCP) server built with FastMCP for interacting with Google Sheets.
"""

import base64
import os
import sys
from typing import List, Dict, Any, Optional, Union
import json
from dataclasses import dataclass
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator

# MCP imports
from mcp.server.fastmcp import FastMCP, Context
from mcp.types import ToolAnnotations

# Google API imports
from google.oauth2.credentials import Credentials
from google.oauth2 import service_account
from google.auth.transport.requests import Request
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
import google.auth

# Constants
SCOPES = ['https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive']
CREDENTIALS_CONFIG = os.environ.get('CREDENTIALS_CONFIG')
TOKEN_PATH = os.environ.get('TOKEN_PATH', 'token.json')
CREDENTIALS_PATH = os.environ.get('CREDENTIALS_PATH', 'credentials.json')
SERVICE_ACCOUNT_PATH = os.environ.get('SERVICE_ACCOUNT_PATH', 'service_account.json')
DRIVE_FOLDER_ID = os.environ.get('DRIVE_FOLDER_ID', '')  # Working directory in Google Drive

# Tool filtering configuration
# Parse enabled tools from environment variable or command-line argument
def _parse_enabled_tools() -> Optional[set]:
    """
    Parse enabled tools from ENABLED_TOOLS environment variable or --include-tools argument.
    Returns None if all tools should be enabled (default behavior).
    Returns a set of tool names if filtering is requested.
    """
    # Check command-line arguments first
    enabled_tools_str = None
    for i, arg in enumerate(sys.argv):
        if arg == '--include-tools' and i + 1 < len(sys.argv):
            enabled_tools_str = sys.argv[i + 1]
            break
    
    # Fall back to environment variable
    if not enabled_tools_str:
        enabled_tools_str = os.environ.get('ENABLED_TOOLS')
    
    if not enabled_tools_str:
        return None  # No filtering, enable all tools
    
    # Parse comma-separated list and normalize
    tools = {tool.strip() for tool in enabled_tools_str.split(',') if tool.strip()}
    return tools if tools else None

ENABLED_TOOLS = _parse_enabled_tools()

@dataclass
class SpreadsheetContext:
    """Context for Google Spreadsheet service"""
    sheets_service: Any
    drive_service: Any
    folder_id: Optional[str] = None


@asynccontextmanager
async def spreadsheet_lifespan(server: FastMCP) -> AsyncIterator[SpreadsheetContext]:
    """Manage Google Spreadsheet API connection lifecycle"""
    # Authenticate and build the service
    creds = None

    if CREDENTIALS_CONFIG:
        creds = service_account.Credentials.from_service_account_info(json.loads(base64.b64decode(CREDENTIALS_CONFIG)), scopes=SCOPES)
    
    # Check for explicit service account authentication first (custom SERVICE_ACCOUNT_PATH)
    if not creds and SERVICE_ACCOUNT_PATH and os.path.exists(SERVICE_ACCOUNT_PATH):
        try:
            # Regular service account authentication
            creds = service_account.Credentials.from_service_account_file(
                SERVICE_ACCOUNT_PATH,
                scopes=SCOPES
            )
            print("Using service account authentication")
            print(f"Working with Google Drive folder ID: {DRIVE_FOLDER_ID or 'Not specified'}")
        except Exception as e:
            print(f"Error using service account authentication: {e}")
            creds = None
    
    # Fall back to OAuth flow if service account auth failed or not configured
    if not creds:
        print("Trying OAuth authentication flow")
        if os.path.exists(TOKEN_PATH):
            with open(TOKEN_PATH, 'r') as token:
                creds = Credentials.from_authorized_user_info(json.load(token), SCOPES)
                
        # If credentials are not valid or don't exist, get new ones
        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                try:
                    print("Attempting to refresh expired token...")
                    creds.refresh(Request())
                    print("Token refreshed successfully")
                    # Save the refreshed token
                    with open(TOKEN_PATH, 'w') as token:
                        token.write(creds.to_json())
                except Exception as refresh_error:
                    print(f"Token refresh failed: {refresh_error}")
                    print("Triggering reauthentication flow...")
                    creds = None  # Clear creds to trigger OAuth flow below

            # If refresh failed or creds don't exist, run OAuth flow
            if not creds:
                try:
                    flow = InstalledAppFlow.from_client_secrets_file(CREDENTIALS_PATH, SCOPES)
                    creds = flow.run_local_server(port=0)

                    # Save the credentials for the next run
                    with open(TOKEN_PATH, 'w') as token:
                        token.write(creds.to_json())
                    print("Successfully authenticated using OAuth flow")
                except Exception as e:
                    print(f"Error with OAuth flow: {e}")
                    creds = None
    
    # Try Application Default Credentials if no creds thus far
    # This will automatically check GOOGLE_APPLICATION_CREDENTIALS, gcloud auth, and metadata service
    if not creds:
        try:
            print("Attempting to use Application Default Credentials (ADC)")
            print("ADC will check: GOOGLE_APPLICATION_CREDENTIALS, gcloud auth, and metadata service")
            creds, project = google.auth.default(
                scopes=SCOPES
            )
            print(f"Successfully authenticated using ADC for project: {project}")
        except Exception as e:
            print(f"Error using Application Default Credentials: {e}")
            print("WARNING: No authentication configured. Server will start but API calls will fail.")
            print("Please configure SERVICE_ACCOUNT_PATH, CREDENTIALS_PATH, or gcloud ADC.")

    # Build the services (or None if no creds)
    sheets_service = build('sheets', 'v4', credentials=creds) if creds else None
    drive_service = build('drive', 'v3', credentials=creds) if creds else None

    try:
        # Provide the service in the context
        yield SpreadsheetContext(
            sheets_service=sheets_service,
            drive_service=drive_service,
            folder_id=DRIVE_FOLDER_ID if DRIVE_FOLDER_ID else None
        )
    finally:
        # No explicit cleanup needed for Google APIs
        pass


# Initialize the MCP server with lifespan management
# Resolve host/port from environment variables with flexible names
_resolved_host = os.environ.get('HOST') or os.environ.get('FASTMCP_HOST') or "0.0.0.0"
_resolved_port_str = os.environ.get('PORT') or os.environ.get('FASTMCP_PORT') or "8000"
try:
    _resolved_port = int(_resolved_port_str)
except ValueError:
    _resolved_port = 8000

# Initialize the MCP server with explicit host/port to ensure binding as configured
mcp = FastMCP("Google Spreadsheet",
              dependencies=["google-auth", "google-auth-oauthlib", "google-api-python-client"],
              lifespan=spreadsheet_lifespan,
              host=_resolved_host,
              port=_resolved_port)


def tool(annotations: Optional[ToolAnnotations] = None):
    """
    Conditional tool decorator that only registers tools if they're enabled.
    
    This wrapper checks ENABLED_TOOLS configuration and only applies the @mcp.tool
    decorator if the tool should be enabled. If ENABLED_TOOLS is None (default),
    all tools are enabled.
    
    Args:
        annotations: Optional ToolAnnotations for the tool
    
    Returns:
        Decorator function
    """
    def decorator(func):
        tool_name = func.__name__
        
        # If no filtering is configured, or if this tool is in the enabled list
        if ENABLED_TOOLS is None or tool_name in ENABLED_TOOLS:
            # Apply the mcp.tool decorator
            if annotations:
                return mcp.tool(annotations=annotations)(func)
            else:
                return mcp.tool()(func)
        else:
            # Don't register this tool - return the function undecorated
            return func
    
    return decorator


@tool(
    annotations=ToolAnnotations(
        title="Get Sheet Data",
        readOnlyHint=True,
    ),
)
def get_sheet_data(spreadsheet_id: str,
                   sheet: str,
                   range: Optional[str] = None,
                   include_grid_data: bool = False,
                   ctx: Context = None) -> Dict[str, Any]:
    """
    Get data from a specific sheet in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet
        range: Optional cell range in A1 notation (e.g., 'A1:C10'). If not provided, gets all data.
        include_grid_data: If True, includes cell formatting and other metadata in the response.
            Note: Setting this to True will significantly increase the response size and token usage
            when parsing the response, as it includes detailed cell formatting information.
            Default is False (returns values only, more efficient).
    
    Returns:
        Grid data structure with either full metadata or just values from Google Sheets API, depending on include_grid_data parameter
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service

    # Construct the range - keep original API behavior
    if range:
        full_range = f"{sheet}!{range}"
    else:
        full_range = sheet
    
    if include_grid_data:
        # Use full API to get all grid data including formatting
        result = sheets_service.spreadsheets().get(
            spreadsheetId=spreadsheet_id,
            ranges=[full_range],
            includeGridData=True
        ).execute()
    else:
        # Use values API to get cell values only (more efficient)
        values_result = sheets_service.spreadsheets().values().get(
            spreadsheetId=spreadsheet_id,
            range=full_range
        ).execute()
        
        # Format the response to match expected structure
        result = {
            'spreadsheetId': spreadsheet_id,
            'valueRanges': [{
                'range': full_range,
                'values': values_result.get('values', [])
            }]
        }

    return result

@tool(
    annotations=ToolAnnotations(
        title="Get Sheet Formulas",
        readOnlyHint=True,
    ),
)
def get_sheet_formulas(spreadsheet_id: str,
                       sheet: str,
                       range: Optional[str] = None,
                       ctx: Context = None) -> List[List[Any]]:
    """
    Get formulas from a specific sheet in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet
        range: Optional cell range in A1 notation (e.g., 'A1:C10'). If not provided, gets all formulas from the sheet.
    
    Returns:
        A 2D array of the sheet formulas.
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Construct the range
    if range:
        full_range = f"{sheet}!{range}"
    else:
        full_range = sheet  # Get all formulas in the specified sheet
    
    # Call the Sheets API
    result = sheets_service.spreadsheets().values().get(
        spreadsheetId=spreadsheet_id,
        range=full_range,
        valueRenderOption='FORMULA'  # Request formulas
    ).execute()
    
    # Get the formulas from the response
    formulas = result.get('values', [])
    return formulas

@tool(
    annotations=ToolAnnotations(
        title="Update Cells",
        destructiveHint=True,
    ),
)
def update_cells(spreadsheet_id: str,
                sheet: str,
                range: str,
                data: List[List[Any]],
                ctx: Context = None) -> Dict[str, Any]:
    """
    Update cells in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet
        range: Cell range in A1 notation (e.g., 'A1:C10')
        data: 2D array of values to update
    
    Returns:
        Result of the update operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Construct the range
    full_range = f"{sheet}!{range}"
    
    # Prepare the value range object
    value_range_body = {
        'values': data
    }
    
    # Call the Sheets API to update values
    result = sheets_service.spreadsheets().values().update(
        spreadsheetId=spreadsheet_id,
        range=full_range,
        valueInputOption='USER_ENTERED',
        body=value_range_body
    ).execute()
    
    return result


@tool(
    annotations=ToolAnnotations(
        title="Batch Update Cells",
        destructiveHint=True,
    ),
)
def batch_update_cells(spreadsheet_id: str,
                       sheet: str,
                       ranges: Dict[str, List[List[Any]]],
                       ctx: Context = None) -> Dict[str, Any]:
    """
    Batch update multiple ranges in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet
        ranges: Dictionary mapping range strings to 2D arrays of values
               e.g., {'A1:B2': [[1, 2], [3, 4]], 'D1:E2': [['a', 'b'], ['c', 'd']]}
    
    Returns:
        Result of the batch update operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Prepare the batch update request
    data = []
    for range_str, values in ranges.items():
        full_range = f"{sheet}!{range_str}"
        data.append({
            'range': full_range,
            'values': values
        })
    
    batch_body = {
        'valueInputOption': 'USER_ENTERED',
        'data': data
    }
    
    # Call the Sheets API to perform batch update
    result = sheets_service.spreadsheets().values().batchUpdate(
        spreadsheetId=spreadsheet_id,
        body=batch_body
    ).execute()
    
    return result


@tool(
    annotations=ToolAnnotations(
        title="Add Rows",
        destructiveHint=True,
    ),
)
def add_rows(spreadsheet_id: str,
             sheet: str,
             count: int,
             start_row: Optional[int] = None,
             ctx: Context = None) -> Dict[str, Any]:
    """
    Add rows to a sheet in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet
        count: Number of rows to add
        start_row: 0-based row index to start adding. If not provided, adds at the beginning.
    
    Returns:
        Result of the operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Get sheet ID
    spreadsheet = sheets_service.spreadsheets().get(spreadsheetId=spreadsheet_id).execute()
    sheet_id = None
    
    for s in spreadsheet['sheets']:
        if s['properties']['title'] == sheet:
            sheet_id = s['properties']['sheetId']
            break
            
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}
    
    # Prepare the insert rows request
    request_body = {
        "requests": [
            {
                "insertDimension": {
                    "range": {
                        "sheetId": sheet_id,
                        "dimension": "ROWS",
                        "startIndex": start_row if start_row is not None else 0,
                        "endIndex": (start_row if start_row is not None else 0) + count
                    },
                    "inheritFromBefore": start_row is not None and start_row > 0
                }
            }
        ]
    }
    
    # Execute the request
    result = sheets_service.spreadsheets().batchUpdate(
        spreadsheetId=spreadsheet_id,
        body=request_body
    ).execute()
    
    return result


@tool(
    annotations=ToolAnnotations(
        title="Add Columns",
        destructiveHint=True,
    ),
)
def add_columns(spreadsheet_id: str,
                sheet: str,
                count: int,
                start_column: Optional[int] = None,
                ctx: Context = None) -> Dict[str, Any]:
    """
    Add columns to a sheet in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet
        count: Number of columns to add
        start_column: 0-based column index to start adding. If not provided, adds at the beginning.
    
    Returns:
        Result of the operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Get sheet ID
    spreadsheet = sheets_service.spreadsheets().get(spreadsheetId=spreadsheet_id).execute()
    sheet_id = None
    
    for s in spreadsheet['sheets']:
        if s['properties']['title'] == sheet:
            sheet_id = s['properties']['sheetId']
            break
            
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}
    
    # Prepare the insert columns request
    request_body = {
        "requests": [
            {
                "insertDimension": {
                    "range": {
                        "sheetId": sheet_id,
                        "dimension": "COLUMNS",
                        "startIndex": start_column if start_column is not None else 0,
                        "endIndex": (start_column if start_column is not None else 0) + count
                    },
                    "inheritFromBefore": start_column is not None and start_column > 0
                }
            }
        ]
    }
    
    # Execute the request
    result = sheets_service.spreadsheets().batchUpdate(
        spreadsheetId=spreadsheet_id,
        body=request_body
    ).execute()
    
    return result


@tool(
    annotations=ToolAnnotations(
        title="List Sheets",
        readOnlyHint=True,
    ),
)
def list_sheets(spreadsheet_id: str, ctx: Context = None) -> List[str]:
    """
    List all sheets in a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
    
    Returns:
        List of sheet names
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Get spreadsheet metadata
    spreadsheet = sheets_service.spreadsheets().get(spreadsheetId=spreadsheet_id).execute()
    
    # Extract sheet names
    sheet_names = [sheet['properties']['title'] for sheet in spreadsheet['sheets']]
    
    return sheet_names


@tool(
    annotations=ToolAnnotations(
        title="Copy Sheet",
        destructiveHint=True,
    ),
)
def copy_sheet(src_spreadsheet: str,
               src_sheet: str,
               dst_spreadsheet: str,
               dst_sheet: str,
               ctx: Context = None) -> Dict[str, Any]:
    """
    Copy a sheet from one spreadsheet to another.
    
    Args:
        src_spreadsheet: Source spreadsheet ID
        src_sheet: Source sheet name
        dst_spreadsheet: Destination spreadsheet ID
        dst_sheet: Destination sheet name
    
    Returns:
        Result of the operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Get source sheet ID
    src = sheets_service.spreadsheets().get(spreadsheetId=src_spreadsheet).execute()
    src_sheet_id = None
    
    for s in src['sheets']:
        if s['properties']['title'] == src_sheet:
            src_sheet_id = s['properties']['sheetId']
            break
            
    if src_sheet_id is None:
        return {"error": f"Source sheet '{src_sheet}' not found"}
    
    # Copy the sheet to destination spreadsheet
    copy_result = sheets_service.spreadsheets().sheets().copyTo(
        spreadsheetId=src_spreadsheet,
        sheetId=src_sheet_id,
        body={
            "destinationSpreadsheetId": dst_spreadsheet
        }
    ).execute()
    
    # If destination sheet name is different from the default copied name, rename it
    if 'title' in copy_result and copy_result['title'] != dst_sheet:
        # Get the ID of the newly copied sheet
        copy_sheet_id = copy_result['sheetId']
        
        # Rename the copied sheet
        rename_request = {
            "requests": [
                {
                    "updateSheetProperties": {
                        "properties": {
                            "sheetId": copy_sheet_id,
                            "title": dst_sheet
                        },
                        "fields": "title"
                    }
                }
            ]
        }
        
        rename_result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=dst_spreadsheet,
            body=rename_request
        ).execute()
        
        return {
            "copy": copy_result,
            "rename": rename_result
        }
    
    return {"copy": copy_result}


@tool(
    annotations=ToolAnnotations(
        title="Rename Sheet",
        destructiveHint=True,
    ),
)
def rename_sheet(spreadsheet: str,
                 sheet: str,
                 new_name: str,
                 ctx: Context = None) -> Dict[str, Any]:
    """
    Rename a sheet in a Google Spreadsheet.
    
    Args:
        spreadsheet: Spreadsheet ID
        sheet: Current sheet name
        new_name: New sheet name
    
    Returns:
        Result of the operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Get sheet ID
    spreadsheet_data = sheets_service.spreadsheets().get(spreadsheetId=spreadsheet).execute()
    sheet_id = None
    
    for s in spreadsheet_data['sheets']:
        if s['properties']['title'] == sheet:
            sheet_id = s['properties']['sheetId']
            break
            
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}
    
    # Prepare the rename request
    request_body = {
        "requests": [
            {
                "updateSheetProperties": {
                    "properties": {
                        "sheetId": sheet_id,
                        "title": new_name
                    },
                    "fields": "title"
                }
            }
        ]
    }
    
    # Execute the request
    result = sheets_service.spreadsheets().batchUpdate(
        spreadsheetId=spreadsheet,
        body=request_body
    ).execute()
    
    return result


@tool(
    annotations=ToolAnnotations(
        title="Get Multiple Sheet Data",
        readOnlyHint=True,
    ),
)
def get_multiple_sheet_data(queries: List[Dict[str, str]],
                            ctx: Context = None) -> List[Dict[str, Any]]:
    """
    Get data from multiple specific ranges in Google Spreadsheets.
    
    Args:
        queries: A list of dictionaries, each specifying a query. 
                 Each dictionary should have 'spreadsheet_id', 'sheet', and 'range' keys.
                 Example: [{'spreadsheet_id': 'abc', 'sheet': 'Sheet1', 'range': 'A1:B5'}, 
                           {'spreadsheet_id': 'xyz', 'sheet': 'Data', 'range': 'C1:C10'}]
    
    Returns:
        A list of dictionaries, each containing the original query parameters 
        and the fetched 'data' or an 'error'.
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    results = []
    
    for query in queries:
        spreadsheet_id = query.get('spreadsheet_id')
        sheet = query.get('sheet')
        range_str = query.get('range')
        
        if not all([spreadsheet_id, sheet, range_str]):
            results.append({**query, 'error': 'Missing required keys (spreadsheet_id, sheet, range)'})
            continue

        try:
            # Construct the range
            full_range = f"{sheet}!{range_str}"
            
            # Call the Sheets API
            result = sheets_service.spreadsheets().values().get(
                spreadsheetId=spreadsheet_id,
                range=full_range
            ).execute()
            
            # Get the values from the response
            values = result.get('values', [])
            results.append({**query, 'data': values})

        except Exception as e:
            results.append({**query, 'error': str(e)})
            
    return results


@tool(
    annotations=ToolAnnotations(
        title="Get Multiple Spreadsheet Summary",
        readOnlyHint=True,
    ),
)
def get_multiple_spreadsheet_summary(spreadsheet_ids: List[str],
                                   rows_to_fetch: int = 5,
                                   ctx: Context = None) -> List[Dict[str, Any]]:
    """
    Get a summary of multiple Google Spreadsheets, including sheet names, 
    headers, and the first few rows of data for each sheet.
    
    Args:
        spreadsheet_ids: A list of spreadsheet IDs to summarize.
        rows_to_fetch: The number of rows (including header) to fetch for the summary (default: 5).
    
    Returns:
        A list of dictionaries, each representing a spreadsheet summary. 
        Includes spreadsheet title, sheet summaries (title, headers, first rows), or an error.
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    summaries = []
    
    for spreadsheet_id in spreadsheet_ids:
        summary_data = {
            'spreadsheet_id': spreadsheet_id,
            'title': None,
            'sheets': [],
            'error': None
        }
        try:
            # Get spreadsheet metadata
            spreadsheet = sheets_service.spreadsheets().get(
                spreadsheetId=spreadsheet_id,
                fields='properties.title,sheets(properties(title,sheetId))'
            ).execute()
            
            summary_data['title'] = spreadsheet.get('properties', {}).get('title', 'Unknown Title')
            
            sheet_summaries = []
            for sheet in spreadsheet.get('sheets', []):
                sheet_title = sheet.get('properties', {}).get('title')
                sheet_id = sheet.get('properties', {}).get('sheetId')
                sheet_summary = {
                    'title': sheet_title,
                    'sheet_id': sheet_id,
                    'headers': [],
                    'first_rows': [],
                    'error': None
                }
                
                if not sheet_title:
                    sheet_summary['error'] = 'Sheet title not found'
                    sheet_summaries.append(sheet_summary)
                    continue
                    
                try:
                    # Fetch the first few rows (e.g., A1:Z5)
                    # Adjust range if fewer rows are requested
                    max_row = max(1, rows_to_fetch) # Ensure at least 1 row is fetched
                    range_to_get = f"{sheet_title}!A1:{max_row}" # Fetch all columns up to max_row
                    
                    result = sheets_service.spreadsheets().values().get(
                        spreadsheetId=spreadsheet_id,
                        range=range_to_get
                    ).execute()
                    
                    values = result.get('values', [])
                    
                    if values:
                        sheet_summary['headers'] = values[0]
                        if len(values) > 1:
                            sheet_summary['first_rows'] = values[1:max_row]
                    else:
                        # Handle empty sheets or sheets with less data than requested
                        sheet_summary['headers'] = []
                        sheet_summary['first_rows'] = []

                except Exception as sheet_e:
                    sheet_summary['error'] = f'Error fetching data for sheet {sheet_title}: {sheet_e}'
                
                sheet_summaries.append(sheet_summary)
            
            summary_data['sheets'] = sheet_summaries
            
        except Exception as e:
            summary_data['error'] = f'Error fetching spreadsheet {spreadsheet_id}: {e}'
            
        summaries.append(summary_data)
        
    return summaries


@mcp.resource("spreadsheet://{spreadsheet_id}/info")
def get_spreadsheet_info(spreadsheet_id: str) -> str:
    """
    Get basic information about a Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet
    
    Returns:
        JSON string with spreadsheet information
    """
    # Access the context through mcp.get_lifespan_context() for resources
    context = mcp.get_lifespan_context()
    sheets_service = context.sheets_service
    
    # Get spreadsheet metadata
    spreadsheet = sheets_service.spreadsheets().get(spreadsheetId=spreadsheet_id).execute()
    
    # Extract relevant information
    info = {
        "title": spreadsheet.get('properties', {}).get('title', 'Unknown'),
        "sheets": [
            {
                "title": sheet['properties']['title'],
                "sheetId": sheet['properties']['sheetId'],
                "gridProperties": sheet['properties'].get('gridProperties', {})
            }
            for sheet in spreadsheet.get('sheets', [])
        ]
    }
    
    return json.dumps(info, indent=2)


@tool(
    annotations=ToolAnnotations(
        title="Create Spreadsheet",
        destructiveHint=True,
    ),
)
def create_spreadsheet(title: str, folder_id: Optional[str] = None, ctx: Context = None) -> Dict[str, Any]:
    """
    Create a new Google Spreadsheet.
    
    Args:
        title: The title of the new spreadsheet
        folder_id: Optional Google Drive folder ID where the spreadsheet should be created.
                  If not provided, uses the configured default folder or creates in root.
    
    Returns:
        Information about the newly created spreadsheet including its ID
    """
    drive_service = ctx.request_context.lifespan_context.drive_service
    # Use provided folder_id or fall back to configured default
    target_folder_id = folder_id or ctx.request_context.lifespan_context.folder_id

    # Create the spreadsheet
    file_body = {
        'name': title,
        'mimeType': 'application/vnd.google-apps.spreadsheet',
    }
    if target_folder_id:
        file_body['parents'] = [target_folder_id]
    
    spreadsheet = drive_service.files().create(
        supportsAllDrives=True,
        body=file_body,
        fields='id, name, parents'
    ).execute()

    spreadsheet_id = spreadsheet.get('id')
    parents = spreadsheet.get('parents')
    folder_info = f" in folder {target_folder_id}" if target_folder_id else " in root"
    print(f"Spreadsheet created with ID: {spreadsheet_id}{folder_info}")

    return {
        'spreadsheetId': spreadsheet_id,
        'title': spreadsheet.get('name', title),
        'folder': parents[0] if parents else 'root',
    }


@tool(
    annotations=ToolAnnotations(
        title="Create Sheet",
        destructiveHint=True,
    ),
)
def create_sheet(spreadsheet_id: str,
                title: str,
                ctx: Context = None) -> Dict[str, Any]:
    """
    Create a new sheet tab in an existing Google Spreadsheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet
        title: The title for the new sheet
    
    Returns:
        Information about the newly created sheet
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Define the add sheet request
    request_body = {
        "requests": [
            {
                "addSheet": {
                    "properties": {
                        "title": title
                    }
                }
            }
        ]
    }
    
    # Execute the request
    result = sheets_service.spreadsheets().batchUpdate(
        spreadsheetId=spreadsheet_id,
        body=request_body
    ).execute()
    
    # Extract the new sheet information
    new_sheet_props = result['replies'][0]['addSheet']['properties']
    
    return {
        'sheetId': new_sheet_props['sheetId'],
        'title': new_sheet_props['title'],
        'index': new_sheet_props.get('index'),
        'spreadsheetId': spreadsheet_id
    }


@tool(
    annotations=ToolAnnotations(
        title="List Spreadsheets",
        readOnlyHint=True,
    ),
)
def list_spreadsheets(folder_id: Optional[str] = None, ctx: Context = None) -> List[Dict[str, str]]:
    """
    List all spreadsheets in the specified Google Drive folder.
    If no folder is specified, uses the configured default folder or lists from 'My Drive'.
    
    Args:
        folder_id: Optional Google Drive folder ID to search in.
                  If not provided, uses the configured default folder or searches 'My Drive'.
    
    Returns:
        List of spreadsheets with their ID and title
    """
    drive_service = ctx.request_context.lifespan_context.drive_service
    # Use provided folder_id or fall back to configured default
    target_folder_id = folder_id or ctx.request_context.lifespan_context.folder_id
    
    query = "mimeType='application/vnd.google-apps.spreadsheet'"
    
    # If a specific folder is provided or configured, search only in that folder
    if target_folder_id:
        query += f" and '{target_folder_id}' in parents"
        print(f"Searching for spreadsheets in folder: {target_folder_id}")
    else:
        print("Searching for spreadsheets in 'My Drive'")
    
    # List spreadsheets
    results = drive_service.files().list(
        q=query,
        spaces='drive',
        includeItemsFromAllDrives=True,
        supportsAllDrives=True,
        fields='files(id, name)',
        orderBy='modifiedTime desc'
    ).execute()
    
    spreadsheets = results.get('files', [])
    
    return [{'id': sheet['id'], 'title': sheet['name']} for sheet in spreadsheets]


@tool(
    annotations=ToolAnnotations(
        title="Share Spreadsheet",
        destructiveHint=True,
    ),
)
def share_spreadsheet(spreadsheet_id: str,
                      recipients: List[Dict[str, str]],
                      send_notification: bool = True,
                      ctx: Context = None) -> Dict[str, List[Dict[str, Any]]]:
    """
    Share a Google Spreadsheet with multiple users via email, assigning specific roles.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet to share.
        recipients: A list of dictionaries, each containing 'email_address' and 'role'.
                    The role should be one of: 'reader', 'commenter', 'writer'.
                    Example: [
                        {'email_address': 'user1@example.com', 'role': 'writer'},
                        {'email_address': 'user2@example.com', 'role': 'reader'}
                    ]
        send_notification: Whether to send a notification email to the users. Defaults to True.

    Returns:
        A dictionary containing lists of 'successes' and 'failures'. 
        Each item in the lists includes the email address and the outcome.
    """
    drive_service = ctx.request_context.lifespan_context.drive_service
    successes = []
    failures = []
    
    for recipient in recipients:
        email_address = recipient.get('email_address')
        role = recipient.get('role', 'writer') # Default to writer if role is missing for an entry
        
        if not email_address:
            failures.append({
                'email_address': None,
                'error': 'Missing email_address in recipient entry.'
            })
            continue
            
        if role not in ['reader', 'commenter', 'writer']:
             failures.append({
                'email_address': email_address,
                'error': f"Invalid role '{role}'. Must be 'reader', 'commenter', or 'writer'."
            })
             continue

        permission = {
            'type': 'user',
            'role': role,
            'emailAddress': email_address
        }
        
        try:
            result = drive_service.permissions().create(
                fileId=spreadsheet_id,
                body=permission,
                sendNotificationEmail=send_notification,
                fields='id'
            ).execute()
            successes.append({
                'email_address': email_address, 
                'role': role, 
                'permissionId': result.get('id')
            })
        except Exception as e:
            # Try to provide a more informative error message
            error_details = str(e)
            if hasattr(e, 'content'):
                try:
                    error_content = json.loads(e.content)
                    error_details = error_content.get('error', {}).get('message', error_details)
                except json.JSONDecodeError:
                    pass # Keep the original error string
            failures.append({
                'email_address': email_address,
                'error': f"Failed to share: {error_details}"
            })
            
    return {"successes": successes, "failures": failures}


@tool(
    annotations=ToolAnnotations(
        title="List Folders",
        readOnlyHint=True,
    ),
)
def list_folders(parent_folder_id: Optional[str] = None, ctx: Context = None) -> List[Dict[str, str]]:
    """
    List all folders in the specified Google Drive folder.
    If no parent folder is specified, lists folders from 'My Drive' root.
    
    Args:
        parent_folder_id: Optional Google Drive folder ID to search within.
                         If not provided, searches the root of 'My Drive'.
    
    Returns:
        List of folders with their ID, name, and parent information
    """
    drive_service = ctx.request_context.lifespan_context.drive_service
    
    query = "mimeType='application/vnd.google-apps.folder'"
    
    # If a specific parent folder is provided, search only within that folder
    if parent_folder_id:
        query += f" and '{parent_folder_id}' in parents"
        print(f"Searching for folders in parent folder: {parent_folder_id}")
    else:
        # Search in root of My Drive (folders that don't have any parent folders)
        query += " and 'root' in parents"
        print("Searching for folders in 'My Drive' root")
    
    # List folders
    results = drive_service.files().list(
        q=query,
        spaces='drive',
        includeItemsFromAllDrives=True,
        supportsAllDrives=True,
        fields='files(id, name, parents)',
        orderBy='name'
    ).execute()
    
    folders = results.get('files', [])
    
    return [
        {
            'id': folder['id'], 
            'name': folder['name'],
            'parent': folder.get('parents', ['root'])[0] if folder.get('parents') else 'root'
        } 
        for folder in folders
    ]




@tool(
    annotations=ToolAnnotations(
        title="Search Spreadsheets by Name or Content",
        readOnlyHint=True,
    ),
)
def search_spreadsheets(query: str,
                        max_results: int = 20,
                        ctx: Context = None) -> List[Dict[str, Any]]:
    """
    Search for spreadsheets in Google Drive by name or content.

    Args:
        query: Search query string. Searches in file name and content.
               Examples: "budget 2024", "sales report", "project tracker"
        max_results: Maximum number of results to return (default 20, max 100)

    Returns:
        List of matching spreadsheets with their ID, name, and metadata
    """
    drive_service = ctx.request_context.lifespan_context.drive_service

    # Limit max_results to reasonable bounds
    max_results = min(max(1, max_results), 100)

    # Build the search query for Google Drive
    # Search only for spreadsheets and match the query in name or fullText
    search_query = (
        f"mimeType='application/vnd.google-apps.spreadsheet' and "
        f"(name contains '{query}' or fullText contains '{query}')"
    )

    try:
        results = drive_service.files().list(
            q=search_query,
            pageSize=max_results,
            spaces='drive',
            includeItemsFromAllDrives=True,
            supportsAllDrives=True,
            fields='files(id, name, createdTime, modifiedTime, owners, webViewLink)',
            orderBy='modifiedTime desc'
        ).execute()

        files = results.get('files', [])

        return [
            {
                'id': f['id'],
                'name': f['name'],
                'created_time': f.get('createdTime'),
                'modified_time': f.get('modifiedTime'),
                'owners': [owner.get('emailAddress') for owner in f.get('owners', [])],
                'web_link': f.get('webViewLink')
            }
            for f in files
        ]
    except Exception as e:
        return [{'error': f'Search failed: {str(e)}'}]


def _column_index_to_letter(index: int) -> str:
    """Convert 0-based column index to A1 notation letter (0='A', 25='Z', 26='AA', etc.)"""
    result = ""
    while index >= 0:
        result = chr(index % 26 + ord('A')) + result
        index = index // 26 - 1
    return result


def _letter_to_column_index(letter: str) -> int:
    """Convert A1 notation letter to 0-based column index ('A'=0, 'Z'=25, 'AA'=26, etc.)"""
    result = 0
    for char in letter.upper():
        result = result * 26 + (ord(char) - ord('A') + 1)
    return result - 1


def _parse_a1_notation(range_str: str) -> Dict[str, int]:
    """
    Parse A1 notation range to row/column indices.
    
    Args:
        range_str: A1 notation range (e.g., 'A1:C10')
    
    Returns:
        Dictionary containing applicable indices based on the range format.
        May include: startRowIndex, endRowIndex, startColumnIndex, endColumnIndex.
        Not all keys are present for all range formats (e.g., 'A:B' has no row indices).
    """
    import re
    
    # Match patterns like A1, A1:B2, A:B, 1:10
    match = re.match(r'^([A-Z]+)?(\d+)?(?::([A-Z]+)?(\d+)?)?$', range_str.upper())
    
    if not match:
        raise ValueError(f"Invalid A1 notation: {range_str}")
    
    start_col, start_row, end_col, end_row = match.groups()
    
    result = {}
    
    # Start column
    if start_col:
        result['startColumnIndex'] = _letter_to_column_index(start_col)
    
    # Start row (A1 notation is 1-based, convert to 0-based)
    if start_row:
        result['startRowIndex'] = int(start_row) - 1
    
    # End column (exclusive in API, so add 1)
    if end_col:
        result['endColumnIndex'] = _letter_to_column_index(end_col) + 1
    elif start_col:
        result['endColumnIndex'] = result['startColumnIndex'] + 1
    
    # End row (exclusive in API, so no -1 needed)
    if end_row:
        result['endRowIndex'] = int(end_row)
    elif start_row:
        result['endRowIndex'] = result['startRowIndex'] + 1
    
    return result


def _get_sheet_id(sheets_service: Any, spreadsheet_id: str, sheet_name: str) -> Optional[int]:
    """
    Get the sheet ID for a given sheet name.
    
    Args:
        sheets_service: Google Sheets service instance
        spreadsheet_id: The spreadsheet ID
        sheet_name: The name of the sheet
    
    Returns:
        The sheet ID, or None if not found
    """
    try:
        spreadsheet = sheets_service.spreadsheets().get(
            spreadsheetId=spreadsheet_id,
            fields='sheets(properties(title,sheetId))'
        ).execute()
        
        for sheet in spreadsheet.get('sheets', []):
            if sheet['properties']['title'] == sheet_name:
                return sheet['properties']['sheetId']
        
        return None
    except Exception:
        return None


@tool(
    annotations=ToolAnnotations(
        title="Find Cells",
        readOnlyHint=True,
    ),
)
def find_in_spreadsheet(spreadsheet_id: str,
                        query: str,
                        sheet: Optional[str] = None,
                        case_sensitive: bool = False,
                        max_results: int = 50,
                        ctx: Context = None) -> List[Dict[str, Any]]:
    """
    Find cells containing a specific value in a Google Spreadsheet.

    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        query: The text to search for in cell values
        sheet: Optional sheet name to search in. If not provided, searches all sheets.
        case_sensitive: Whether the search should be case-sensitive (default False)
        max_results: Maximum number of results to return (default 50)

    Returns:
        List of found cells with their location (sheet, cell in A1 notation) and value
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    results = []

    try:
        # Get spreadsheet metadata to find all sheets
        spreadsheet = sheets_service.spreadsheets().get(
            spreadsheetId=spreadsheet_id,
            fields='sheets(properties(title,sheetId))'
        ).execute()

        sheets_to_search = []
        for s in spreadsheet.get('sheets', []):
            sheet_title = s.get('properties', {}).get('title')
            if sheet is None or sheet_title == sheet:
                sheets_to_search.append(sheet_title)

        if not sheets_to_search:
            return [{'error': f"Sheet '{sheet}' not found"}]

        search_query = query if case_sensitive else query.lower()

        for sheet_name in sheets_to_search:
            if len(results) >= max_results:
                break

            # Get all data from the sheet
            response = sheets_service.spreadsheets().values().get(
                spreadsheetId=spreadsheet_id,
                range=sheet_name
            ).execute()

            values = response.get('values', [])

            for row_idx, row in enumerate(values):
                if len(results) >= max_results:
                    break

                for col_idx, cell_value in enumerate(row):
                    if len(results) >= max_results:
                        break

                    cell_str = str(cell_value)
                    compare_value = cell_str if case_sensitive else cell_str.lower()

                    if search_query in compare_value:
                        cell_ref = f"{_column_index_to_letter(col_idx)}{row_idx + 1}"
                        results.append({
                            'sheet': sheet_name,
                            'cell': cell_ref,
                            'value': cell_value
                        })

        return results

    except Exception as e:
        return [{'error': f'Search failed: {str(e)}'}]


@tool(
    annotations=ToolAnnotations(
        title="Batch Update",
        destructiveHint=True,
    ),
)
def batch_update(spreadsheet_id: str,
                 requests: List[Dict[str, Any]],
                 ctx: Context = None) -> Dict[str, Any]:
    """
    Execute a batch update on a Google Spreadsheet using the full batchUpdate endpoint.
    This provides access to all batchUpdate operations including adding sheets, updating properties,
    inserting/deleting dimensions, formatting, and more.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        requests: A list of request objects. Each request object can contain any valid batchUpdate operation.
                 Common operations include:
                 - addSheet: Add a new sheet
                 - updateSheetProperties: Update sheet properties (title, grid properties, etc.)
                 - insertDimension: Insert rows or columns
                 - deleteDimension: Delete rows or columns
                 - updateCells: Update cell values and formatting
                 - updateBorders: Update cell borders
                 - addConditionalFormatRule: Add conditional formatting
                 - deleteConditionalFormatRule: Remove conditional formatting
                 - updateDimensionProperties: Update row/column properties
                 - and many more...
                 
                 Example requests:
                 [
                     {
                         "addSheet": {
                             "properties": {
                                 "title": "New Sheet"
                             }
                         }
                     },
                     {
                         "updateSheetProperties": {
                             "properties": {
                                 "sheetId": 0,
                                 "title": "Renamed Sheet"
                             },
                             "fields": "title"
                         }
                     },
                     {
                         "insertDimension": {
                             "range": {
                                 "sheetId": 0,
                                 "dimension": "ROWS",
                                 "startIndex": 1,
                                 "endIndex": 3
                             }
                         }
                     }
                 ]
    
    Returns:
        Result of the batch update operation, including replies for each request
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Validate input
    if not requests:
        return {"error": "requests list cannot be empty"}
    
    if not all(isinstance(req, dict) for req in requests):
        return {"error": "Each request must be a dictionary"}
    
    # Prepare the batch update request body
    request_body = {
        "requests": requests
    }
    
    # Execute the batch update
    result = sheets_service.spreadsheets().batchUpdate(
        spreadsheetId=spreadsheet_id,
        body=request_body
    ).execute()
    
    return result


@tool(
    annotations=ToolAnnotations(
        title="Add Chart",
        destructiveHint=True,
    ),
)
def add_chart(spreadsheet_id: str,
              sheet: str,
              chart_type: str,
              data_range: str,
              title: Optional[str] = None,
              x_axis_label: Optional[str] = None,
              y_axis_label: Optional[str] = None,
              position_x: int = 0,
              position_y: int = 0,
              width: int = 600,
              height: int = 400,
              ctx: Context = None) -> Dict[str, Any]:
    """
    Add a chart to a Google Spreadsheet.
    
    Creates a chart from the specified data range with customizable type, title, and positioning.
    The chart is added as a floating element on the sheet.
    
    Args:
        spreadsheet_id: The ID of the spreadsheet (found in the URL)
        sheet: The name of the sheet containing the data
        chart_type: Type of chart to create. Supported types:
                   - COLUMN: Vertical bar chart
                   - BAR: Horizontal bar chart
                   - LINE: Line chart
                   - AREA: Area chart
                   - PIE: Pie chart
                   - SCATTER: Scatter plot
                   - COMBO: Combination chart
                   - HISTOGRAM: Histogram
        data_range: A1 notation range for chart data (e.g., 'A1:C10'). 
                   The first row is typically treated as headers.
        title: Optional title for the chart
        x_axis_label: Optional label for the X axis (bottom axis)
        y_axis_label: Optional label for the Y axis (left axis)
        position_x: Horizontal position offset in pixels from the top-left corner (default: 0)
        position_y: Vertical position offset in pixels from the top-left corner (default: 0)
        width: Width of the chart in pixels (default: 600)
        height: Height of the chart in pixels (default: 400)
    
    Returns:
        Result of the chart creation operation
    
    Examples:
        Create a column chart showing sales data:
        add_chart(
            spreadsheet_id="abc123",
            sheet="Sales",
            chart_type="COLUMN",
            data_range="A1:B13",
            title="Monthly Sales",
            x_axis_label="Month",
            y_axis_label="Revenue ($)"
        )
        
        Create a pie chart for market share:
        add_chart(
            spreadsheet_id="abc123",
            sheet="Market",
            chart_type="PIE",
            data_range="A1:B5",
            title="Market Share by Product"
        )
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    
    # Validate chart type
    valid_chart_types = ['COLUMN', 'BAR', 'LINE', 'AREA', 'PIE', 'SCATTER', 'COMBO', 'HISTOGRAM']
    if chart_type.upper() not in valid_chart_types:
        return {
            "error": f"Invalid chart type '{chart_type}'. Must be one of: {', '.join(valid_chart_types)}"
        }
    
    chart_type = chart_type.upper()
    
    # Get sheet ID
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found in spreadsheet"}
    
    # Parse the A1 notation range
    try:
        range_indices = _parse_a1_notation(data_range)
    except ValueError as e:
        return {"error": str(e)}
    
    # Build the source range for the chart
    source_range = {
        "sheetId": sheet_id,
        **range_indices
    }
    
    # Build chart specification based on chart type
    # Note: For basic charts, using the same source_range for both domains and series
    # allows the API to automatically interpret the first column as the domain (X-axis labels)
    # and subsequent columns as data series (Y-axis values). This is the standard behavior
    # for most chart types and works correctly for typical use cases.
    if chart_type == "PIE":
        # Pie charts use a different spec structure
        chart_spec = {
            "pieChart": {
                "legendPosition": "RIGHT_LEGEND",
                "domain": {
                    "sourceRange": {
                        "sources": [source_range]
                    }
                },
                "series": {
                    "sourceRange": {
                        "sources": [source_range]
                    }
                }
            }
        }
        if title:
            chart_spec["title"] = title
    else:
        # All other chart types use basicChart spec
        chart_spec = {
            "basicChart": {
                "chartType": chart_type,
                "legendPosition": "RIGHT_LEGEND",
                "axis": [],
                "domains": [{
                    "domain": {
                        "sourceRange": {
                            "sources": [source_range]
                        }
                    }
                }],
                "series": [{
                    "series": {
                        "sourceRange": {
                            "sources": [source_range]
                        }
                    },
                    "targetAxis": "LEFT_AXIS"
                }],
                "headerCount": 1
            }
        }
        
        # Add title if provided
        if title:
            chart_spec["title"] = title
        
        # Add axis labels if provided
        if x_axis_label:
            chart_spec["basicChart"]["axis"].append({
                "position": "BOTTOM_AXIS",
                "title": x_axis_label
            })
        else:
            chart_spec["basicChart"]["axis"].append({
                "position": "BOTTOM_AXIS"
            })
        
        if y_axis_label:
            chart_spec["basicChart"]["axis"].append({
                "position": "LEFT_AXIS",
                "title": y_axis_label
            })
        else:
            chart_spec["basicChart"]["axis"].append({
                "position": "LEFT_AXIS"
            })
    
    # Build the add chart request
    request_body = {
        "requests": [{
            "addChart": {
                "chart": {
                    "spec": chart_spec,
                    "position": {
                        "overlayPosition": {
                            "anchorCell": {
                                "sheetId": sheet_id,
                                "rowIndex": 0,
                                "columnIndex": 0
                            },
                            "offsetXPixels": position_x,
                            "offsetYPixels": position_y,
                            "widthPixels": width,
                            "heightPixels": height
                        }
                    }
                }
            }
        }]
    }
    
    # Execute the request
    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body=request_body
        ).execute()
        
        return {
            "success": True,
            "message": f"Chart '{title or chart_type}' added successfully",
            "chartId": result.get('replies', [{}])[0].get('addChart', {}).get('chart', {}).get('chartId'),
            "result": result
        }
    except Exception as e:
        return {
            "error": f"Failed to add chart: {str(e)}"
        }


# =============================================================================
# Extended tools (added by mcp-google-sheets-extended)
# =============================================================================

@tool(
    annotations=ToolAnnotations(
        title="Upload Excel File",
        destructiveHint=True,
    ),
)
def upload_excel(
    file_path: str,
    title: Optional[str] = None,
    folder_id: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Upload an Excel (.xlsx/.xls) or CSV file and convert it to a Google Spreadsheet.

    Args:
        file_path: Local file path to the Excel or CSV file
        title: Optional title for the new spreadsheet (defaults to the filename)
        folder_id: Optional Google Drive folder ID to place the file in

    Returns:
        The created spreadsheet's ID, title, and URL
    """
    from googleapiclient.http import MediaFileUpload
    import mimetypes

    drive_service = ctx.request_context.lifespan_context.drive_service
    target_folder = folder_id or ctx.request_context.lifespan_context.folder_id

    if not os.path.exists(file_path):
        return {"error": f"File not found: {file_path}"}

    basename = os.path.basename(file_path)
    name = title or os.path.splitext(basename)[0]

    mime_map = {
        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        ".xls": "application/vnd.ms-excel",
        ".csv": "text/csv",
    }
    ext = os.path.splitext(basename)[1].lower()
    upload_mime = mime_map.get(ext) or mimetypes.guess_type(file_path)[0] or "application/octet-stream"

    file_metadata: Dict[str, Any] = {
        "name": name,
        "mimeType": "application/vnd.google-apps.spreadsheet",  # auto-convert
    }
    if target_folder:
        file_metadata["parents"] = [target_folder]

    try:
        media = MediaFileUpload(file_path, mimetype=upload_mime, resumable=True)
        created = drive_service.files().create(
            body=file_metadata,
            media_body=media,
            fields="id,name,webViewLink",
        ).execute()
        return {
            "success": True,
            "spreadsheet_id": created["id"],
            "title": created["name"],
            "url": created.get("webViewLink"),
        }
    except Exception as e:
        return {"error": f"Upload failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Export Sheet",
        readOnlyHint=True,
    ),
)
def export_sheet(
    spreadsheet_id: str,
    format: str = "xlsx",
    output_path: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Export (download) a Google Spreadsheet as Excel, CSV, or PDF.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        format: Export format — 'xlsx', 'csv', 'pdf', or 'ods'
        output_path: Local file path to save to. Auto-generated if omitted.

    Returns:
        The local file path where the file was saved
    """
    drive_service = ctx.request_context.lifespan_context.drive_service

    mime_map = {
        "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        "csv": "text/csv",
        "pdf": "application/pdf",
        "ods": "application/vnd.oasis.opendocument.spreadsheet",
    }
    export_mime = mime_map.get(format.lower())
    if not export_mime:
        return {"error": f"Unsupported format: {format}. Use xlsx, csv, pdf, or ods."}

    try:
        # Get spreadsheet title for default filename
        meta = drive_service.files().get(fileId=spreadsheet_id, fields="name").execute()
        default_name = meta.get("name", spreadsheet_id)

        if not output_path:
            output_path = f"{default_name}.{format.lower()}"

        from googleapiclient.http import MediaIoBaseDownload
        import io

        request = drive_service.files().export_media(fileId=spreadsheet_id, mimeType=export_mime)
        fh = io.BytesIO()
        downloader = MediaIoBaseDownload(fh, request)
        done = False
        while not done:
            _, done = downloader.next_chunk()

        with open(output_path, "wb") as f:
            f.write(fh.getvalue())

        return {"success": True, "file_path": os.path.abspath(output_path), "format": format}
    except Exception as e:
        return {"error": f"Export failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Format Cells",
        destructiveHint=True,
    ),
)
def format_cells(
    spreadsheet_id: str,
    sheet: str,
    range: str,
    bold: Optional[bool] = None,
    italic: Optional[bool] = None,
    font_size: Optional[int] = None,
    font_color: Optional[str] = None,
    background_color: Optional[str] = None,
    number_format: Optional[str] = None,
    horizontal_alignment: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Apply formatting to a range of cells.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range (e.g., 'A1:C10')
        bold: Make text bold
        italic: Make text italic
        font_size: Font size in pt
        font_color: Hex color for text (e.g., '#FF0000')
        background_color: Hex color for cell background (e.g., '#FFFF00')
        number_format: Number format pattern (e.g., '#,##0.00', '0%', 'yyyy-mm-dd')
        horizontal_alignment: 'LEFT', 'CENTER', or 'RIGHT'

    Returns:
        Result of the formatting operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    def hex_to_color(hex_str: str) -> Dict[str, float]:
        h = hex_str.lstrip("#")
        return {
            "red": int(h[0:2], 16) / 255.0,
            "green": int(h[2:4], 16) / 255.0,
            "blue": int(h[4:6], 16) / 255.0,
        }

    cell_format: Dict[str, Any] = {}
    fields = []

    text_format: Dict[str, Any] = {}
    if bold is not None:
        text_format["bold"] = bold
        fields.append("userEnteredFormat.textFormat.bold")
    if italic is not None:
        text_format["italic"] = italic
        fields.append("userEnteredFormat.textFormat.italic")
    if font_size is not None:
        text_format["fontSize"] = font_size
        fields.append("userEnteredFormat.textFormat.fontSize")
    if font_color is not None:
        text_format["foregroundColorStyle"] = {"rgbColor": hex_to_color(font_color)}
        fields.append("userEnteredFormat.textFormat.foregroundColorStyle")
    if text_format:
        cell_format["textFormat"] = text_format

    if background_color is not None:
        cell_format["backgroundColorStyle"] = {"rgbColor": hex_to_color(background_color)}
        fields.append("userEnteredFormat.backgroundColorStyle")

    if number_format is not None:
        cell_format["numberFormat"] = {"type": "NUMBER", "pattern": number_format}
        fields.append("userEnteredFormat.numberFormat")

    if horizontal_alignment is not None:
        cell_format["horizontalAlignment"] = horizontal_alignment.upper()
        fields.append("userEnteredFormat.horizontalAlignment")

    if not fields:
        return {"error": "No formatting options specified"}

    grid_range = {"sheetId": sheet_id}
    grid_range.update(_parse_a1_notation(range))

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "repeatCell": {
                        "range": grid_range,
                        "cell": {"userEnteredFormat": cell_format},
                        "fields": ",".join(fields),
                    }
                }]
            },
        ).execute()
        return {"success": True, "updatedCells": result.get("replies", [{}])[0]}
    except Exception as e:
        return {"error": f"Format failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Merge Cells",
        destructiveHint=True,
    ),
)
def merge_cells(
    spreadsheet_id: str,
    sheet: str,
    range: str,
    merge_type: str = "MERGE_ALL",
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Merge (or unmerge) cells in a range.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range (e.g., 'A1:C3')
        merge_type: 'MERGE_ALL', 'MERGE_COLUMNS', 'MERGE_ROWS', or 'UNMERGE'

    Returns:
        Result of the merge operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    grid_range = {"sheetId": sheet_id}
    grid_range.update(_parse_a1_notation(range))

    if merge_type.upper() == "UNMERGE":
        request = {"unmergeCells": {"range": grid_range}}
    else:
        request = {"mergeCells": {"range": grid_range, "mergeType": merge_type.upper()}}

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id, body={"requests": [request]}
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Merge failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Auto Resize Columns/Rows",
        destructiveHint=True,
    ),
)
def auto_resize(
    spreadsheet_id: str,
    sheet: str,
    dimension: str = "COLUMNS",
    start_index: int = 0,
    end_index: Optional[int] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Auto-resize columns or rows to fit content.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        dimension: 'COLUMNS' or 'ROWS'
        start_index: 0-based start index
        end_index: 0-based end index (exclusive). Omit to resize all.

    Returns:
        Result of the resize operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    dims: Dict[str, Any] = {
        "sheetId": sheet_id,
        "dimension": dimension.upper(),
        "startIndex": start_index,
    }
    if end_index is not None:
        dims["endIndex"] = end_index

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={"requests": [{"autoResizeDimensions": {"dimensions": dims}}]},
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Auto-resize failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Sort Range",
        destructiveHint=True,
    ),
)
def sort_range(
    spreadsheet_id: str,
    sheet: str,
    range: str,
    sort_column: str,
    order: str = "ASCENDING",
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Sort a range of cells by a column.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range to sort (e.g., 'A1:D100')
        sort_column: Column letter to sort by (e.g., 'B')
        order: 'ASCENDING' or 'DESCENDING'

    Returns:
        Result of the sort operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    grid_range = {"sheetId": sheet_id}
    grid_range.update(_parse_a1_notation(range))

    col_index = _letter_to_column_index(sort_column)

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "sortRange": {
                        "range": grid_range,
                        "sortSpecs": [{
                            "dimensionIndex": col_index,
                            "sortOrder": order.upper(),
                        }],
                    }
                }]
            },
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Sort failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Add Basic Filter",
        destructiveHint=True,
    ),
)
def add_filter(
    spreadsheet_id: str,
    sheet: str,
    range: str,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Add a basic filter (auto-filter) to a range.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range (e.g., 'A1:E100')

    Returns:
        Result of the filter operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    grid_range = {"sheetId": sheet_id}
    grid_range.update(_parse_a1_notation(range))

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={"requests": [{"setBasicFilter": {"filter": {"range": grid_range}}}]},
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Add filter failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Add Pivot Table",
        destructiveHint=True,
    ),
)
def add_pivot_table(
    spreadsheet_id: str,
    source_sheet: str,
    source_range: str,
    target_sheet: str,
    target_cell: str = "A1",
    row_columns: Optional[List[int]] = None,
    value_columns: Optional[List[int]] = None,
    aggregation: str = "SUM",
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Create a pivot table from source data.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        source_sheet: Sheet name containing source data
        source_range: A1 range of source data (e.g., 'A1:E100')
        target_sheet: Sheet name to place the pivot table
        target_cell: Cell to anchor the pivot table (default 'A1')
        row_columns: List of 0-based column indices for row grouping (e.g., [0, 1])
        value_columns: List of 0-based column indices to aggregate (e.g., [3, 4])
        aggregation: Aggregation function — 'SUM', 'COUNTA', 'AVERAGE', 'MAX', 'MIN'

    Returns:
        Result of the pivot table creation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    source_sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, source_sheet)
    target_sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, target_sheet)

    if source_sheet_id is None:
        return {"error": f"Source sheet '{source_sheet}' not found"}
    if target_sheet_id is None:
        return {"error": f"Target sheet '{target_sheet}' not found"}

    source_grid = {"sheetId": source_sheet_id}
    source_grid.update(_parse_a1_notation(source_range))

    target_parsed = _parse_a1_notation(target_cell)

    rows = [{"sourceColumnOffset": c, "showTotals": True, "sortOrder": "ASCENDING"} for c in (row_columns or [0])]
    values = [{"sourceColumnOffset": c, "summarizeFunction": aggregation.upper()} for c in (value_columns or [1])]

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "updateCells": {
                        "rows": [{"values": [{"pivotTable": {
                            "source": source_grid,
                            "rows": rows,
                            "values": values,
                        }}]}],
                        "start": {
                            "sheetId": target_sheet_id,
                            "rowIndex": target_parsed.get("startRowIndex", 0),
                            "columnIndex": target_parsed.get("startColumnIndex", 0),
                        },
                        "fields": "pivotTable",
                    }
                }]
            },
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Pivot table failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Add Conditional Formatting",
        destructiveHint=True,
    ),
)
def add_conditional_format(
    spreadsheet_id: str,
    sheet: str,
    range: str,
    rule_type: str = "NUMBER_GREATER",
    values: Optional[List[str]] = None,
    background_color: str = "#FF0000",
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Add a conditional formatting rule to a range.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range (e.g., 'B2:B100')
        rule_type: Condition type — 'NUMBER_GREATER', 'NUMBER_LESS', 'NUMBER_EQ',
                   'TEXT_CONTAINS', 'TEXT_NOT_CONTAINS', 'BLANK', 'NOT_BLANK',
                   'NUMBER_BETWEEN', 'NUMBER_NOT_BETWEEN'
        values: Condition values (e.g., ['100'] for greater than 100, ['10','50'] for between)
        background_color: Hex color to apply when condition is met (e.g., '#FF0000')

    Returns:
        Result of the conditional formatting operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    grid_range = {"sheetId": sheet_id}
    grid_range.update(_parse_a1_notation(range))

    def hex_to_color(hex_str: str) -> Dict[str, float]:
        h = hex_str.lstrip("#")
        return {
            "red": int(h[0:2], 16) / 255.0,
            "green": int(h[2:4], 16) / 255.0,
            "blue": int(h[4:6], 16) / 255.0,
        }

    condition_values = [{"userEnteredValue": v} for v in (values or [])]

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "addConditionalFormatRule": {
                        "rule": {
                            "ranges": [grid_range],
                            "booleanRule": {
                                "condition": {
                                    "type": rule_type.upper(),
                                    "values": condition_values,
                                },
                                "format": {
                                    "backgroundColor": hex_to_color(background_color),
                                },
                            },
                        },
                        "index": 0,
                    }
                }]
            },
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Conditional format failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Duplicate Sheet",
        destructiveHint=True,
    ),
)
def duplicate_sheet(
    spreadsheet_id: str,
    sheet: str,
    new_name: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Duplicate a sheet/tab within the same spreadsheet.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Name of the sheet to duplicate
        new_name: Name for the duplicated sheet (defaults to 'Copy of <sheet>')

    Returns:
        The new sheet's properties
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    try:
        requests = [{"duplicateSheet": {"sourceSheetId": sheet_id, "newSheetName": new_name or f"Copy of {sheet}"}}]
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id, body={"requests": requests}
        ).execute()
        new_props = result.get("replies", [{}])[0].get("duplicateSheet", {}).get("properties", {})
        return {"success": True, "new_sheet": new_props}
    except Exception as e:
        return {"error": f"Duplicate failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Move Sheet (Reorder Tab)",
        destructiveHint=True,
    ),
)
def move_sheet(
    spreadsheet_id: str,
    sheet: str,
    index: int,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Move a sheet/tab to a different position (reorder tabs).

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Name of the sheet to move
        index: New 0-based position index

    Returns:
        Result of the move operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "updateSheetProperties": {
                        "properties": {"sheetId": sheet_id, "index": index},
                        "fields": "index",
                    }
                }]
            },
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Move failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Protect Range",
        destructiveHint=True,
    ),
)
def protect_range(
    spreadsheet_id: str,
    sheet: str,
    range: Optional[str] = None,
    description: str = "",
    warning_only: bool = False,
    editors: Optional[List[str]] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Protect a range or entire sheet from editing.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range to protect. Omit to protect the entire sheet.
        description: Description of the protection
        warning_only: If True, shows a warning but allows editing
        editors: List of email addresses allowed to edit

    Returns:
        Result of the protection operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    protected: Dict[str, Any] = {
        "description": description,
        "warningOnly": warning_only,
    }

    if range:
        grid_range = {"sheetId": sheet_id}
        grid_range.update(_parse_a1_notation(range))
        protected["range"] = grid_range
    else:
        protected["range"] = {"sheetId": sheet_id}

    if editors:
        protected["editors"] = {"users": editors}

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={"requests": [{"addProtectedRange": {"protectedRange": protected}}]},
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Protect failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Add Comment",
        destructiveHint=True,
    ),
)
def add_comment(
    spreadsheet_id: str,
    content: str,
    anchor_cell: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Add a comment to a spreadsheet. Optionally anchor it to a specific cell.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        content: The comment text
        anchor_cell: Cell reference in 'Sheet1!A1' format. If omitted, adds a file-level comment.

    Returns:
        The created comment
    """
    drive_service = ctx.request_context.lifespan_context.drive_service

    body: Dict[str, Any] = {"content": content}
    if anchor_cell:
        body["anchor"] = f"cell={anchor_cell}"

    try:
        result = drive_service.comments().create(
            fileId=spreadsheet_id,
            body=body,
            fields="id,content,author,createdTime,anchor",
        ).execute()
        return {"success": True, "comment": result}
    except Exception as e:
        return {"error": f"Add comment failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Get Revision History",
        readOnlyHint=True,
    ),
)
def get_revision_history(
    spreadsheet_id: str,
    max_results: int = 20,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Get the revision (edit) history of a spreadsheet.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        max_results: Maximum number of revisions to return (default 20)

    Returns:
        List of revisions with author, date, and modification details
    """
    drive_service = ctx.request_context.lifespan_context.drive_service

    try:
        revisions = drive_service.revisions().list(
            fileId=spreadsheet_id,
            pageSize=max_results,
            fields="revisions(id,modifiedTime,lastModifyingUser(displayName,emailAddress))",
        ).execute()
        return {
            "success": True,
            "revisions": revisions.get("revisions", []),
            "count": len(revisions.get("revisions", [])),
        }
    except Exception as e:
        return {"error": f"Get revisions failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Delete Rows or Columns",
        destructiveHint=True,
    ),
)
def delete_rows_columns(
    spreadsheet_id: str,
    sheet: str,
    dimension: str,
    start_index: int,
    end_index: int,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Delete rows or columns from a sheet.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        dimension: 'ROWS' or 'COLUMNS'
        start_index: 0-based start index (inclusive)
        end_index: 0-based end index (exclusive)

    Returns:
        Result of the delete operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "deleteDimension": {
                        "range": {
                            "sheetId": sheet_id,
                            "dimension": dimension.upper(),
                            "startIndex": start_index,
                            "endIndex": end_index,
                        }
                    }
                }]
            },
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Delete failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Clear Range",
        destructiveHint=True,
    ),
)
def clear_range(
    spreadsheet_id: str,
    sheet: str,
    range: str,
    clear_format: bool = False,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Clear data from a range of cells. Optionally clear formatting too.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        range: A1 notation range (e.g., 'A1:C10')
        clear_format: If True, also clears formatting. If False (default), only clears values.

    Returns:
        Result of the clear operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    full_range = f"{sheet}!{range}"

    try:
        if clear_format:
            # Use batchUpdate to clear both values and formatting
            sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
            if sheet_id is None:
                return {"error": f"Sheet '{sheet}' not found"}

            grid_range = {"sheetId": sheet_id}
            grid_range.update(_parse_a1_notation(range))

            result = sheets_service.spreadsheets().batchUpdate(
                spreadsheetId=spreadsheet_id,
                body={
                    "requests": [{
                        "updateCells": {
                            "range": grid_range,
                            "fields": "userEnteredValue,userEnteredFormat",
                        }
                    }]
                },
            ).execute()
        else:
            # Only clear values, keep formatting
            result = sheets_service.spreadsheets().values().clear(
                spreadsheetId=spreadsheet_id,
                range=full_range,
            ).execute()

        return {"success": True, "clearedRange": full_range, "result": result}
    except Exception as e:
        return {"error": f"Clear failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Get Comments",
        readOnlyHint=True,
    ),
)
def get_comments(
    spreadsheet_id: str,
    include_resolved: bool = False,
    max_results: int = 50,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Get all comments on a spreadsheet.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        include_resolved: Whether to include resolved comments (default False)
        max_results: Maximum number of comments to return (default 50)

    Returns:
        List of comments with author, content, anchor cell, and resolved status
    """
    drive_service = ctx.request_context.lifespan_context.drive_service

    try:
        comments = drive_service.comments().list(
            fileId=spreadsheet_id,
            pageSize=max_results,
            includeDeleted=False,
            fields="comments(id,content,author(displayName,emailAddress),createdTime,modifiedTime,resolved,anchor,replies(id,content,author(displayName,emailAddress),createdTime))",
        ).execute()

        result_comments = []
        for c in comments.get("comments", []):
            if not include_resolved and c.get("resolved"):
                continue
            result_comments.append({
                "id": c.get("id"),
                "content": c.get("content"),
                "author": c.get("author", {}).get("displayName"),
                "author_email": c.get("author", {}).get("emailAddress"),
                "created": c.get("createdTime"),
                "resolved": c.get("resolved", False),
                "anchor": c.get("anchor"),
                "replies": [
                    {
                        "content": r.get("content"),
                        "author": r.get("author", {}).get("displayName"),
                        "created": r.get("createdTime"),
                    }
                    for r in c.get("replies", [])
                ],
            })

        return {
            "success": True,
            "comments": result_comments,
            "count": len(result_comments),
        }
    except Exception as e:
        return {"error": f"Get comments failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Resolve Comment",
        destructiveHint=True,
    ),
)
def resolve_comment(
    spreadsheet_id: str,
    comment_id: str,
    reply_content: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Resolve a comment, optionally adding a reply before resolving.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        comment_id: The comment ID (from get_comments)
        reply_content: Optional reply to add before resolving

    Returns:
        Result of the resolve operation
    """
    drive_service = ctx.request_context.lifespan_context.drive_service

    try:
        # Add reply if provided
        if reply_content:
            drive_service.replies().create(
                fileId=spreadsheet_id,
                commentId=comment_id,
                body={"content": reply_content},
                fields="id",
            ).execute()

        # Resolve the comment by updating it
        result = drive_service.comments().update(
            fileId=spreadsheet_id,
            commentId=comment_id,
            body={"resolved": True},
            fields="id,resolved",
        ).execute()

        return {"success": True, "comment_id": comment_id, "resolved": True}
    except Exception as e:
        return {"error": f"Resolve comment failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Set Tab Color",
        destructiveHint=True,
    ),
)
def set_tab_color(
    spreadsheet_id: str,
    sheet: str,
    color: str,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Change the color of a sheet tab.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        color: Hex color (e.g., '#FF0000' for red, '#00FF00' for green). Use '#FFFFFF' or 'none' to remove color.

    Returns:
        Result of the color change operation
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service
    sheet_id = _get_sheet_id(sheets_service, spreadsheet_id, sheet)
    if sheet_id is None:
        return {"error": f"Sheet '{sheet}' not found"}

    if color.lower() == "none":
        tab_color = {"red": 1.0, "green": 1.0, "blue": 1.0}
    else:
        h = color.lstrip("#")
        tab_color = {
            "red": int(h[0:2], 16) / 255.0,
            "green": int(h[2:4], 16) / 255.0,
            "blue": int(h[4:6], 16) / 255.0,
        }

    try:
        result = sheets_service.spreadsheets().batchUpdate(
            spreadsheetId=spreadsheet_id,
            body={
                "requests": [{
                    "updateSheetProperties": {
                        "properties": {
                            "sheetId": sheet_id,
                            "tabColorStyle": {"rgbColor": tab_color},
                        },
                        "fields": "tabColorStyle",
                    }
                }]
            },
        ).execute()
        return {"success": True, "result": result}
    except Exception as e:
        return {"error": f"Set tab color failed: {str(e)}"}


@tool(
    annotations=ToolAnnotations(
        title="Append Rows",
        destructiveHint=True,
    ),
)
def append_rows(
    spreadsheet_id: str,
    sheet: str,
    values: List[List[Any]],
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Append rows of data after the last row with content in a sheet.
    Useful for adding log entries or new records without knowing the exact row number.

    Args:
        spreadsheet_id: The ID of the spreadsheet
        sheet: Sheet/tab name
        values: List of rows, each row is a list of cell values.
                Example: [["Alice", 30, "Seoul"], ["Bob", 25, "Busan"]]

    Returns:
        The range where data was appended and number of rows added
    """
    sheets_service = ctx.request_context.lifespan_context.sheets_service

    try:
        result = sheets_service.spreadsheets().values().append(
            spreadsheetId=spreadsheet_id,
            range=f"{sheet}!A1",
            valueInputOption="USER_ENTERED",
            insertDataOption="INSERT_ROWS",
            body={"values": values},
        ).execute()

        updates = result.get("updates", {})
        return {
            "success": True,
            "updatedRange": updates.get("updatedRange"),
            "updatedRows": updates.get("updatedRows"),
            "updatedCells": updates.get("updatedCells"),
        }
    except Exception as e:
        return {"error": f"Append rows failed: {str(e)}"}


def main():
    # Log tool filtering configuration if enabled
    if ENABLED_TOOLS is not None:
        print(f"Tool filtering enabled. Active tools: {', '.join(sorted(ENABLED_TOOLS))}")
    else:
        print("Tool filtering disabled. All tools are enabled.")
    
    # Run the server
    transport = "stdio"
    for i, arg in enumerate(sys.argv):
        if arg == "--transport" and i + 1 < len(sys.argv):
            transport = sys.argv[i + 1]
            break

    mcp.run(transport=transport)
