# Cognito configuration
import base64
import json
# from utils.logger import ServiceLogger
import logging
import os
from urllib.parse import urlencode

import aiohttp
import boto3
import jwt
import requests
from fastapi import HTTPException, status
from jwt.algorithms import RSAAlgorithm
from pydantic import BaseModel, EmailStr

from configs.config import (ACCESS_TOKEN, COGNITO_CLIENT_ID,
                            COGNITO_CLIENT_SECRET, COGNITO_DOMAIN,
                            COGNITO_LOGOUT_REDIRECT_URI, COGNITO_REDIRECT_URI,
                            COGNITO_REGION, COGNITO_USER_POOL_ID, ID_TOKEN,
                            REFRESH_TOKEN)
from services.user_service import UserService

XCM_logger = logging.getLogger()


# Initialize Cognito client
cognito_client = boto3.client("cognito-idp", region_name=COGNITO_REGION)


class CreateUserRequest(BaseModel):
    email: EmailStr
    first_name: str
    last_name: str
    role: str


class TokenResponse(BaseModel):
    access_token: str
    id_token: str
    refresh_token: str
    expires_in: int
    token_type: str


def get_cognito_public_keys():
    keys_url = f"https://cognito-idp.{COGNITO_REGION}.amazonaws.com/{COGNITO_USER_POOL_ID}/.well-known/jwks.json"
    response = requests.get(keys_url)
    return response.json()["keys"]


def validate_jwt_token(token: str):
    try:
        headers = jwt.get_unverified_header(token)
        key_id = headers["kid"]

        # Get the public keys
        public_keys = get_cognito_public_keys()
        key = next((k for k in public_keys if k["kid"] == key_id), None)

        if not key:
            raise HTTPException(status_code=401, detail="Invalid token")

        public_key = RSAAlgorithm.from_jwk(json.dumps(key))

        import time

        # Verify and decode the token
        payload = jwt.decode(
            token,
            public_key,
            algorithms=["RS256"],
            audience=COGNITO_CLIENT_ID,
            leeway=10,
        )

        iat = payload.get("iat", 0)  # Issued At time
        current_time = int(time.time())
        time_difference = iat - current_time
        if time_difference > 0:
            # TODO log this for the specific user and send an email
            print(
                "⚠️ WARNING: Your server clock is behind by", time_difference, "seconds."
            )
        return payload

    except Exception as e:
        raise HTTPException(status_code=401, detail=str(e)) from e


async def create_user_cognito_pool(
    user_data: CreateUserRequest, temporary_password: str
):
    """
    Creates a new user in Cognito user pool with the specified email and temporary password.
    The user will be required to change their password on first login.
    """
    try:
        # Create user in Cognito
        response = cognito_client.admin_create_user(
            UserPoolId=COGNITO_USER_POOL_ID,
            Username=user_data.email,
            UserAttributes=[
                {"Name": "email", "Value": user_data.email},
                {"Name": "email_verified", "Value": "true"},
                *(
                    [{"Name": "given_name", "Value": user_data.first_name}]
                    if user_data.first_name
                    else []
                ),
                *(
                    [{"Name": "family_name", "Value": user_data.last_name}]
                    if user_data.last_name
                    else []
                ),
            ],
            TemporaryPassword=temporary_password,
        )

        # Extract user information from the response under the "User" key
        user = response.get("User", {})
        attributes_list = user.get("Attributes", [])
        # Convert the attributes list to a dictionary for easier lookup
        user_attributes = {attr["Name"]: attr["Value"] for attr in attributes_list}

        XCM_logger.info(f"Created new user: {user_data.email}")
        return {
            "message": "User created successfully",
            "username": user.get("Username"),
            "email": user_attributes.get("email"),
            "first_name": user_attributes.get("given_name"),
            "last_name": user_attributes.get("family_name"),
        }

    except cognito_client.exceptions.UsernameExistsException as e:
        raise HTTPException(
            status_code=400, detail="User with this email already exists"
        ) from e
    except Exception as e:
        XCM_logger.error(f"Error creating user: {str(e)}")
        raise HTTPException(
            status_code=500, detail=f"Error creating user: {str(e)}"
        ) from e


class AuthService:
    def __init__(self):
        pass

    async def login(self):
        """Generates Cognito login URL"""
        params = {
            "redirect_uri": COGNITO_REDIRECT_URI,
            "client_id": COGNITO_CLIENT_ID,
            "response_type": "code",
            "scope": "openid email profile",
        }
        return f"{COGNITO_DOMAIN}/login?{urlencode(params)}"

    async def handle_callback(self, code: str):
        """Handles the OAuth callback and token exchange"""
        if not code:
            raise HTTPException(
                status_code=400, detail="Authorization code not provided."
            )
        # import anyio
        # await anyio.sleep(2)
        token_url = f"{COGNITO_DOMAIN}/oauth2/token"
        auth_header = self._create_basic_auth_header()

        payload = {
            "grant_type": "authorization_code",
            "client_id": COGNITO_CLIENT_ID,
            "code": code,
            "redirect_uri": COGNITO_REDIRECT_URI,
        }

        headers = {
            "Content-Type": "application/x-www-form-urlencoded",
            "Authorization": f"Basic {auth_header}",
        }

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    token_url, data=payload, headers=headers
                ) as response:
                    if response.status != 200:
                        error_detail = await response.json()
                        XCM_logger.error(f"Token exchange failed: {error_detail}")
                        raise HTTPException(
                            status_code=response.status,
                            detail="Failed to retrieve tokens.",
                        )

                    tokens = await response.json()
                    return await self._process_tokens(tokens)

        except aiohttp.ClientError as e:
            XCM_logger.error("Cognito API error: %s", str(e))
            print(e)
            raise HTTPException(
                status_code=500,
                detail=f"Error connecting to Cognito API: {str(e)}",
            ) from e

    async def refresh_token(self, refresh_token: str):
        """Refreshes the access token using a refresh token"""
        if not refresh_token:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Refresh token not provided",
            )

        token_url = f"{COGNITO_DOMAIN}/oauth2/token"
        auth_header = self._create_basic_auth_header()

        payload = {
            "grant_type": REFRESH_TOKEN,
            "client_id": COGNITO_CLIENT_ID,
            "refresh_token": refresh_token,
        }

        headers = {
            "Content-Type": "application/x-www-form-urlencoded",
            "Authorization": f"Basic {auth_header}",
        }

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    token_url, data=payload, headers=headers
                ) as response:
                    response.raise_for_status()
                    tokens = await response.json()
                    return tokens
        except aiohttp.ClientError as e:
            XCM_logger.error("Token refresh failed: %s", str(e))
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Failed to refresh token",
            ) from e

    def get_logout_url(self):
        """Generates the Cognito logout URL"""
        return f"{COGNITO_DOMAIN}/logout?client_id={COGNITO_CLIENT_ID}&logout_uri={COGNITO_LOGOUT_REDIRECT_URI}"

    def _create_basic_auth_header(self):
        """Creates the Basic Auth header for Cognito API calls"""
        auth_string = f"{COGNITO_CLIENT_ID}:{COGNITO_CLIENT_SECRET}"
        auth_bytes = auth_string.encode("ascii")
        return base64.b64encode(auth_bytes).decode("ascii")

    async def _process_tokens(self, tokens):
        """Process and validate tokens, create user if needed"""
        id_token = tokens.get(ID_TOKEN)
        access_token = tokens.get(ACCESS_TOKEN)
        refresh_token = tokens.get(REFRESH_TOKEN)

        user_info = validate_jwt_token(id_token)

        user_service = UserService()
        user_response = user_service.get_or_create_user(user_info)
        is_new_user = user_response.get("isNewUser")

        frontend_url = (
            f"/account/setup/user?id_token={id_token}&access_token={access_token}"
            if is_new_user
            else f"/project/listing?id_token={id_token}&access_token={access_token}"
        )

        return {
            "tokens": tokens,
            "frontend_url": frontend_url,
            "refresh_token": refresh_token,
        }