from fastapi import Depends, HTTPException, status, Request
from fastapi.security import OAuth2AuthorizationCodeBearer
from fastapi.security.utils import get_authorization_scheme_param
import os
import jwt

REGION = "us-east-2"
USER_POOL_ID = os.getenv("COGNITO_USER_POOL_ID")
CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")
DOMAIN = f"https://iptv-updater.auth.{REGION}.amazoncognito.com"

class CustomOAuth2(OAuth2AuthorizationCodeBearer):
    async def __call__(self, request: Request) -> str:
        # Check if this is a browser request
        is_browser = "text/html" in request.headers.get("accept", "")
        
        authorization = request.cookies.get("token")
        if not authorization:
            authorization = request.headers.get("Authorization")
        
        scheme, param = get_authorization_scheme_param(authorization)
        
        if not authorization or scheme.lower() != "bearer":
            if is_browser:
                redirect_uri = str(request.base_url)[:-1] + "/auth/callback"  # Remove trailing slash
                raise HTTPException(
                    status_code=302,
                    headers={
                        "Location": f"{DOMAIN}/login?client_id={CLIENT_ID}"
                        f"&response_type=code"
                        f"&scope=openid+email+profile"
                        f"&redirect_uri={redirect_uri}"
                    }
                )
            else:
                raise HTTPException(
                    status_code=401,
                    detail="Not authenticated",
                    headers={"WWW-Authenticate": "Bearer"},
                )
                
        return param

oauth2_scheme = CustomOAuth2(
    authorizationUrl=f"{DOMAIN}/oauth2/authorize",
    tokenUrl=f"{DOMAIN}/oauth2/token",
)

async def get_current_user(request: Request, token: str = Depends(oauth2_scheme)):
    try:
        decoded = jwt.decode(
            token,
            options={"verify_signature": False}
        )
        return {
            "Username": decoded.get("email") or decoded.get("sub"),
            "UserAttributes": [
                {"Name": k, "Value": v}
                for k, v in decoded.items()
            ]
        }
    except Exception as e:
        print(f"Token verification failed: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )