import os
import json
import boto3
from jose import jwt
from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import OAuth2AuthorizationCodeBearer

app = FastAPI()

# Get Cognito info from environment (set by userdata.sh)
REGION = "us-east-2"
USER_POOL_ID = os.getenv("COGNITO_USER_POOL_ID")
CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")

# OAuth2 scheme for authorization code flow
oauth2_scheme = OAuth2AuthorizationCodeBearer(
    authorizationUrl=f"https://cognito-idp.{REGION}.amazonaws.com/{USER_POOL_ID}/oauth2/authorize",
    tokenUrl=f"https://cognito-idp.{REGION}.amazonaws.com/{USER_POOL_ID}/oauth2/token"
)

async def get_current_user(token: str = Depends(oauth2_scheme)):
    try:
        # Verify the JWT token with Cognito
        cognito_idp = boto3.client('cognito-idp', region_name=REGION)
        response = cognito_idp.get_user(
            AccessToken=token
        )
        return response
    except Exception as e:
        raise HTTPException(
            status_code=401,
            detail="Invalid authentication credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )

@app.get("/")
async def root():
    return {"message": "IPTV Updater API"}

@app.get("/health")
async def health():
    return {"status": "healthy"}

@app.get("/protected")
async def protected_route(user = Depends(get_current_user)):
    return {"message": "This is a protected route", "user": user['Username']}