import jwt
from flask import (
    Blueprint,
    current_app,
    make_response,
    redirect,
    request,
    session,
    url_for,
    render_template,
    flash,
)
from werkzeug.exceptions import BadRequest

from atlas_um import consts
from atlas_um import tasks
from atlas_um.analytics import user_events
from atlas_um.audit import user_audit
from atlas_um.auth.utils import no_mfa_is_active
from atlas_um.consts import SystemEvents
from atlas_um.exceptions import AccountIsNotActive
from atlas_um.extensions import usm
from atlas_um.helpers.domains import (
    get_root_domain,
    get_allowed_root_domain,
    get_allowed_subdomain,
)
from atlas_um.helpers.tokens import get_safe_repr
from atlas_um.logs import logger
from atlas_um.pgdb import pgdb, DNAAccount, DNAAccountActivity, ResourceGroup
from atlas_um.pgdb.dna_account import DNAAccountStatuses
from atlas_um.helpers.auth_proxy import get_auth_proxy
from atlas_um.tokens.services import (
    CreateBearerTokenService,
    CreateRefreshTokenService,
)

login = Blueprint("usm_login", __name__)


@login.route("", methods=["GET"])
def get_login():
    """
    Main endpoint to trigger the USM OIDC login flow.

    It is possible to omit the whole USM flow for developer needs
    on dev envs by passing the `no_mfa_account` param with needed account email
    and dev m2m token with `atlasum|no_mfa` aud  claim in `dev_token` param.
    """

    #  values here to reuse them value in both USM and no_mfa cases
    session[consts.REFERRER_SESSION_NAME] = (
        request.args.get("next") or request.referrer
    )
    session[consts.AUTH_PROXY_SESSION_NAME] = get_auth_proxy()

    if no_mfa_is_active() and request.args.get(consts.NO_MFA_ACCOUNT_PARAM):
        dna_account = DNAAccount.query.no_mfa_by_email(
            request.args.get(consts.NO_MFA_ACCOUNT_PARAM)
        ).first()
        if dna_account:
            logger.bind(dna_account=dna_account).info(
                "Processing no mfa login"
            )
            return process_login(dna_account)

    nonce = request.args.get("nonce", None)
    if nonce is None:
        nonce = usm.generate_nonce()
    state = usm.generate_state()
    prompt = request.args.get("prompt")
    usm_url = usm.authentication_url(state, nonce, prompt)
    session[consts.STATE_SESSION_NAME] = state
    session[consts.NONCE_SESSION_NAME] = nonce
    return redirect(usm_url)


@login.route("/callback", methods=["GET"])
def get_login_callback():
    """
    USM OIDC callback URL endpoint:

    - validates standard OIDC session protection params
    - completes OIDC flow, retrieving the ID token
    - detects the DNA account, and
    - executes the Atlas own login procedures
    """

    state_param = request.args.get("state", None)
    if state_param is None:
        logger.warning("No state in query on login callback")
        raise BadRequest("No state in query")
    session_state = session.get(consts.STATE_SESSION_NAME, None)
    if session_state is None:
        logger.warning("State param is missing from session on login callback")
        raise BadRequest(f"{consts.STATE_SESSION_NAME} missing from session")
    if session_state != state_param:
        logger.warning("No match for state in session/query on login callback")
        raise BadRequest(
            f"{consts.STATE_SESSION_NAME} "
            f"in session does not match state in query"
        )
    code_param = request.args.get("code", None)
    if code_param is None:
        logger.warning("No code in query on login callback")
        raise BadRequest("No code in query")
    e_access_token = usm.fetch_access_token(code_param)
    if e_access_token.is_left:
        logger.bind(error=e_access_token.value).warning(
            "Error fetching USM access token on login callback"
        )
        flash("Sorry, something went wrong. Please try again.", "negative")
        return redirect(url_for("common.userinfo"))
    id_token = e_access_token.value["id_token"]
    usm_claimset = jwt.decode(id_token, options={"verify_signature": False})
    dna_account = DNAAccount.selsert_from_usm_claimset(usm_claimset)
    if dna_account.status != DNAAccountStatuses.ACTIVE:
        logger.bind(dna_account=dna_account, usm_claimset=usm_claimset).info(
            "The account is not active on login callback"
        )
        if dna_account.expiration_date:
            raise AccountIsNotActive(
                "Your access to the Sony Analytics products has expired"
            )
        else:
            # fallback for future cases where we might have other cases
            raise AccountIsNotActive()
    pgdb.session.commit()
    tasks.sync_dna_account_with_usm.delay(dna_account.id)

    return process_login(dna_account, id_token)


def process_login(dna_account, usm_id_token=None):
    """
    Main Atlas login logic.

    Sets all required cookies, logs all needed events,
    redirects to auth proxy if needed, or to the final target.

    If the product auth proxy is being used, then also issuing
    separate product token with claims relevant only for the specific
    product. This allows to resolve possible issues when the same
    token would be refreshed both in atlas and product, and also
    make most tokens lighter.
    """

    cookie_max_age = current_app.config.get("TOKEN_COOKIES_MAX_AGE")

    authentication_referrer = session.get(consts.REFERRER_SESSION_NAME, None)
    if authentication_referrer is None:
        back_to = url_for("common.userinfo", _external=True)
    else:
        back_to = authentication_referrer

    domain = get_allowed_root_domain()
    resource_group = ResourceGroup.query.by_namespace_url(
        current_app.config.get("RELATED_CLAIMS_NAMESPACE")
    ).first()
    dna_bearer_token = CreateBearerTokenService.execute(dna_account).value
    dna_refresh_token = CreateRefreshTokenService.execute(
        dna_account, resource_group
    ).value

    auth_proxy = session.get(consts.AUTH_PROXY_SESSION_NAME)
    if auth_proxy:
        referrer_subdomain = (
            get_allowed_subdomain(authentication_referrer)
            if authentication_referrer
            else None
        )
        product_domain = (
            referrer_subdomain or get_root_domain(auth_proxy) or domain
        )
        product_resource_group = ResourceGroup.query.by_domain(
            product_domain
        ).first()
        product_bearer_token = CreateBearerTokenService.execute(
            dna_account, product_resource_group
        ).value
        product_refresh_token = CreateRefreshTokenService.execute(
            dna_account, product_resource_group
        ).value
        DNAAccountActivity.register_for_resource_group(
            dna_account, product_resource_group
        )
        resp = make_response(
            render_template(
                "auth_proxy.html",
                auth_proxy=auth_proxy,
                access_token=product_bearer_token,
                refresh_token=product_refresh_token,
                back_to=back_to,
            )
        )
    elif (
        authentication_referrer
        and (
            referrer_subdomain := get_allowed_subdomain(
                authentication_referrer
            )
        )
        != domain
    ):
        # atlas subdomain product flow
        product_domain = referrer_subdomain or domain  # used for logging
        product_resource_group = (
            ResourceGroup.query.by_domain(product_domain, True).first()
            or resource_group
        )
        product_bearer_token = dna_bearer_token
        product_refresh_token = dna_refresh_token
        DNAAccountActivity.register_for_resource_group(
            dna_account, product_resource_group
        )
        resp = make_response(redirect(back_to))
    else:
        # atlas flow
        product_domain = domain  # used for logging
        product_resource_group = resource_group
        product_bearer_token = dna_bearer_token
        product_refresh_token = dna_refresh_token
        resp = make_response(redirect(back_to))

    domain_cookie = f".{domain}" if domain else None  # used for cookies

    if usm_id_token:
        resp.set_cookie(
            consts.USM_ID_TOKEN_COOKIE_NAME,
            usm_id_token,
            httponly=True,
            domain=domain_cookie,
            secure=True,
            max_age=cookie_max_age,
        )
    resp.set_cookie(
        current_app.config.get("DNA_BEARER_TOKEN_COOKIE_NAME"),
        dna_bearer_token,
        httponly=True,
        domain=domain_cookie,
        secure=True,
        max_age=cookie_max_age,
    )
    resp.set_cookie(
        current_app.config.get("DNA_REFRESH_TOKEN_COOKIE_NAME"),
        dna_refresh_token,
        httponly=True,
        domain=domain_cookie,
        secure=True,
        max_age=cookie_max_age,
    )
    DNAAccountActivity.register_for_resource_group(dna_account, resource_group)
    user_audit.log(
        SystemEvents.successful_login,
        domain=product_domain,
        user_id=dna_account.sub,
        access_token=get_safe_repr(dna_bearer_token),
        refresh_token=get_safe_repr(dna_refresh_token),
        product_access_token=get_safe_repr(product_bearer_token),
        product_refresh_token=get_safe_repr(product_refresh_token),
    )
    user_events.log(
        SystemEvents.successful_login,
        domain=product_domain,
        user_id=dna_account.sub,
    )
    logger.bind(
        dna_account=dna_account,
        auth_proxy=auth_proxy,
        back_to=back_to,
        domain=domain,
        domain_cookie=domain_cookie,
        product_domain=product_domain,
        resource_group=resource_group,
        access_token=get_safe_repr(dna_bearer_token),
        refresh_token=get_safe_repr(dna_refresh_token),
        product_resource_group=product_resource_group,
        product_access_token=get_safe_repr(dna_bearer_token),
        product_refresh_token=get_safe_repr(dna_refresh_token),
    ).info("Processed user login")
    return resp