/**
 * bedrock-login.ts — Makes Amazon Bedrock available via /login
 *
 * Registers an OAuth-compatible provider for amazon-bedrock that:
 * 1. Prompts for AWS credentials via /login UI
 * 2. Stores them in auth.json (0600 permissions, same as other providers)
 * 3. Injects them as process.env vars on extension load and session start
 *
 * Usage: /login amazon-bedrock
 * Supported credential formats:
 *   - Bedrock API key (bearer token): full string including any colons
 *     e.g. BedrockAPIKey-<user>-<account>:<secret>
 *   - IAM user keys: ACCESS_KEY_ID:SECRET_ACCESS_KEY (AKIA... prefix)
 *
 * Key format detection:
 *   - Starts with "BedrockAPIKey-" → bearer token (entire string)
 *   - Starts with "AKIA" or no colon → bearer token / IAM access key
 *   - Otherwise with colon → IAM ACCESS_KEY_ID:SECRET_ACCESS_KEY split
 *
 * The Pi runtime stores credentials as an internal binary envelope in
 * auth.json. loadStoredCredentials() handles both the legacy readable
 * {type,access,refresh} format AND the Pi binary-blob format (where Pi
 * decodes it and calls getApiKey/refreshToken with plain credential objects).
 */

import type { ExtensionAPI } from '@helios-agent/cli';
import { join } from 'path';
import { homedir } from 'os';
import { readFileSync, existsSync } from 'fs';
const { heliosPath } = require('./lib/helios-root');

const AUTH_PATH = process.env.HELIOS_AGENT_PACKAGE
  ? join(process.env.HELIOS_AGENT_PACKAGE, 'auth.json')
  : existsSync(join(homedir(), '.pi', 'agent', 'auth.json'))
    ? join(homedir(), '.pi', 'agent', 'auth.json')
    : heliosPath('auth.json');

/** Detect whether a credential string is a Bedrock API key (bearer token). */
function isBedrockApiKey(s: string): boolean {
  // Bedrock service-specific credential format: "BedrockAPIKey-<user>-<account>:<secret>"
  // The entire string (including the colon) is the bearer token.
  return s.startsWith('BedrockAPIKey-');
}

function injectEnvVars(access: string, refresh: string): void {
  // If the access string itself is a Bedrock API key (contains the full bearer token),
  // OR if refresh sentinel is 'bearer-token', inject as bearer token.
  if (refresh === 'bearer-token' || isBedrockApiKey(access)) {
    process.env.AWS_BEARER_TOKEN_BEDROCK = wrapAsBearerToken(access);
  } else {
    process.env.AWS_ACCESS_KEY_ID = access;
    process.env.AWS_SECRET_ACCESS_KEY = refresh;
  }
}

/**
 * Wrap a raw BedrockAPIKey string in the Pi binary envelope format that
 * AWS_BEARER_TOKEN_BEDROCK actually requires.
 *
 * AWS rejects the raw "BedrockAPIKey-..." string with 403 "Invalid API Key format".
 * The accepted form is: base64([0x00, 0x14, 0x8a] + utf8(rawKey))
 *
 * If the string is already wrapped (valid base64 that decodes to the 3-byte prefix)
 * it is returned as-is.
 */
function wrapAsBearerToken(rawKey: string): string {
  if (!rawKey.startsWith('BedrockAPIKey-')) return rawKey;
  // Check if already wrapped (idempotent)
  try {
    const decoded = Buffer.from(rawKey, 'base64');
    if (decoded[0] === 0x00 && decoded[1] === 0x14 && decoded[2] === 0x8a) return rawKey;
  } catch (_) { /* not base64, continue */ }
  const keyBytes = Buffer.from(rawKey, 'utf-8');
  const prefix = Buffer.from([0x00, 0x14, 0x8a]);
  return Buffer.concat([prefix, keyBytes]).toString('base64');
}

function loadStoredCredentials(): void {
  // P2-13 NX fix: cache result in globalThis with 5-min TTL so repeated calls
  // (multiple worker processes loading this extension) skip the readFileSync.
  // Credentials don't change during a session — caching is safe.
  const CREDS_KEY = '__helios_bedrock_creds_ts';
  const g = globalThis as any;
  const TTL = 5 * 60 * 1000;
  if (g[CREDS_KEY] && Date.now() - g[CREDS_KEY] < TTL) return; // cache hit
  g[CREDS_KEY] = Date.now(); // mark as attempted (even on failure)

  try {
    const auth = JSON.parse(readFileSync(AUTH_PATH, 'utf-8'));
    const cred = auth['amazon-bedrock'];
    if (!cred) return;

    if (typeof cred === 'object' && cred.type === 'api_key' && cred.key) {
      // New pi format: {"type":"api_key","key":"BedrockAPIKey-..."} — wrap in Pi binary envelope
      const key = String(cred.key).trim();
      if (key) process.env.AWS_BEARER_TOKEN_BEDROCK = wrapAsBearerToken(key);
    } else if (typeof cred === 'object' && cred.type === 'oauth' && cred.access && cred.refresh) {
      // Readable {type, access, refresh} format (written by older login handler)
      injectEnvVars(cred.access, cred.refresh);
    } else if (typeof cred === 'string') {
      // Pi runtime stores credentials as a base64-encoded binary envelope.
      // The framework expects the envelope format (not the decoded raw key).
      // Use the stored value directly; decode only to detect IAM key pairs.
      try {
        const raw = Buffer.from(cred, 'base64');
        // Pi envelope: [uint16 length (2 bytes)] [type byte (1 byte)] [UTF-8 payload]
        const payload = raw.slice(3).toString('utf-8');
        if (payload && isBedrockApiKey(payload)) {
          // Pass the original envelope — framework rejects the raw decoded key
          process.env.AWS_BEARER_TOKEN_BEDROCK = cred;
        } else if (payload && payload.includes(':')) {
          const ci = payload.indexOf(':');
          injectEnvVars(payload.substring(0, ci), payload.substring(ci + 1));
        } else if (payload && payload.length > 20) {
          // Pass the original envelope for any other opaque credential
          process.env.AWS_BEARER_TOKEN_BEDROCK = cred;
        }
      } catch (_) {
        // Not base64 — try using as raw bearer token (wrap it first)
        if (isBedrockApiKey(cred)) process.env.AWS_BEARER_TOKEN_BEDROCK = wrapAsBearerToken(cred);
      }
    }
  } catch (e) { process.stderr.write(`[bedrock-login.ts] loadStoredCredentials failed: ${String(e)}\n`); }
}

export default function bedrockLogin(pi: ExtensionAPI): void {
  // Inject stored credentials immediately on extension load
  // so model discovery sees bedrock as authenticated
  loadStoredCredentials();

  pi.registerProvider('amazon-bedrock', {
    oauth: {
      name: 'Amazon Bedrock (AWS Credentials)',

      async login(callbacks: any) {
        // onAuth shows the AWS console URL for users to find/create credentials
        callbacks.onAuth({
          url: 'https://console.aws.amazon.com/iam/home#/security_credentials',
          instructions: [
            'For a Bedrock API key (recommended): paste the full key string,',
            'e.g. BedrockAPIKey-<user>-<account>:<secret>',
            '',
            'For IAM user keys: paste ACCESS_KEY_ID:SECRET_ACCESS_KEY',
          ].join('\n'),
        });

        // onPrompt prompts the user for their credentials
        const input = await callbacks.onPrompt({
          message: 'Paste your Bedrock API key OR ACCESS_KEY_ID:SECRET_ACCESS_KEY',
          placeholder: 'BedrockAPIKey-... or AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI...',
        });

        if (!input) {
          throw new Error('No credentials provided. Run /login amazon-bedrock to try again.');
        }

        const trimmed = input.trim();

        // Bedrock API key: entire string is the bearer token, even if it contains a colon
        if (isBedrockApiKey(trimmed)) {
          process.env.AWS_BEARER_TOKEN_BEDROCK = wrapAsBearerToken(trimmed);
          return {
            access: trimmed,
            refresh: 'bearer-token',
            expires: Date.now() + 10 * 365 * 24 * 60 * 60 * 1000,
          };
        }

        const colonIndex = trimmed.indexOf(':');
        if (colonIndex === -1) {
          // No colon — treat as a bearer token if long enough
          if (trimmed.length > 50) {
            process.env.AWS_BEARER_TOKEN_BEDROCK = trimmed;
            return {
              access: trimmed,
              refresh: 'bearer-token',
              expires: Date.now() + 10 * 365 * 24 * 60 * 60 * 1000,
            };
          }
          throw new Error(
            'Invalid format. Expected a Bedrock API key (BedrockAPIKey-...) or\n' +
            'IAM keys as ACCESS_KEY_ID:SECRET_ACCESS_KEY'
          );
        }

        const accessKeyId = trimmed.substring(0, colonIndex).trim();
        const secretAccessKey = trimmed.substring(colonIndex + 1).trim();

        if (!accessKeyId || !secretAccessKey) {
          throw new Error('Both access key ID and secret access key are required.');
        }

        // Inject immediately so bedrock works this session
        process.env.AWS_ACCESS_KEY_ID = accessKeyId;
        process.env.AWS_SECRET_ACCESS_KEY = secretAccessKey;

        return {
          access: accessKeyId,
          refresh: secretAccessKey,
          expires: Date.now() + 10 * 365 * 24 * 60 * 60 * 1000,
        };
      },

      async refreshToken(credentials: any) {
        // IAM keys don't expire — re-inject env vars and pass through
        injectEnvVars(credentials.access, credentials.refresh);
        return {
          ...credentials,
          expires: Date.now() + 10 * 365 * 24 * 60 * 60 * 1000,
        };
      },

      getApiKey(credentials: any) {
        // Inject env vars every time Pi resolves the API key
        injectEnvVars(credentials.access, credentials.refresh);
        return '<authenticated>';
      },
    },
  });

  // Safety net: also inject on session_start
  pi.on('session_start', async () => {
    loadStoredCredentials();
  });
}