import grp
import os
import pwd
import uuid
from datetime import datetime, timedelta, UTC
from typing import List
import time

import jwt
from fastapi import Depends, HTTPException, Request, status
from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
from pydantic import BaseModel

from aare.common.auth_models import BatonRequestStatus, BatonTransferQueue, BatonRequest, BatonStatus
from aare.common.models import SessionsStateEnum
from aare.daq.config import BeamlineConfig

from aare.common.exception_handler import AuthenticationException, UserRightsException, AuthErrorCode

if os.environ.get("JWT_AAREDAQ_KEY") is None:
    raise Exception("JWT_AAREDAQ_KEY environment variable not set, cannot guarantee safe authentication.")

SECRET_KEY = os.environ.get("JWT_AAREDAQ_KEY")

ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 24 * 60 * 7 # 1 week
SESSION_EXPIRE_SECONDS = 60 * 10
BATON_REQUEST_TIMEOUT_SECONDS = 30

STAFF_GROUP = "unx-MXgroup"
SUPER_USERS = ["e10019", "e11206", "e18147"]

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

class TokenData(BaseModel):
    sub: str # Username
    pgroups: List[str]
    session: int
    staff: bool = False


def create_access_token(token: TokenData):
    to_encode = token.model_dump()
    to_encode.update({"exp": datetime.now(UTC) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt


def authenticate_from_proxy_header(request: Request) -> str:
    """
    Extract the pre-authenticated username from the trusted Apache proxy header.
    Only trust this header when the request originates from localhost (127.0.0.1 / ::1),
    meaning only the Apache proxy running on the same host can supply it.
    """
    client_host = request.client.host if request.client else None
    if client_host not in ("127.0.0.1", "::1"):
        raise AuthenticationException(
            message="X-Remote-User header is only trusted from the localhost proxy",
            status_code=401,
            headers={"WWW-Authenticate": "Bearer"},
            code=AuthErrorCode.INVALID_TOKEN,
        )
    remote_user = request.headers.get("X-Remote-User")
    if not remote_user:
        raise AuthenticationException(
            message="Missing X-Remote-User header from proxy",
            status_code=401,
            headers={"WWW-Authenticate": "Bearer"},
            code=AuthErrorCode.INVALID_TOKEN,
        )
    return remote_user


def authenticate_user(cfg: BeamlineConfig, username: str) -> str:
    user_info = pwd.getpwnam(username)

    groups = os.getgrouplist(user_info.pw_name, user_info.pw_gid)
    supplementary_groups = [grp.getgrgid(g).gr_name.lower() for g in groups]
    super_user = username in SUPER_USERS
    pgroups = [group for group in supplementary_groups if group.startswith('p')]
    staff = "unx-mxgroup" in supplementary_groups or "unx-sls_mx" in supplementary_groups or super_user
    token = TokenData(sub=username,
                      pgroups=pgroups,
                      staff=staff,
                      session=cfg.generate_session())
    return create_access_token(token)


def parse_token(token: str = Depends(oauth2_scheme)) -> TokenData:
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        token = TokenData(**payload)
        return token
    except jwt.PyJWTError as e:
        raise AuthenticationException(
            message="Invalid token",
            status_code=401,
            headers={"WWW-Authenticate": "Bearer"},
            code=AuthErrorCode.INVALID_TOKEN,
        ) from e


def check_jwt_ro(cfg: BeamlineConfig, data: TokenData) -> None:
    active_pgroup = cfg.pgroup
    if not data.staff and (active_pgroup is None or active_pgroup not in data.pgroups):
        raise UserRightsException(
            message="Not member of a currently active p-group.",
            status_code=403,
            code=AuthErrorCode.NOT_IN_ACTIVE_PGROUP,
        )


def check_jwt_rw(cfg: BeamlineConfig, data: TokenData) -> None:
    check_jwt_ro(cfg, data)

    holder = cfg.baton_holder
    if holder is None or holder.session != data.session:
        raise UserRightsException(
            message="You do not hold the baton.",
            status_code=403,
            code=AuthErrorCode.NOT_BATON_HOLDER,
        )

    try:
        cfg.try_extend_active_session(data.session, SESSION_EXPIRE_SECONDS)
    except Exception as e:
        # In case something is wrong but you are holder (maybe redis expiry?)
        cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS)

def check_jwt_staff_only(data: TokenData) -> None:
    if not data.staff:
        raise UserRightsException(
            message="Not member of the MX staff.",
            status_code=403,
            code=AuthErrorCode.NOT_STAFF,
        )

def check_jwt_staff(cfg: BeamlineConfig, data: TokenData) -> None:
    check_jwt_staff_only(data)

    holder = cfg.baton_holder
    if holder and holder.session != data.session:
         # Staff can take over if they don't have it, but they need to use force_current_session
         # or request_baton (which does staff override).
         # If they are calling a RW endpoint, they SHOULD already be the holder.
         pass

    try:
        cfg.try_extend_active_session(data.session, SESSION_EXPIRE_SECONDS)
    except Exception:
        cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS)

def force_current_sesion(cfg: BeamlineConfig, data: TokenData) -> None:
    cfg.execute_baton_transfer(
        to_session=data.session,
        to_username=data.sub,
        to_is_staff=data.staff,
        to_pgroup=cfg.pgroup,
        expiry_sec=SESSION_EXPIRE_SECONDS
    )

def _finalize_expired_baton_request(cfg: BeamlineConfig, pending: BatonRequest) -> BatonStatus:
    """
    Resolve an expired baton request in one place.

    Rules:
      - staff can override immediately if beamline is free
      - otherwise the transfer is queued
      - pending request is cleared once it is no longer pending
    """
    requester_is_staff = bool(pending.requester_is_staff)

    if cfg.can_transfer_baton_now():
        cfg.execute_baton_transfer(
            to_session=pending.requester_session,
            to_username=pending.requester_username,
            to_is_staff=requester_is_staff,
            to_pgroup=cfg.pgroup,
            expiry_sec=SESSION_EXPIRE_SECONDS
        )
    else:
        cfg.queued_baton_transfer = BatonTransferQueue(
            target_session=pending.requester_session,
            target_username=pending.requester_username,
            target_is_staff=requester_is_staff,
            target_pgroup=cfg.pgroup,
            queued_at=time.time(),
            reason="timeout_beamline_busy"
        )

    cfg.clear_pending_baton_request()
    return get_baton_status(cfg, TokenData(
        sub=pending.requester_username,
        pgroups=[],
        session=pending.requester_session,
        staff=requester_is_staff
    ))

def resolve_baton_timeout_if_needed(cfg: BeamlineConfig) -> BatonStatus | None:
    """
    Check the current pending request and resolve it if expired.
    Returns the updated BatonStatus when a timeout was processed, else None.
    """
    pending = cfg.pending_baton_request
    if pending is None or pending.status != BatonRequestStatus.PENDING:
        return None

    elapsed = time.time() - pending.created_at
    if elapsed < pending.timeout_seconds:
        return None

    return _finalize_expired_baton_request(cfg, pending)

def get_baton_status(cfg: BeamlineConfig, data: TokenData) -> BatonStatus:
    """
    Build baton status scoped to the requesting session.

    Important:
    - requester sees you_have_pending_request
    - holder sees incoming_request
    - nobody else sees the request as actionable
    """
    holder = cfg.baton_holder
    pending = cfg.pending_baton_request

    is_requester = bool(
        pending
        and pending.status in (BatonRequestStatus.PENDING, BatonRequestStatus.REFUSED)
        and pending.requester_session == data.session
    )

    is_holder = bool(
        pending
        and pending.status == BatonRequestStatus.PENDING
        and pending.holder_session == data.session
    )

    # Only expose the pending request object to the two relevant sessions
    scoped_pending = pending if (is_requester or is_holder) else None

    return BatonStatus(
        holder=holder,
        pending_request=scoped_pending,
        queued_transfer=cfg.queued_baton_transfer,
        you_are_holder=bool(holder and holder.session == data.session),
        you_have_pending_request=is_requester,
        incoming_request=is_holder,
        allow_non_staff_request=cfg.allow_non_staff_request_from_staff,
    )

def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict:
    resolve_baton_timeout_if_needed(cfg)

    session_state = cfg.session_state(data.session)

    if session_state == SessionsStateEnum.Vacant:
        # State: VACANT -> OWNED (B)
        cfg.execute_baton_transfer(
            to_session=data.session,
            to_username=data.sub,
            to_is_staff=data.staff,
            to_pgroup=cfg.pgroup,
            expiry_sec=SESSION_EXPIRE_SECONDS,
        )
        return {"granted": True, "message": "Baton acquired (beamline was vacant)"}

    if session_state == SessionsStateEnum.OwnedByYou:
        cfg.try_extend_active_session(data.session, SESSION_EXPIRE_SECONDS)
        return {"already_holder": True, "message": "You already hold the baton"}

    if session_state == SessionsStateEnum.PendingYouToElse:
        # Already requested, waiting for holder response or timeout
        existing_request = cfg.pending_baton_request
        if existing_request:
            elapsed = time.time() - existing_request.created_at
            remaining = existing_request.timeout_seconds - elapsed
            return {
                "pending": True,
                "existing": True,
                "remaining_seconds": max(0, remaining),
                "message": f"Request already pending ({remaining:.0f}s remaining)",
            }

    holder = cfg.baton_holder
    # Handle Staff override
    if data.staff:
        if not cfg.can_transfer_baton_now():
            # State: OWNED (A) -> PENDING (A -> B, B=staff) -> Queue -> OWNED (B)
            cfg.queued_baton_transfer = BatonTransferQueue(
                target_session=data.session,
                target_username=data.sub,
                target_is_staff=data.staff,
                target_pgroup=cfg.pgroup,
                queued_at=time.time(),
                reason="beamline_busy_staff_override",
            )
            return {
                "queued": True,
                "message": "Staff override queued - will transfer when beamline is available",
            }

        # State: OWNED (A) -> OWNED (B, B=staff)
        cfg.execute_baton_transfer(
            to_session=data.session,
            to_username=data.sub,
            to_is_staff=data.staff,
            to_pgroup=cfg.pgroup,
            expiry_sec=SESSION_EXPIRE_SECONDS,
        )
        return {"granted": True, "override": True, "message": "Staff override - baton acquired"}

    # Policy check: non-staff requesting from staff
    if holder and holder.is_staff and not data.staff and not cfg.allow_non_staff_request_from_staff:
        return {
            "error": True,
            "message": "Requesting baton from staff is disabled by backend policy.",
        }

    # State: OWNED (A) -> PENDING (A -> B)
    existing_request = cfg.pending_baton_request
    if existing_request and existing_request.status == BatonRequestStatus.PENDING:
        # Someone else already has a request pending
        return {
            "error": True,
            "message": f"Another user ({existing_request.requester_username}) already has a pending request",
        }

    # Beamline busy notification
    is_busy = not cfg.can_transfer_baton_now()

    request = BatonRequest(
        request_id=str(uuid.uuid4()),
        requester_username=data.sub,
        requester_session=data.session,
        requester_is_staff=data.staff,
        holder_username=holder.username if holder else None,
        holder_session=holder.session if holder else None,
        created_at=time.time(),
        timeout_seconds=BATON_REQUEST_TIMEOUT_SECONDS,
        status=BatonRequestStatus.PENDING,
    )
    cfg.set_pending_baton_request(request, timeout_sec=BATON_REQUEST_TIMEOUT_SECONDS)

    msg = f"Request sent to {holder.username if holder else 'current holder'}"
    if is_busy:
        msg += " (Note: beamline is currently busy, transfer will be queued if accepted)"

    return {
        "pending": True,
        "request_id": request.request_id,
        "timeout_seconds": BATON_REQUEST_TIMEOUT_SECONDS,
        "message": msg,
        "beamline_busy": is_busy,
    }

def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) -> dict:
    """
    Current baton holder responds to a pending request.
    """
    resolve_baton_timeout_if_needed(cfg)

    holder = cfg.baton_holder
    if holder is None or holder.session != data.session:
        return {"error": True, "message": "You are not the current baton holder"}

    pending = cfg.pending_baton_request
    if pending is None or pending.status != BatonRequestStatus.PENDING:
        return {"error": True, "message": "No pending request to respond to"}

    if accept:
        # State: PENDING (A -> B) -> OWNED (B)
        if cfg.can_transfer_baton_now():
            cfg.execute_baton_transfer(
                to_session=pending.requester_session,
                to_username=pending.requester_username,
                to_is_staff=pending.requester_is_staff,
                to_pgroup=cfg.pgroup,
                expiry_sec=SESSION_EXPIRE_SECONDS
            )
            return {"accepted": True, "transferred": True, "message": "Baton transferred"}
        else:
            # Beamline busy, queue the transfer
            cfg.queued_baton_transfer = BatonTransferQueue(
                target_session=pending.requester_session,
                target_username=pending.requester_username,
                target_is_staff=pending.requester_is_staff,
                target_pgroup=cfg.pgroup,
                queued_at=time.time(),
                reason="accepted_beamline_busy"
            )
            cfg.clear_pending_baton_request()
            return {
                "accepted": True,
                "queued": True,
                "message": "Request accepted. Baton will be transferred as soon as beamline is available.",
            }
    else:
        # State: PENDING (A -> B) -> OWNED (A)
        pending.status = BatonRequestStatus.REFUSED
        cfg.set_pending_baton_request(pending, timeout_sec=10)
        return {"refused": True, "message": "Baton request refused"}

def release_baton(cfg: BeamlineConfig, data: TokenData) -> dict:
    """
    Voluntarily release the baton (set session to free).
    """
    resolve_baton_timeout_if_needed(cfg)

    holder = cfg.baton_holder
    if holder is None:
        return {"released": True, "message": "Baton was already vacant"}

    if holder.session != data.session:
        return {"info": True, "message": "You don't hold the baton"}

    cfg.end_active_session(data.session)
    cfg.baton_holder = None
    cfg.clear_pending_baton_request()

    return {"released": True, "message": "Baton released - beamline is now vacant"}

def cancel_baton_request(cfg: BeamlineConfig, data: TokenData) -> dict:
    """
    Cancel your own pending baton request.
    """
    resolve_baton_timeout_if_needed(cfg)

    pending = cfg.pending_baton_request
    if pending is None:
        return {"error": True, "message": "No pending request to cancel"}

    if pending.requester_session != data.session:
        return {"error": True, "message": "You can only cancel your own request"}

    cfg.clear_pending_baton_request()
    return {"cancelled": True, "message": "Request cancelled"}