from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import secrets
import time
from typing import Any

from grantflow.api.redis_store import redis_available, redis_jti_is_revoked, redis_jti_revoke

from fastapi import FastAPI, HTTPException, Request
from fastapi.openapi.utils import get_openapi

from grantflow.api.errors import agent_error_detail
from grantflow.api.openapi_examples import install_agent_openapi_examples

PROTECTED_OPERATIONS = {
    ("post", "/generate"),
    ("post", "/generate/from-preset"),
    ("post", "/generate/from-preset/batch"),
    ("post", "/generate/preflight"),
    ("post", "/ingest"),
    ("post", "/ingest/text"),
    ("post", "/ingest/readiness"),
    ("get", "/ingest/recent"),
    ("get", "/ingest/inventory"),
    ("get", "/ingest/inventory/export"),
    ("get", "/queue/worker-heartbeat"),
    ("get", "/queue/dead-letter"),
    ("get", "/queue/dead-letter/export"),
    ("post", "/queue/dead-letter/requeue"),
    ("delete", "/queue/dead-letter"),
    ("post", "/cancel/{job_id}"),
    ("post", "/resume/{job_id}"),
    ("post", "/hitl/approve"),
    ("post", "/export"),
    ("get", "/status/{job_id}"),
    ("get", "/status/{job_id}/citations"),
    ("get", "/status/{job_id}/export-payload"),
    ("get", "/status/{job_id}/pilot-quick-report/export"),
    ("get", "/status/{job_id}/pilot-evidence-pack/export"),
    ("get", "/status/{job_id}/versions"),
    ("get", "/status/{job_id}/diff"),
    ("get", "/status/{job_id}/events"),
    ("get", "/status/{job_id}/events/export"),
    ("get", "/status/{job_id}/hitl/history"),
    ("get", "/status/{job_id}/hitl/history/export"),
    ("get", "/status/{job_id}/metrics"),
    ("get", "/status/{job_id}/quality"),
    ("get", "/status/{job_id}/pilot-quick-report"),
    ("get", "/status/{job_id}/grounding-gate"),
    ("get", "/status/{job_id}/critic"),
    ("get", "/status/{job_id}/critic/findings"),
    ("get", "/status/{job_id}/critic/findings/{finding_id}"),
    ("post", "/status/{job_id}/critic/findings/{finding_id}/ack"),
    ("post", "/status/{job_id}/critic/findings/{finding_id}/open"),
    ("post", "/status/{job_id}/critic/findings/{finding_id}/resolve"),
    ("post", "/status/{job_id}/critic/findings/bulk-status"),
    ("get", "/status/{job_id}/comments"),
    ("get", "/status/{job_id}/comments/export"),
    ("get", "/status/{job_id}/review/workflow"),
    ("get", "/status/{job_id}/review/workflow/trends"),
    ("get", "/status/{job_id}/review/workflow/sla"),
    ("get", "/status/{job_id}/review/workflow/sla/hotspots"),
    ("get", "/status/{job_id}/review/workflow/sla/hotspots/trends"),
    ("get", "/status/{job_id}/review/workflow/sla/hotspots/export"),
    ("get", "/status/{job_id}/review/workflow/sla/hotspots/trends/export"),
    ("get", "/status/{job_id}/review/workflow/sla/trends"),
    ("get", "/status/{job_id}/review/workflow/sla/export"),
    ("get", "/status/{job_id}/review/workflow/sla/trends/export"),
    ("get", "/status/{job_id}/review/workflow/sla/profile"),
    ("post", "/status/{job_id}/review/workflow/sla/recompute"),
    ("get", "/status/{job_id}/review/workflow/export"),
    ("get", "/status/{job_id}/review/workflow/trends/export"),
    ("post", "/status/{job_id}/comments"),
    ("post", "/status/{job_id}/comments/bulk-status"),
    ("post", "/status/{job_id}/comments/{comment_id}/ack"),
    ("post", "/status/{job_id}/comments/{comment_id}/resolve"),
    ("post", "/status/{job_id}/comments/{comment_id}/reopen"),
    ("get", "/portfolio/metrics"),
    ("get", "/portfolio/metrics/export"),
    ("get", "/portfolio/quality"),
    ("get", "/portfolio/quality/export"),
    ("get", "/portfolio/review-workflow"),
    ("get", "/portfolio/review-workflow/export"),
    ("get", "/portfolio/review-workflow/sla"),
    ("get", "/portfolio/review-workflow/sla/export"),
    ("get", "/portfolio/review-workflow/sla/hotspots"),
    ("get", "/portfolio/review-workflow/sla/hotspots/export"),
    ("get", "/portfolio/review-workflow/sla/hotspots/trends"),
    ("get", "/portfolio/review-workflow/sla/hotspots/trends/export"),
    ("get", "/portfolio/review-workflow/trends"),
    ("get", "/portfolio/review-workflow/trends/export"),
    ("get", "/portfolio/review-workflow/sla/trends"),
    ("get", "/portfolio/review-workflow/sla/trends/export"),
    ("get", "/hitl/pending"),
}

_REVOKED_AGENT_CREDENTIAL_JTIS: set[str] = set()


def api_key_configured() -> str | None:
    return os.getenv("GRANTFLOW_API_KEY") or os.getenv("API_KEY")


def self_serve_agent_keys_enabled() -> bool:
    return os.getenv("GRANTFLOW_AGENT_SELF_SERVE_KEYS_ENABLED", "false").strip().lower() == "true"


def create_self_serve_agent_api_key(
    *,
    agent_id: str,
    tenant_id: str,
    scopes: list[str],
    ttl_seconds: int = 86_400,
) -> dict[str, Any]:
    if not self_serve_agent_keys_enabled():
        raise ValueError("Self-serve agent API keys are disabled")
    secret = _agent_key_signing_secret()
    now = int(time.time())
    payload = {
        "agent_id": agent_id,
        "tenant_id": tenant_id,
        "scopes": sorted(set(scopes)),
        "jti": f"agt_{secrets.token_urlsafe(16)}",
        "iat": now,
        "exp": now + max(60, min(int(ttl_seconds), 604_800)),
    }
    encoded_payload = _b64url_json(payload)
    signature = _sign_agent_key(encoded_payload, secret)
    return {
        "type": "api_key",
        "header": "X-API-Key",
        "value": f"gf_agent_{encoded_payload}.{signature}",
        "expires_at": payload["exp"],
        "scopes": payload["scopes"],
        "jti": payload["jti"],
    }


def create_self_serve_oauth_client_secret(
    *,
    agent_id: str,
    tenant_id: str,
    scopes: list[str],
    ttl_seconds: int = 604_800,
) -> dict[str, Any]:
    if not self_serve_agent_keys_enabled():
        raise ValueError("Self-serve agent OAuth clients are disabled")
    secret = _agent_key_signing_secret()
    now = int(time.time())
    client_id = f"gf_client_{agent_id}"
    payload = {
        "agent_id": agent_id,
        "client_id": client_id,
        "tenant_id": tenant_id,
        "scopes": sorted(set(scopes)),
        "jti": f"oauth_{secrets.token_urlsafe(16)}",
        "iat": now,
        "exp": now + max(300, min(int(ttl_seconds), 2_592_000)),
        "credential_type": "oauth_client_secret",
    }
    encoded_payload = _b64url_json(payload)
    signature = _sign_agent_key(encoded_payload, secret)
    return {
        "client_id": client_id,
        "client_secret": f"gf_oauth_client_{encoded_payload}.{signature}",
        "client_secret_expires_at": payload["exp"],
        "token_endpoint": "/agents/oauth/token",
        "grant_type": "client_credentials",
        "requested_scopes": payload["scopes"],
        "tenant_id": tenant_id,
        "jti": payload["jti"],
    }


def create_self_serve_agent_access_token(
    *,
    client_id: str,
    client_secret: str,
    requested_scopes: list[str],
    ttl_seconds: int = 3_600,
) -> dict[str, Any] | None:
    claims = self_serve_oauth_client_secret_claims(client_secret)
    if claims is None or claims.get("client_id") != client_id:
        return None
    allowed_scopes = {str(scope) for scope in claims.get("scopes") or []}
    scopes = sorted(set(requested_scopes or allowed_scopes))
    if not set(scopes).issubset(allowed_scopes):
        return None
    credential = create_self_serve_agent_api_key(
        agent_id=str(claims["agent_id"]),
        tenant_id=str(claims["tenant_id"]),
        scopes=scopes,
        ttl_seconds=ttl_seconds,
    )
    return {
        "access_token": credential["value"],
        "token_type": "Bearer",
        "expires_in": max(1, int(credential["expires_at"]) - int(time.time())),
        "expires_at": credential["expires_at"],
        "scope": " ".join(credential["scopes"]),
        "tenant_id": claims["tenant_id"],
        "agent_id": claims["agent_id"],
    }


def read_auth_required() -> bool:
    return os.getenv("GRANTFLOW_REQUIRE_AUTH_FOR_READS", "false").strip().lower() == "true"


def require_api_key_if_configured(
    request: Request,
    *,
    for_read: bool = False,
    required_scope: str | None = None,
    tenant_id: str | None = None,
) -> None:
    expected = api_key_configured()
    if not expected:
        return
    if for_read and not read_auth_required():
        return

    provided = credential_from_request(request)
    claims = _api_key_claims(provided, expected) if provided else None
    if claims is None:
        raise HTTPException(
            status_code=401,
            detail=agent_error_detail(
                code="unauthorized",
                message="A valid X-API-Key header is required for this endpoint.",
                retryable=False,
                next_action="Send X-API-Key with a valid GrantFlow API key, or use public discovery endpoints first.",
            ),
        )
    if required_scope:
        _require_agent_scope(claims, required_scope)
    if tenant_id:
        _require_agent_tenant(claims, tenant_id)


def _valid_api_key(provided: str, expected: str) -> bool:
    return _api_key_claims(provided, expected) is not None


def _api_key_claims(provided: str, expected: str) -> dict[str, Any] | None:
    if secrets.compare_digest(provided, expected):
        return {"type": "operator", "scopes": ["*"]}
    return self_serve_agent_api_key_claims(provided)


def validate_self_serve_agent_api_key(provided: str) -> bool:
    return self_serve_agent_api_key_claims(provided) is not None


def introspect_api_key(provided: str | None) -> dict[str, Any]:
    if not provided:
        return {
            "active": False,
            "credential_type": None,
            "next_action": "Send X-API-Key with an operator key or self-serve agent key.",
        }
    expected = api_key_configured()
    if expected and secrets.compare_digest(provided, expected):
        return {
            "active": True,
            "credential_type": "operator",
            "agent_id": None,
            "tenant_id": None,
            "scopes": ["*"],
            "issued_at": None,
            "expires_at": None,
            "next_action": "Use operator credentials only for trusted automation and gateway services.",
        }
    claims = self_serve_agent_api_key_claims(provided)
    if claims is None:
        return {
            "active": False,
            "credential_type": None,
            "next_action": "Request a fresh agent credential with /agents/session or /agents/onboarding.",
        }
    return {
        "active": True,
        "credential_type": "self_serve_agent",
        "agent_id": claims.get("agent_id"),
        "tenant_id": claims.get("tenant_id"),
        "scopes": claims.get("scopes") or [],
        "issued_at": claims.get("iat"),
        "expires_at": claims.get("exp"),
        "jti": claims.get("jti"),
        "next_action": "Use this credential only for the returned tenant_id and scopes.",
    }


def self_serve_agent_api_key_claims(provided: str) -> dict[str, Any] | None:
    if not self_serve_agent_keys_enabled() or not provided.startswith("gf_agent_"):
        return None
    try:
        token = provided.removeprefix("gf_agent_")
        encoded_payload, signature = token.split(".", 1)
        expected_signature = _sign_agent_key(encoded_payload, _agent_key_signing_secret())
        if not secrets.compare_digest(signature, expected_signature):
            return None
        payload = _decode_b64url_json(encoded_payload)
        if _credential_jti_revoked(payload):
            return None
        if int(payload.get("exp") or 0) < int(time.time()):
            return None
        payload["type"] = "self_serve_agent"
        return payload
    except Exception:
        return None


def self_serve_oauth_client_secret_claims(provided: str) -> dict[str, Any] | None:
    if not self_serve_agent_keys_enabled() or not provided.startswith("gf_oauth_client_"):
        return None
    try:
        token = provided.removeprefix("gf_oauth_client_")
        encoded_payload, signature = token.split(".", 1)
        expected_signature = _sign_agent_key(encoded_payload, _agent_key_signing_secret())
        if not secrets.compare_digest(signature, expected_signature):
            return None
        payload = _decode_b64url_json(encoded_payload)
        if payload.get("credential_type") != "oauth_client_secret":
            return None
        if _credential_jti_revoked(payload):
            return None
        if int(payload.get("exp") or 0) < int(time.time()):
            return None
        return payload
    except Exception:
        return None


def revoke_self_serve_agent_credential(provided: str) -> dict[str, Any]:
    claims = _self_serve_agent_api_key_claims_unchecked(provided)
    credential_type = "self_serve_agent"
    if claims is None:
        claims = _self_serve_oauth_client_secret_claims_unchecked(provided)
        credential_type = "oauth_client_secret"
    if claims is None or not claims.get("jti"):
        return {
            "status": "not_revokable",
            "credential_type": None,
            "next_action": "Provide a valid self-serve API key, Bearer token, or OAuth client secret.",
        }
    revoked_jti = str(claims["jti"])
    _REVOKED_AGENT_CREDENTIAL_JTIS.add(revoked_jti)
    jti_ttl = max(86_400, int(claims.get("exp") or 0) - int(time.time()) + 86_400)
    redis_jti_revoke(revoked_jti, ttl_seconds=jti_ttl)
    scope = "redis_and_in_process" if redis_available() else "in_process"
    return {
        "status": "revoked",
        "credential_type": credential_type,
        "agent_id": claims.get("agent_id"),
        "tenant_id": claims.get("tenant_id"),
        "revoked_jti": revoked_jti,
        "revocation_scope": scope,
        "restart_safe": redis_available(),
        "next_action": "Discard this credential and request a fresh one before the next protected call.",
        "revocation_note": (
            "Revocation is enforced via Redis and an in-process JTI denylist."
            if redis_available()
            else "Revocation is enforced via an in-process JTI denylist and does not persist across "
            "process restarts. For production deployments, enforce revocation at the gateway layer."
        ),
    }


def rotate_self_serve_agent_credential(provided: str, *, ttl_seconds: int) -> dict[str, Any] | None:
    api_key_claims = _self_serve_agent_api_key_claims_unchecked(provided)
    if api_key_claims is not None and api_key_claims.get("jti") and not _credential_jti_revoked(api_key_claims):
        old_jti = str(api_key_claims["jti"])
        _REVOKED_AGENT_CREDENTIAL_JTIS.add(old_jti)
        jti_ttl = max(86_400, int(api_key_claims.get("exp") or 0) - int(time.time()) + 86_400)
        redis_jti_revoke(old_jti, ttl_seconds=jti_ttl)
        credential = create_self_serve_agent_api_key(
            agent_id=str(api_key_claims["agent_id"]),
            tenant_id=str(api_key_claims["tenant_id"]),
            scopes=[str(scope) for scope in api_key_claims.get("scopes") or []],
            ttl_seconds=ttl_seconds,
        )
        return {
            "status": "rotated",
            "credential_type": "self_serve_agent",
            "agent_id": api_key_claims["agent_id"],
            "tenant_id": api_key_claims["tenant_id"],
            "revoked_jti": api_key_claims["jti"],
            "credential": credential,
            "next_actions": [
                "Store credential.value securely; it is returned once.",
                "Replace the old X-API-Key or Bearer token immediately.",
            ],
        }
    oauth_claims = _self_serve_oauth_client_secret_claims_unchecked(provided)
    if oauth_claims is None or not oauth_claims.get("jti") or _credential_jti_revoked(oauth_claims):
        return None
    old_oauth_jti = str(oauth_claims["jti"])
    _REVOKED_AGENT_CREDENTIAL_JTIS.add(old_oauth_jti)
    oauth_jti_ttl = max(86_400, int(oauth_claims.get("exp") or 0) - int(time.time()) + 86_400)
    redis_jti_revoke(old_oauth_jti, ttl_seconds=oauth_jti_ttl)
    oauth_client = create_self_serve_oauth_client_secret(
        agent_id=str(oauth_claims["agent_id"]),
        tenant_id=str(oauth_claims["tenant_id"]),
        scopes=[str(scope) for scope in oauth_claims.get("scopes") or []],
        ttl_seconds=ttl_seconds,
    )
    return {
        "status": "rotated",
        "credential_type": "oauth_client_secret",
        "agent_id": oauth_claims["agent_id"],
        "tenant_id": oauth_claims["tenant_id"],
        "revoked_jti": oauth_claims["jti"],
        "oauth_client": oauth_client,
        "next_actions": [
            "Store oauth_client.client_secret securely; it is returned once.",
            "Exchange the new client secret at /agents/oauth/token before protected calls.",
        ],
    }


def require_self_serve_agent_tenant_access(request: Request, tenant_id: str | None) -> None:
    if not tenant_id:
        return
    expected = api_key_configured()
    provided = credential_from_request(request)
    if not expected or not provided:
        return
    claims = _api_key_claims(provided, expected)
    if claims is not None:
        _require_agent_tenant(claims, tenant_id)


def _require_agent_scope(claims: dict[str, Any], required_scope: str) -> None:
    scopes = {str(scope) for scope in (claims.get("scopes") or [])}
    if "*" in scopes or required_scope in scopes:
        return
    raise HTTPException(
        status_code=403,
        detail=agent_error_detail(
            code="insufficient_agent_scope",
            message=f"Agent key is missing required scope: {required_scope}.",
            retryable=False,
            next_action="Request a new agent credential with the required scope.",
            required_scope=required_scope,
            scopes=sorted(scopes),
        ),
    )


def _require_agent_tenant(claims: dict[str, Any], tenant_id: str) -> None:
    if claims.get("type") == "operator":
        return
    token = str(tenant_id or "").strip()
    claim_tenant = str(claims.get("tenant_id") or "").strip()
    if not token or not claim_tenant or token == claim_tenant:
        return
    raise HTTPException(
        status_code=403,
        detail=agent_error_detail(
            code="agent_tenant_mismatch",
            message="Agent key is not authorized for the requested tenant_id.",
            retryable=False,
            next_action="Use the tenant_id returned by /agents/onboarding or request a new credential.",
            tenant_id=token,
            credential_tenant_id=claim_tenant,
        ),
    )


def _agent_key_signing_secret() -> str:
    return (
        os.getenv("GRANTFLOW_AGENT_KEY_SIGNING_SECRET")
        or os.getenv("GRANTFLOW_API_KEY")
        or os.getenv("API_KEY")
        or "grantflow-dev-agent-key-signing-secret"
    )


def _credential_jti_revoked(payload: dict[str, Any]) -> bool:
    jti = payload.get("jti")
    if not jti:
        return False
    jti_str = str(jti)
    # Check in-process denylist first (fastest path), then Redis.
    if jti_str in _REVOKED_AGENT_CREDENTIAL_JTIS:
        return True
    return redis_jti_is_revoked(jti_str)


def _self_serve_agent_api_key_claims_unchecked(provided: str) -> dict[str, Any] | None:
    if not provided.startswith("gf_agent_"):
        return None
    try:
        token = provided.removeprefix("gf_agent_")
        encoded_payload, signature = token.split(".", 1)
        expected_signature = _sign_agent_key(encoded_payload, _agent_key_signing_secret())
        if not secrets.compare_digest(signature, expected_signature):
            return None
        payload = _decode_b64url_json(encoded_payload)
        if int(payload.get("exp") or 0) < int(time.time()):
            return None
        payload["type"] = "self_serve_agent"
        return payload
    except Exception:
        return None


def _self_serve_oauth_client_secret_claims_unchecked(provided: str) -> dict[str, Any] | None:
    if not provided.startswith("gf_oauth_client_"):
        return None
    try:
        token = provided.removeprefix("gf_oauth_client_")
        encoded_payload, signature = token.split(".", 1)
        expected_signature = _sign_agent_key(encoded_payload, _agent_key_signing_secret())
        if not secrets.compare_digest(signature, expected_signature):
            return None
        payload = _decode_b64url_json(encoded_payload)
        if payload.get("credential_type") != "oauth_client_secret":
            return None
        if int(payload.get("exp") or 0) < int(time.time()):
            return None
        return payload
    except Exception:
        return None


def credential_from_request(request: Request) -> str | None:
    api_key = request.headers.get("x-api-key")
    if api_key:
        return api_key
    authorization = request.headers.get("authorization") or ""
    scheme, _, value = authorization.partition(" ")
    if scheme.lower() == "bearer" and value.strip():
        return value.strip()
    return None


def _b64url_json(payload: dict[str, Any]) -> str:
    raw = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def _decode_b64url_json(value: str) -> dict[str, Any]:
    padded = value + ("=" * (-len(value) % 4))
    decoded = base64.urlsafe_b64decode(padded.encode("ascii"))
    payload = json.loads(decoded.decode("utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("agent key payload must be an object")
    return payload


def _sign_agent_key(encoded_payload: str, secret: str) -> str:
    digest = hmac.new(secret.encode("utf-8"), encoded_payload.encode("ascii"), hashlib.sha256).digest()
    return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")


def install_openapi_api_key_security(app: FastAPI) -> None:
    def custom_openapi():
        if app.openapi_schema:
            return app.openapi_schema

        schema = get_openapi(
            title=app.title,
            version=app.version,
            description=(
                f"{app.description}\n\n"
                "Optional API key auth: set `GRANTFLOW_API_KEY` on the server and send "
                "`X-API-Key` on protected endpoints. Reads can also require auth when "
                "`GRANTFLOW_REQUIRE_AUTH_FOR_READS=true`."
            ),
            routes=app.routes,
        )

        components = schema.setdefault("components", {})
        security_schemes = components.setdefault("securitySchemes", {})
        security_schemes["ApiKeyAuth"] = {"type": "apiKey", "in": "header", "name": "X-API-Key"}
        security_schemes["BearerAuth"] = {"type": "http", "scheme": "bearer"}

        for path, methods in (schema.get("paths") or {}).items():
            for method_name, operation in methods.items():
                if (method_name.lower(), path) in PROTECTED_OPERATIONS and isinstance(operation, dict):
                    operation["security"] = [{"ApiKeyAuth": []}, {"BearerAuth": []}]

        install_agent_openapi_examples(schema)
        app.openapi_schema = schema
        return app.openapi_schema

    app.openapi = custom_openapi
