#!/usr/bin/env python3
"""
Archie Backend with WorkOS AuthKit
Proper implementation using WorkOS User Management API
"""
from fastapi import FastAPI, Depends, HTTPException, status, Request, Cookie
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel
from workos import WorkOSClient
import os
from typing import Optional
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration
WORKOS_API_KEY = os.getenv("WORKOS_API_KEY")
WORKOS_CLIENT_ID = os.getenv("WORKOS_CLIENT_ID")
REDIRECT_URI = os.getenv("REDIRECT_URI", "https://auth-dev2.arch.ie/callback")

# Initialize WorkOS Client
workos = WorkOSClient(
    api_key=WORKOS_API_KEY,
    client_id=WORKOS_CLIENT_ID
)

app = FastAPI(title="Archie API with AuthKit", version="2.0.0")

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Models
class QueryRequest(BaseModel):
    query: str

class QueryResponse(BaseModel):
    response: str
    query_id: str

# Auth dependency using AuthKit session
async def get_current_user(
    wos_session: Optional[str] = Cookie(None)
):
    """
    Validate WorkOS AuthKit session cookie.

    AuthKit uses secure HTTP-only cookies for session management.
    """
    if not wos_session:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="No session cookie found"
        )

    try:
        # Authenticate with WorkOS session
        auth_response = workos.user_management.authenticate_with_session_cookie(
            session_cookie=wos_session
        )

        user = auth_response.user
        logger.info(f"Authenticated user: {user.id} ({user.email})")

        return {
            "id": user.id,
            "email": user.email,
            "first_name": user.first_name,
            "last_name": user.last_name,
        }

    except Exception as e:
        logger.error(f"Session validation failed: {e}")
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid session"
        )

# Health endpoint
@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "version": "2.0.0",
        "auth": "WorkOS AuthKit",
        "redirect_uri": REDIRECT_URI
    }

# AuthKit Authorization (starts the flow)
@app.get("/api/v1/auth/workos/authorize")
async def authkit_authorize(redirect_uri: Optional[str] = None):
    """
    Initiate WorkOS AuthKit authentication flow.

    AuthKit will:
    1. Show a hosted login page
    2. Handle Google OAuth
    3. Create a session
    4. Redirect back with authorization code
    """
    uri = redirect_uri or REDIRECT_URI
    logger.info(f"Starting AuthKit flow with redirect_uri: {uri}")

    try:
        # Get authorization URL from AuthKit
        authorization_url = workos.user_management.get_authorization_url(
            provider="authkit",  # Use AuthKit's hosted auth
            redirect_uri=uri,
            client_id=WORKOS_CLIENT_ID
        )

        return RedirectResponse(url=authorization_url)

    except Exception as e:
        logger.error(f"Failed to create authorization URL: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# AuthKit Callback Handler
@app.get("/callback")
async def authkit_callback(code: str):
    """
    Handle AuthKit OAuth callback.

    AuthKit simplified flow:
    1. Receive authorization code
    2. Exchange for access token + refresh token
    3. Get user profile
    4. Create session cookie (HTTP-only, secure)
    """
    logger.info(f"Handling AuthKit callback with code: {code[:20]}...")

    try:
        # Exchange code for tokens using AuthKit
        auth_response = workos.user_management.authenticate_with_code(
            code=code,
            session_options={
                "cookie_password": os.getenv("COOKIE_PASSWORD", "replace-with-32-char-secret-key!!"),
                "seal_session": True,  # HTTP-only, secure cookie
            }
        )

        user = auth_response.user
        sealed_session = auth_response.sealed_session

        logger.info(f"AuthKit authentication successful for user: {user.email}")

        # Create response with session cookie
        response = JSONResponse({
            "success": True,
            "user": {
                "id": user.id,
                "email": user.email,
                "first_name": user.first_name,
                "last_name": user.last_name,
            },
            "message": "Authentication successful"
        })

        # Set secure session cookie
        response.set_cookie(
            key="wos_session",
            value=sealed_session,
            httponly=True,  # Prevents JavaScript access
            secure=True,  # HTTPS only
            samesite="lax",  # CSRF protection
            max_age=604800  # 7 days
        )

        return response

    except Exception as e:
        logger.error(f"AuthKit callback error: {e}")
        raise HTTPException(
            status_code=400,
            detail=f"Authentication failed: {str(e)}"
        )

# Alternative callback path
@app.get("/api/v1/auth/workos/callback")
async def authkit_callback_full(code: str):
    """Handle AuthKit callback - full path"""
    return await authkit_callback(code)

# Query endpoint (main CLI endpoint)
@app.post("/api/v1/query")
async def create_query(
    request: QueryRequest,
    user=Depends(get_current_user)
):
    """
    Process user query with AuthKit session authentication.

    The session cookie is automatically validated by the dependency.
    """
    logger.info(f"Query from {user['email']}: {request.query}")

    return QueryResponse(
        response=f"Received query from {user['first_name']}: '{request.query}'. AuthKit working!",
        query_id="authkit-123"
    )

# Logout endpoint
@app.post("/api/v1/auth/logout")
async def logout():
    """Clear AuthKit session"""
    response = JSONResponse({"success": True, "message": "Logged out"})
    response.delete_cookie(key="wos_session")
    return response

# WebSocket endpoint stub
@app.get("/api/v1/ws")
async def websocket_info():
    return {"message": "WebSocket endpoint - not implemented yet"}

if __name__ == "__main__":
    import uvicorn
    logger.info(f"Starting backend with WorkOS AuthKit")
    logger.info(f"REDIRECT_URI: {REDIRECT_URI}")
    uvicorn.run(app, host="0.0.0.0", port=7000)