#!/usr/bin/env python3
"""
Minimal Archie Backend with WorkOS OAuth
Supports CLI authentication and browser-based OAuth
"""
from fastapi import FastAPI, Depends, HTTPException, status, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel
import jwt
import httpx
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", "sk_test_a2V5XzAxSkFaWktHRUE2WVRFMTUxUzE3NFg4NVRNLGFJUDN1eHY0TFhiWVoyVm54a2o5SEpEOGw")
WORKOS_CLIENT_ID = os.getenv("WORKOS_CLIENT_ID", "client_01JAZZKGWFWTWZMDCZGE24VFC1")
# FIXED: Use ngrok domain for redirect_uri
REDIRECT_URI = os.getenv("REDIRECT_URI", "https://auth-dev.arch.ie/callback")

app = FastAPI(title="Archie API", version="1.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
async def get_current_user(request: Request):
    """Validate WorkOS JWT token"""
    auth_header = request.headers.get("Authorization")
    if not auth_header or not auth_header.startswith("Bearer "):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication token"
        )

    token = auth_header.split(" ")[1]

    try:
        # Decode without verification for now (add JWKS validation in production)
        payload = jwt.decode(token, options={"verify_signature": False})
        logger.info(f"Authenticated user: {payload.get('sub')}")
        return payload
    except Exception as e:
        logger.error(f"Token validation failed: {e}")
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication token"
        )

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

# WorkOS OAuth Authorization
@app.get("/api/v1/auth/workos/authorize")
async def workos_authorize(redirect_uri: Optional[str] = None):
    """Initiate WorkOS OAuth flow"""
    uri = redirect_uri or REDIRECT_URI
    logger.info(f"Authorizing with redirect_uri: {uri}")
    auth_url = (
        f"https://api.workos.com/sso/authorize"
        f"?client_id={WORKOS_CLIENT_ID}"
        f"&redirect_uri={uri}"
        f"&response_type=code"
        f"&provider=GoogleOAuth"
    )
    return RedirectResponse(url=auth_url)

# WorkOS OAuth Callback Handler (shared implementation)
async def handle_workos_callback(code: str):
    """Handle WorkOS OAuth callback - shared implementation"""
    logger.info(f"Handling callback with code: {code[:20]}...")
    logger.info(f"Using redirect_uri: {REDIRECT_URI}")

    try:
        async with httpx.AsyncClient() as client:
            # Exchange code for token
            # CRITICAL: client_secret must be in form data, not HTTP Basic Auth
            response = await client.post(
                "https://api.workos.com/sso/token",
                data={
                    "client_id": WORKOS_CLIENT_ID,
                    "client_secret": WORKOS_API_KEY,  # ← FIXED: Send as form parameter
                    "code": code,
                    "grant_type": "authorization_code",
                    "redirect_uri": REDIRECT_URI,
                },
            )

            logger.info(f"WorkOS token response status: {response.status_code}")

            if response.status_code != 200:
                error_detail = response.text
                logger.error(f"Token exchange failed: {error_detail}")
                raise HTTPException(
                    status_code=400,
                    detail=f"Token exchange failed: {error_detail}"
                )

            token_data = response.json()
            logger.info("Token exchange successful!")

            return {
                "access_token": token_data.get("access_token"),
                "refresh_token": token_data.get("refresh_token"),
                "token_type": "Bearer",
                "user": token_data.get("user"),  # Include user info
                "expires_in": token_data.get("expires_in"),
            }
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Callback error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# WorkOS OAuth Callback (full path)
@app.get("/api/v1/auth/workos/callback")
async def workos_callback_full(code: str):
    """Handle WorkOS OAuth callback - full path"""
    return await handle_workos_callback(code)

# WorkOS OAuth Callback (short path for direct redirects)
@app.get("/callback")
async def workos_callback_short(code: str):
    """Handle WorkOS OAuth callback - short path"""
    return await handle_workos_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"""
    logger.info(f"Query from {user.get('sub')}: {request.query}")

    # Simple echo response for testing
    return QueryResponse(
        response=f"Received your query: '{request.query}'. Backend is working with WorkOS auth!",
        query_id="test-123"
    )

# WebSocket endpoint stub (for future)
@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 REDIRECT_URI: {REDIRECT_URI}")
    uvicorn.run(app, host="0.0.0.0", port=7000)