#!/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")
REDIRECT_URI = os.getenv("REDIRECT_URI", "http://localhost:3000/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"
    }

# 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
    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"""
    try:
        async with httpx.AsyncClient() as client:
            # Exchange code for token (using client_credentials auth)
            response = await client.post(
                "https://api.workos.com/sso/token",
                auth=(WORKOS_API_KEY, ""),
                data={
                    "client_id": WORKOS_CLIENT_ID,
                    "code": code,
                    "grant_type": "authorization_code",
                },
            )

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

            token_data = response.json()
            return {
                "access_token": token_data.get("access_token"),
                "refresh_token": token_data.get("refresh_token"),
                "token_type": "Bearer",
            }
    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
    uvicorn.run(app, host="0.0.0.0", port=7000)