#!/usr/bin/env python3
"""
JARVIS Google Sheets — Dynamic Discovery Engine
No hardcoded sheet IDs. JARVIS discovers sheets per Chrome profile account,
asks the user which one, then remembers forever.
"""

import json, os
from pathlib import Path
from typing import Optional

import memory_store as mem
from chrome_manager import discover_chrome_profiles

CREDS_DIR = Path(__file__).parent.parent / ".jarvis_tokens"
CREDS_DIR.mkdir(exist_ok=True)

SCOPES = [
    "https://www.googleapis.com/auth/spreadsheets.readonly",
    "https://www.googleapis.com/auth/drive.readonly",
]


# ═══════════════════════════════════════════════════════════════
#  OAUTH TOKENS PER PROFILE
# ═══════════════════════════════════════════════════════════════

def _token_path(email: str) -> Path:
    safe = email.replace("@", "_at_").replace(".", "_")
    return CREDS_DIR / f"token_{safe}.json"


def _has_token(email: str) -> bool:
    return _token_path(email).exists()


def _get_credentials(email: str):
    """Return OAuth credentials for the given email, refreshing if needed."""
    from google.oauth2.credentials import Credentials
    from google.auth.transport.requests import Request
    from google_auth_oauthlib.flow import InstalledAppFlow

    token_path = _token_path(email)

    creds = None
    if token_path.exists():
        creds = Credentials.from_authorized_user_file(str(token_path), SCOPES)

    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            # Need to do initial OAuth — requires client_secrets.json
            client_secrets = CREDS_DIR / "client_secrets.json"
            if not client_secrets.exists():
                return None, "google_oauth_needed"
            flow  = InstalledAppFlow.from_client_secrets_file(str(client_secrets), SCOPES)
            creds = flow.run_local_server(port=0)
        token_path.write_text(creds.to_json())

    return creds, "ok"


# ═══════════════════════════════════════════════════════════════
#  SHEET DISCOVERY
# ═══════════════════════════════════════════════════════════════

def list_sheets_for_profile(email: str) -> dict:
    """
    List all Google Sheets in the Drive of the given email.
    Requires prior OAuth authorization.
    """
    creds, status = _get_credentials(email)
    if not creds:
        return {"status": status, "sheets": [], "email": email}

    from googleapiclient.discovery import build

    try:
        service = build("drive", "v3", credentials=creds)
        results = service.files().list(
            q="mimeType='application/vnd.google-apps.spreadsheet' and trashed=false",
            pageSize=50,
            fields="files(id, name, modifiedTime, owners)",
        ).execute()
        files   = results.get("files", [])
        sheets  = [{"id": f["id"], "name": f["name"],
                    "modified": f.get("modifiedTime", "")} for f in files]
        # Cache in memory
        for sh in sheets:
            mem.save_sheet_mapping(email, sh["id"], sh["name"])
        return {"status": "ok", "email": email, "sheets": sheets}
    except Exception as e:
        return {"status": "error", "error": str(e), "sheets": []}


def read_sheet(sheet_id: str, email: str, range_: str = "A1:Z1000") -> list[dict]:
    """
    Read a Google Sheet and return rows as list of dicts using first row as headers.
    Remembers sheet access in memory_store.
    """
    creds, status = _get_credentials(email)
    if not creds:
        return [{"error": f"Not authorized: {status}"}]

    from googleapiclient.discovery import build

    try:
        service = build("sheets", "v4", credentials=creds)
        result  = service.spreadsheets().values().get(
            spreadsheetId=sheet_id, range=range_
        ).execute()
        values  = result.get("values", [])
        if not values:
            return []
        headers = values[0]
        rows    = []
        for row in values[1:]:
            padded = row + [""] * (len(headers) - len(row))
            rows.append(dict(zip(headers, padded)))
        mem.touch_sheet(sheet_id)
        return rows
    except Exception as e:
        return [{"error": str(e)}]


def search_in_sheet(sheet_id: str, email: str, query: str,
                    columns: list = None) -> list[dict]:
    """Search all rows in a sheet for a query string."""
    rows    = read_sheet(sheet_id, email)
    query_l = query.lower()
    matches = []
    for row in rows:
        vals = [str(v).lower() for v in row.values()]
        if any(query_l in v for v in vals):
            matches.append(row)
    # If we found something, learn which sheet to use for this type of query
    if matches:
        mem.learn(f"query_source_{query_l[:30]}", {"sheet_id": sheet_id, "email": email},
                  category="sheet_mapping", source="sheets_search")
    return matches


# ═══════════════════════════════════════════════════════════════
#  SMART LOOKUP (the discovery-driven flow)
# ═══════════════════════════════════════════════════════════════

def smart_lookup(query: str) -> dict:
    """
    Smart contact/data lookup:
    1. Check memory: do we already know which sheet has this?
    2. If yes → search that sheet directly
    3. If no → return instructions for JARVIS to ask the user
    """
    # Step 1: Check memory for known mapping
    known = mem.recall(f"query_source_{query.lower()[:30]}")
    if known:
        rows = search_in_sheet(known["sheet_id"], known["email"], query)
        if rows:
            return {"status": "found", "source": "memory", "results": rows}

    # Step 2: Check known sheet mappings for fuzzy match
    mappings = mem.get_sheet_mappings()
    for mapping in mappings:
        rows = search_in_sheet(mapping["sheet_id"], mapping["chrome_profile"], query)
        if rows:
            return {"status": "found", "source": mapping["sheet_name"], "results": rows}

    # Step 3: Not found anywhere — ask user
    profiles = discover_chrome_profiles()
    profile_list = [f"{p.get('display_name', 'Unknown')} ({p.get('email', 'no email')})"
                    for p in profiles]
    return {
        "status":   "need_profile",
        "message":  "I don't know which sheet has this information yet.",
        "profiles": profile_list,
        "query":    query,
    }


def authorize_profile_sheets(email: str) -> dict:
    """
    Initiate OAuth for a Google profile.
    Returns status and list of sheets found.
    """
    if not (CREDS_DIR / "client_secrets.json").exists():
        return {
            "status": "needs_setup",
            "message": (
                "To access Google Sheets, I need your Google API client credentials. "
                "Please download 'client_secrets.json' from Google Cloud Console "
                f"and place it at: {CREDS_DIR / 'client_secrets.json'}"
            )
        }
    result = list_sheets_for_profile(email)
    if result["status"] == "ok":
        mem.learn(f"sheets_authorized_{email}", True, category="auth", source="sheets")
        mem.learn(f"sheets_list_{email}", result["sheets"], category="sheets", source="sheets")
    return result