"""
OAuth2 authentication with named profile support.

Designed for zero-friction setup:
  1. Claude calls setup_google_workspace() with OAuth credentials + profiles
  2. Claude calls authenticate_account() per profile → browser opens, user consents, done
  3. All subsequent tool calls auto-resolve credentials, auto-refresh tokens

Config stored at ~/.config/google-workspace-mcp/config.json
Tokens stored at ~/.config/google-workspace-mcp/tokens/<profile>.json
"""

from __future__ import annotations

import asyncio
import json
import os
import threading
import webbrowser
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Optional

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
CONFIG_DIR = Path(
    os.environ.get(
        "GW_MCP_CONFIG_DIR",
        Path.home() / ".config" / "google-workspace-mcp",
    )
)
CONFIG_FILE = CONFIG_DIR / "config.json"
TOKENS_DIR = CONFIG_DIR / "tokens"

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

OAUTH_PORT = int(os.environ.get("GW_MCP_OAUTH_PORT", "3456"))


# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------
@dataclass
class AccountProfile:
    name: str
    email: str
    description: str = ""


@dataclass
class Config:
    client_id: str
    client_secret: str
    profiles: list[AccountProfile]


# ---------------------------------------------------------------------------
# Config helpers
# ---------------------------------------------------------------------------
def _ensure_dirs() -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    TOKENS_DIR.mkdir(parents=True, exist_ok=True)


def load_config() -> Config:
    if not CONFIG_FILE.exists():
        raise FileNotFoundError(
            "No config found. Use the setup_google_workspace tool to configure."
        )
    raw = json.loads(CONFIG_FILE.read_text())
    profiles = [
        AccountProfile(
            name=p["name"],
            email=p["email"],
            description=p.get("description", ""),
        )
        for p in raw["profiles"]
    ]
    return Config(
        client_id=raw["clientId"],
        client_secret=raw["clientSecret"],
        profiles=profiles,
    )


def save_config(config: Config) -> None:
    """Save config to disk."""
    _ensure_dirs()
    data = {
        "clientId": config.client_id,
        "clientSecret": config.client_secret,
        "profiles": [asdict(p) for p in config.profiles],
    }
    CONFIG_FILE.write_text(json.dumps(data, indent=2))


def config_exists() -> bool:
    return CONFIG_FILE.exists()


def _token_path(profile_name: str) -> Path:
    return TOKENS_DIR / f"{profile_name}.json"


# ---------------------------------------------------------------------------
# Credential cache
# ---------------------------------------------------------------------------
_creds_cache: dict[str, Credentials] = {}


def get_credentials(profile_name: str) -> Credentials:
    """Return valid OAuth2 credentials for the given profile, refreshing if needed."""
    # Check in-memory cache first
    if profile_name in _creds_cache:
        creds = _creds_cache[profile_name]
        if creds.valid:
            return creds
        if creds.expired and creds.refresh_token:
            creds.refresh(Request())
            _save_token(profile_name, creds)
            return creds

    # Check saved token on disk
    tp = _token_path(profile_name)
    if tp.exists():
        creds = Credentials.from_authorized_user_file(str(tp), SCOPES)
        if creds.expired and creds.refresh_token:
            creds.refresh(Request())
            _save_token(profile_name, creds)
        _creds_cache[profile_name] = creds
        return creds

    raise RuntimeError(
        f'Profile "{profile_name}" is not authenticated. '
        f'Use the authenticate_account tool with profile="{profile_name}" to sign in.'
    )


def _save_token(profile_name: str, creds: Credentials) -> None:
    _ensure_dirs()
    _token_path(profile_name).write_text(creds.to_json())


# ---------------------------------------------------------------------------
# Profile helpers
# ---------------------------------------------------------------------------
def list_profiles() -> list[dict]:
    if not config_exists():
        return []
    config = load_config()
    result = []
    for p in config.profiles:
        authed = _token_path(p.name).exists()
        result.append(
            {
                "name": p.name,
                "email": p.email,
                "description": p.description,
                "authenticated": authed,
            }
        )
    return result


def get_default_profile() -> str:
    config = load_config()
    if not config.profiles:
        raise RuntimeError("No profiles configured.")
    return config.profiles[0].name


def resolve_profile(profile: Optional[str]) -> str:
    """Return the given profile name or fall back to the default."""
    return profile or get_default_profile()


# ---------------------------------------------------------------------------
# Setup helpers (called by MCP tools)
# ---------------------------------------------------------------------------
def create_config(
    client_id: str,
    client_secret: str,
    profiles: list[dict[str, str]],
) -> dict[str, Any]:
    """Create or overwrite the config file. Returns status info."""
    account_profiles = [
        AccountProfile(
            name=p["name"],
            email=p["email"],
            description=p.get("description", ""),
        )
        for p in profiles
    ]
    config = Config(
        client_id=client_id,
        client_secret=client_secret,
        profiles=account_profiles,
    )
    save_config(config)
    return {
        "status": "configured",
        "config_path": str(CONFIG_FILE),
        "profiles": [p["name"] for p in profiles],
        "next_step": "Use authenticate_account for each profile to complete sign-in.",
    }


def add_profile(name: str, email: str, description: str = "") -> dict[str, Any]:
    """Add a new profile to existing config."""
    config = load_config()
    # Check for duplicates
    if any(p.name == name for p in config.profiles):
        return {"status": "error", "message": f'Profile "{name}" already exists.'}
    config.profiles.append(AccountProfile(name=name, email=email, description=description))
    save_config(config)
    return {
        "status": "added",
        "profile": name,
        "email": email,
        "next_step": f'Use authenticate_account with profile="{name}" to sign in.',
    }


# ---------------------------------------------------------------------------
# OAuth flow — runs in a background thread so MCP server isn't blocked
# ---------------------------------------------------------------------------
def _run_oauth_flow(profile_name: str) -> dict[str, Any]:
    """
    Run the full OAuth2 flow:
      1. Start a local HTTP server on OAUTH_PORT
      2. Open the user's browser to Google's consent screen
      3. Catch the callback with the auth code
      4. Exchange for tokens and save them

    This is a BLOCKING call — use authenticate_profile_async for MCP tools.
    """
    config = load_config()
    profile = next((p for p in config.profiles if p.name == profile_name), None)
    if not profile:
        available = ", ".join(p.name for p in config.profiles)
        return {
            "status": "error",
            "message": f'Profile "{profile_name}" not found. Available: {available}',
        }

    client_config = {
        "installed": {
            "client_id": config.client_id,
            "client_secret": config.client_secret,
            "redirect_uris": [f"http://localhost:{OAUTH_PORT}/"],
            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
            "token_uri": "https://oauth2.googleapis.com/token",
        }
    }

    flow = InstalledAppFlow.from_client_config(client_config, SCOPES)

    # run_local_server opens browser automatically and waits for callback
    creds = flow.run_local_server(
        port=OAUTH_PORT,
        open_browser=True,
        prompt="consent",
        login_hint=profile.email,
        success_message=(
            f'Authenticated "{profile_name}" ({profile.email}). '
            "You can close this tab and go back to Claude."
        ),
    )

    _save_token(profile_name, creds)
    _creds_cache[profile_name] = creds

    return {
        "status": "authenticated",
        "profile": profile_name,
        "email": profile.email,
        "message": f"Successfully signed in as {profile.email}",
    }


async def authenticate_profile_async(profile_name: str) -> dict[str, Any]:
    """
    Async wrapper that runs the blocking OAuth flow in a thread executor.
    This is what MCP tools should call.
    """
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, _run_oauth_flow, profile_name)


# ---------------------------------------------------------------------------
# Status check
# ---------------------------------------------------------------------------
def get_setup_status() -> dict[str, Any]:
    """Get the current setup status — useful for Claude to decide what to do."""
    if not config_exists():
        return {
            "configured": False,
            "message": (
                "Not configured yet. Use setup_google_workspace tool with your "
                "Google Cloud OAuth Client ID and Secret to get started."
            ),
            "setup_instructions": {
                "step_1": "Go to https://console.cloud.google.com/apis/credentials",
                "step_2": "Create an OAuth 2.0 Client ID (type: Desktop app)",
                "step_3": "Enable Gmail, Drive, Sheets & Calendar APIs",
                "step_4": "Copy the Client ID and Client Secret",
                "step_5": "Call setup_google_workspace with those credentials and your profiles",
            },
        }

    profiles = list_profiles()
    unauthenticated = [p["name"] for p in profiles if not p["authenticated"]]

    return {
        "configured": True,
        "config_path": str(CONFIG_FILE),
        "profiles": profiles,
        "all_authenticated": len(unauthenticated) == 0,
        "needs_auth": unauthenticated,
    }
