from flask import Flask, redirect, request, session
import asyncio
from mcp_rank import MCPRankClient

app = Flask(__name__)
app.secret_key = "your-secret-key"
client = MCPRankClient(api_key="sk_mcp_rank_...")

@app.route("/connect/google")
def connect_google():
    # Generate auth URL
    auth_url = asyncio.run(client.get_auth_url(
        provider="google",
        redirect_uri="http://localhost:5000/callback",
        state=session.get("csrf_token")
    ))
    return redirect(auth_url)

@app.route("/callback")
def callback():
    # Verify state (CSRF protection)
    if request.args.get("state") != session.get("csrf_token"):
        return "Invalid state", 400
    
    # Check for errors
    if "error" in request.args:
        return f"Auth failed: {request.args['error']}", 400
    
    # Store tokens
    session["access_token"] = request.args["access_token"]
    session["refresh_token"] = request.args["refresh_token"]
    session["user_id"] = request.args["user_id"]
    
    return redirect("/dashboard")

@app.route("/dashboard")
def dashboard():
    if "access_token" not in session:
        return redirect("/connect/google")
    
    # Use the token
    client.set_user_token(session["access_token"], session["refresh_token"])
    
    try:
        profile = asyncio.run(client.proxy_get("google", "gmail/v1/users/me/profile"))
        return f"Welcome {profile['emailAddress']}!"
    except Exception:
        # Token might be expired, refresh it
        return redirect("/connect/google")

from fastapi import FastAPI, Request
from fastapi.responses import RedirectResponse
from mcp_rank import MCPRankClient

app = FastAPI()
client = MCPRankClient(api_key="sk_mcp_rank_...")

@app.get("/connect/google")
async def connect_google(request: Request):
    auth_url = await client.get_auth_url(
        provider="google",
        redirect_uri="http://localhost:8000/callback"
    )
    return RedirectResponse(auth_url)

@app.get("/callback")
async def callback(
    access_token: str,
    refresh_token: str,
    user_id: str,
    state: str = None
):
    # Store tokens in your database
    await save_user_tokens(user_id, access_token, refresh_token)
    return RedirectResponse("/dashboard")

@app.get("/dashboard")
async def dashboard(request: Request):
    user = await get_current_user(request)
    tokens = await get_user_tokens(user.id)
    
    client.set_user_token(tokens.access_token, tokens.refresh_token)
    profile = await client.proxy_get("google", "gmail/v1/users/me/profile")
    
    return {"email": profile["emailAddress"]}

import asyncio
import webbrowser
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from mcp_rank import MCPRankClient

tokens = {}

class CallbackHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        query = parse_qs(urlparse(self.path).query)
        tokens["access_token"] = query.get("access_token", [None])[0]
        tokens["refresh_token"] = query.get("refresh_token", [None])[0]
        
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b"Success! You can close this window.")

async def authenticate():
    client = MCPRankClient(api_key="sk_mcp_rank_...")
    
    # Start local server
    server = HTTPServer(("localhost", 8080), CallbackHandler)
    
    # Get auth URL
    auth_url = await client.get_auth_url(
        provider="google",
        redirect_uri="http://localhost:8080/callback"
    )
    
    # Open browser
    print(f"Opening browser for authentication...")
    webbrowser.open(auth_url)
    
    # Wait for callback
    server.handle_request()
    
    # Use tokens
    client.set_user_token(tokens["access_token"], tokens["refresh_token"])
    profile = await client.proxy_get("google", "gmail/v1/users/me/profile")
    print(f"Authenticated as: {profile['emailAddress']}")
    
    return client

asyncio.run(authenticate())

# models.py
class UserToken(Base):
    __tablename__ = "user_tokens"
    
    user_id = Column(String, primary_key=True)
    mcp_user_id = Column(String)
    access_token = Column(String)  # Consider encrypting
    refresh_token = Column(String)  # Consider encrypting
    expires_at = Column(DateTime)

# Encrypt sensitive fields
from cryptography.fernet import Fernet

cipher = Fernet(os.environ["ENCRYPTION_KEY"])

def save_tokens(user_id, access_token, refresh_token):
    token = UserToken(
        user_id=user_id,
        access_token=cipher.encrypt(access_token.encode()),
        refresh_token=cipher.encrypt(refresh_token.encode())
    )
    db.add(token)
    db.commit()

async def get_client_for_user(user_id: str) -> MCPRankClient:
    tokens = await get_user_tokens(user_id)
    
    client = MCPRankClient(api_key="sk_mcp_rank_...")
    client.set_user_token(tokens.access_token, tokens.refresh_token)
    
    # Check if token needs refresh
    if tokens.expires_at < datetime.utcnow():
        new_tokens = await client.refresh_user_token()
        await save_user_tokens(
            user_id,
            new_tokens["access_token"],
            new_tokens["refresh_token"]
        )
    
    return client