import { db, member, ssoProvider } from '@sim/db'
import { createLogger } from '@sim/logger'
import { and, eq } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { ssoRegistrationContract } from '@/lib/api/contracts/auth'
import { getValidationErrorMessage, parseRequest } from '@/lib/api/server'
import { auth, getSession } from '@/lib/auth'
import { hasSSOAccess } from '@/lib/billing'
import { env } from '@/lib/core/config/env'
import {
  secureFetchWithPinnedIP,
  validateUrlWithDNS,
} from '@/lib/core/security/input-validation.server'
import { REDACTED_MARKER } from '@/lib/core/security/redaction'
import { getBaseUrl } from '@/lib/core/utils/urls'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'

const logger = createLogger('SSORegisterRoute')

export const POST = withRouteHandler(async (request: NextRequest) => {
  try {
    if (!env.SSO_ENABLED) {
      return NextResponse.json({ error: 'SSO is not enabled' }, { status: 400 })
    }

    const session = await getSession()
    if (!session?.user?.id) {
      return NextResponse.json({ error: 'Authentication required' }, { status: 401 })
    }

    const hasAccess = await hasSSOAccess(session.user.id)
    if (!hasAccess) {
      return NextResponse.json({ error: 'SSO requires an Enterprise plan' }, { status: 403 })
    }

    const parsed = await parseRequest(
      ssoRegistrationContract,
      request,
      {},
      {
        validationErrorResponse: (error) => {
          logger.warn('Invalid SSO registration request', { errors: error.issues })
          return NextResponse.json(
            { error: getValidationErrorMessage(error, 'Validation failed') },
            { status: 400 }
          )
        },
      }
    )
    if (!parsed.success) return parsed.response

    const body = parsed.data.body
    const { providerId, issuer, domain, providerType, mapping, orgId } = body

    if (orgId) {
      const [membership] = await db
        .select({ organizationId: member.organizationId, role: member.role })
        .from(member)
        .where(and(eq(member.userId, session.user.id), eq(member.organizationId, orgId)))
        .limit(1)
      if (!membership) {
        return NextResponse.json({ error: 'Forbidden' }, { status: 403 })
      }
      if (membership.role !== 'owner' && membership.role !== 'admin') {
        return NextResponse.json({ error: 'Forbidden' }, { status: 403 })
      }
    }

    const headers: Record<string, string> = {}
    request.headers.forEach((value, key) => {
      headers[key] = value
    })

    const providerConfig: any = {
      providerId,
      issuer,
      domain,
      mapping,
      ...(orgId ? { organizationId: orgId } : {}),
    }

    if (providerType === 'oidc') {
      const {
        clientId,
        clientSecret: rawClientSecret,
        scopes,
        pkce,
        authorizationEndpoint,
        tokenEndpoint,
        userInfoEndpoint,
        jwksEndpoint,
      } = body

      let clientSecret = rawClientSecret
      if (rawClientSecret === REDACTED_MARKER) {
        const ownerClause = orgId
          ? and(eq(ssoProvider.providerId, providerId), eq(ssoProvider.organizationId, orgId))
          : and(eq(ssoProvider.providerId, providerId), eq(ssoProvider.userId, session.user.id))
        const [existing] = await db
          .select({ oidcConfig: ssoProvider.oidcConfig })
          .from(ssoProvider)
          .where(ownerClause)
          .limit(1)
        if (!existing?.oidcConfig) {
          return NextResponse.json(
            { error: 'Cannot update: existing provider not found. Re-enter your client secret.' },
            { status: 400 }
          )
        }
        try {
          clientSecret = JSON.parse(existing.oidcConfig).clientSecret
        } catch {
          return NextResponse.json(
            {
              error: 'Cannot update: failed to read existing secret. Re-enter your client secret.',
            },
            { status: 400 }
          )
        }
      }

      const oidcConfig: any = {
        clientId,
        clientSecret,
        scopes: Array.isArray(scopes)
          ? scopes.filter((s: string) => s !== 'offline_access')
          : ['openid', 'profile', 'email'].filter((s: string) => s !== 'offline_access'),
        pkce: pkce ?? true,
      }

      oidcConfig.authorizationEndpoint = authorizationEndpoint
      oidcConfig.tokenEndpoint = tokenEndpoint
      oidcConfig.userInfoEndpoint = userInfoEndpoint
      oidcConfig.jwksEndpoint = jwksEndpoint

      const userProvidedEndpoints: Record<string, string | undefined> = {
        authorizationEndpoint,
        tokenEndpoint,
        userInfoEndpoint,
        jwksEndpoint,
      }

      for (const [name, endpointUrl] of Object.entries(userProvidedEndpoints)) {
        if (endpointUrl) {
          const endpointValidation = await validateUrlWithDNS(endpointUrl, `OIDC ${name}`)
          if (!endpointValidation.isValid) {
            logger.warn('Explicitly provided OIDC endpoint failed SSRF validation', {
              endpoint: name,
              url: endpointUrl,
              error: endpointValidation.error,
            })
            return NextResponse.json(
              {
                error: `OIDC ${name} failed security validation: ${endpointValidation.error}`,
              },
              { status: 400 }
            )
          }
        }
      }

      const needsDiscovery =
        !oidcConfig.authorizationEndpoint || !oidcConfig.tokenEndpoint || !oidcConfig.jwksEndpoint

      if (needsDiscovery) {
        const discoveryUrl = `${issuer.replace(/\/$/, '')}/.well-known/openid-configuration`
        try {
          logger.info('Fetching OIDC discovery document for missing endpoints', {
            discoveryUrl,
            hasAuthEndpoint: !!oidcConfig.authorizationEndpoint,
            hasTokenEndpoint: !!oidcConfig.tokenEndpoint,
            hasJwksEndpoint: !!oidcConfig.jwksEndpoint,
          })

          const urlValidation = await validateUrlWithDNS(discoveryUrl, 'OIDC discovery URL')
          if (!urlValidation.isValid || !urlValidation.resolvedIP) {
            logger.warn('OIDC discovery URL failed SSRF validation', {
              discoveryUrl,
              error: urlValidation.error,
            })
            return NextResponse.json(
              { error: urlValidation.error ?? 'SSRF validation failed' },
              { status: 400 }
            )
          }

          const discoveryResponse = await secureFetchWithPinnedIP(
            discoveryUrl,
            urlValidation.resolvedIP,
            {
              headers: { Accept: 'application/json' },
            }
          )

          if (!discoveryResponse.ok) {
            logger.error('Failed to fetch OIDC discovery document', {
              status: discoveryResponse.status,
            })
            return NextResponse.json(
              {
                error:
                  'Failed to fetch OIDC discovery document. Provide all endpoints explicitly or verify the issuer URL.',
              },
              { status: 400 }
            )
          }

          const discovery = (await discoveryResponse.json()) as Record<string, unknown>

          const discoveredEndpoints: Record<string, unknown> = {
            authorization_endpoint: discovery.authorization_endpoint,
            token_endpoint: discovery.token_endpoint,
            userinfo_endpoint: discovery.userinfo_endpoint,
            jwks_uri: discovery.jwks_uri,
          }

          for (const [key, value] of Object.entries(discoveredEndpoints)) {
            if (typeof value === 'string') {
              const endpointValidation = await validateUrlWithDNS(value, `OIDC ${key}`)
              if (!endpointValidation.isValid) {
                logger.warn('OIDC discovered endpoint failed SSRF validation', {
                  endpoint: key,
                  url: value,
                  error: endpointValidation.error,
                })
                return NextResponse.json(
                  {
                    error: `Discovered OIDC ${key} failed security validation: ${endpointValidation.error}`,
                  },
                  { status: 400 }
                )
              }
            }
          }

          oidcConfig.authorizationEndpoint =
            oidcConfig.authorizationEndpoint || discovery.authorization_endpoint
          oidcConfig.tokenEndpoint = oidcConfig.tokenEndpoint || discovery.token_endpoint
          oidcConfig.userInfoEndpoint = oidcConfig.userInfoEndpoint || discovery.userinfo_endpoint
          oidcConfig.jwksEndpoint = oidcConfig.jwksEndpoint || discovery.jwks_uri

          logger.info('Merged OIDC endpoints (user-provided + discovery)', {
            providerId,
            issuer,
            authorizationEndpoint: oidcConfig.authorizationEndpoint,
            tokenEndpoint: oidcConfig.tokenEndpoint,
            userInfoEndpoint: oidcConfig.userInfoEndpoint,
            jwksEndpoint: oidcConfig.jwksEndpoint,
          })
        } catch (error) {
          logger.error('Error fetching OIDC discovery document', {
            error: error instanceof Error ? error.message : 'Unknown error',
            discoveryUrl,
          })
          return NextResponse.json(
            {
              error:
                'Failed to fetch OIDC discovery document. Please verify the issuer URL is correct or provide all endpoints explicitly.',
            },
            { status: 400 }
          )
        }
      } else {
        logger.info('Using explicitly provided OIDC endpoints (all present)', {
          providerId,
          issuer,
          authorizationEndpoint: oidcConfig.authorizationEndpoint,
          tokenEndpoint: oidcConfig.tokenEndpoint,
          userInfoEndpoint: oidcConfig.userInfoEndpoint,
          jwksEndpoint: oidcConfig.jwksEndpoint,
        })
      }

      if (
        !oidcConfig.authorizationEndpoint ||
        !oidcConfig.tokenEndpoint ||
        !oidcConfig.jwksEndpoint
      ) {
        const missing: string[] = []
        if (!oidcConfig.authorizationEndpoint) missing.push('authorizationEndpoint')
        if (!oidcConfig.tokenEndpoint) missing.push('tokenEndpoint')
        if (!oidcConfig.jwksEndpoint) missing.push('jwksEndpoint')

        logger.error('Missing required OIDC endpoints after discovery merge', {
          missing,
          authorizationEndpoint: oidcConfig.authorizationEndpoint,
          tokenEndpoint: oidcConfig.tokenEndpoint,
          jwksEndpoint: oidcConfig.jwksEndpoint,
        })
        return NextResponse.json(
          {
            error: `Missing required OIDC endpoints: ${missing.join(', ')}. Please provide these explicitly or verify the issuer supports OIDC discovery.`,
          },
          { status: 400 }
        )
      }

      providerConfig.oidcConfig = oidcConfig
    } else if (providerType === 'saml') {
      const {
        entryPoint,
        cert,
        callbackUrl,
        audience,
        wantAssertionsSigned,
        signatureAlgorithm,
        digestAlgorithm,
        identifierFormat,
        idpMetadata,
      } = body

      const computedCallbackUrl =
        callbackUrl || `${getBaseUrl()}/api/auth/sso/saml2/callback/${providerId}`

      const escapeXml = (str: string) =>
        str.replace(/[<>&"']/g, (c) => {
          switch (c) {
            case '<':
              return '&lt;'
            case '>':
              return '&gt;'
            case '&':
              return '&amp;'
            case '"':
              return '&quot;'
            case "'":
              return '&apos;'
            default:
              return c
          }
        })

      const spMetadataXml = `<?xml version="1.0" encoding="UTF-8"?>
<md:EntityDescriptor xmlns:md="urn:oasis:names:tc:SAML:2.0:metadata" entityID="${escapeXml(getBaseUrl())}">
  <md:SPSSODescriptor AuthnRequestsSigned="false" WantAssertionsSigned="false" protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
    <md:AssertionConsumerService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" Location="${escapeXml(computedCallbackUrl)}" index="1"/>
  </md:SPSSODescriptor>
</md:EntityDescriptor>`

      const certBase64 = cert
        .replace(/-----BEGIN CERTIFICATE-----/g, '')
        .replace(/-----END CERTIFICATE-----/g, '')
        .replace(/\s/g, '')

      const computedIdpMetadataXml =
        idpMetadata ||
        `<?xml version="1.0"?>
<EntityDescriptor xmlns="urn:oasis:names:tc:SAML:2.0:metadata" entityID="${escapeXml(issuer)}">
  <IDPSSODescriptor WantAuthnRequestsSigned="false" protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
    <KeyDescriptor use="signing">
      <ds:KeyInfo xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
        <ds:X509Data>
          <ds:X509Certificate>${certBase64}</ds:X509Certificate>
        </ds:X509Data>
      </ds:KeyInfo>
    </KeyDescriptor>
    <SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" Location="${escapeXml(entryPoint)}"/>
    <SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" Location="${escapeXml(entryPoint)}"/>
  </IDPSSODescriptor>
</EntityDescriptor>`

      const samlConfig: any = {
        entryPoint,
        cert,
        callbackUrl: computedCallbackUrl,
        spMetadata: {
          metadata: spMetadataXml,
        },
        idpMetadata: {
          metadata: computedIdpMetadataXml,
        },
      }

      if (audience) samlConfig.audience = audience
      if (wantAssertionsSigned !== undefined) samlConfig.wantAssertionsSigned = wantAssertionsSigned
      if (signatureAlgorithm) samlConfig.signatureAlgorithm = signatureAlgorithm
      if (digestAlgorithm) samlConfig.digestAlgorithm = digestAlgorithm
      if (identifierFormat) samlConfig.identifierFormat = identifierFormat

      providerConfig.samlConfig = samlConfig
    }

    logger.info('Calling Better Auth registerSSOProvider with config:', {
      providerId: providerConfig.providerId,
      domain: providerConfig.domain,
      hasOidcConfig: !!providerConfig.oidcConfig,
      hasSamlConfig: !!providerConfig.samlConfig,
      samlConfigKeys: providerConfig.samlConfig ? Object.keys(providerConfig.samlConfig) : [],
      fullConfig: JSON.stringify(
        {
          ...providerConfig,
          oidcConfig: providerConfig.oidcConfig
            ? {
                ...providerConfig.oidcConfig,
                clientSecret: REDACTED_MARKER,
              }
            : undefined,
          samlConfig: providerConfig.samlConfig
            ? {
                ...providerConfig.samlConfig,
                cert: REDACTED_MARKER,
              }
            : undefined,
        },
        null,
        2
      ),
    })

    const registration = await auth.api.registerSSOProvider({
      body: providerConfig,
      headers,
    })

    logger.info('SSO provider registered successfully', {
      providerId,
      providerType,
      domain,
    })

    return NextResponse.json({
      success: true,
      providerId: registration.providerId,
      providerType,
      message: `${providerType.toUpperCase()} provider registered successfully`,
    })
  } catch (error) {
    logger.error('Failed to register SSO provider', {
      error,
      errorMessage: error instanceof Error ? error.message : 'Unknown error',
      errorStack: error instanceof Error ? error.stack : undefined,
      errorDetails: JSON.stringify(error),
    })

    return NextResponse.json(
      {
        error: 'Failed to register SSO provider',
        details: error instanceof Error ? error.message : 'Unknown error',
      },
      { status: 500 }
    )
  }
})
