#!/usr/bin/env python3
"""
Create Bedrock Application Inference Profiles for agent-claw and update SSM.

Run after: aws sso login --profile ai1

Usage:
    python3 scripts/create-inference-profiles.py [--dry-run]

Creates:
    agent-claw-opus    — main agent
    agent-claw-sonnet  — aws_agent + coding_agent
    agent-claw-haiku   — document_agent

Then updates SSM:
    /agent-claw/model-id        → agent-claw-opus ARN
    /agent-claw/subagents       → inline model_id fields replaced with profile ARNs
"""

import argparse
import json
import sys
import boto3
from botocore.exceptions import ClientError

PROFILE = 'ai1'
REGION  = 'us-east-1'

BILLING_TAGS = [
    {'key': 'project', 'value': 'agent-claw'},
    {'key': 'env',     'value': 'prod'},
    {'key': 'owner',   'value': 'daniel'},
]

# Map: profile name → cross-region model ID to copy from
PROFILES_TO_CREATE = {
    'agent-claw-opus':   'us.anthropic.claude-opus-4-6-v1:0',
    'agent-claw-sonnet': 'us.anthropic.claude-sonnet-4-6-20251001-v1:0',
    'agent-claw-haiku':  'us.anthropic.claude-haiku-4-5-20251001-v1:0',
}

# SSM subagent model_id values → which profile ARN to swap in
SUBAGENT_MODEL_MAP = {
    'us.anthropic.claude-sonnet-4-6':              'agent-claw-sonnet',
    'us.anthropic.claude-sonnet-4-6-20251001-v1:0':'agent-claw-sonnet',
    'us.anthropic.claude-haiku-4-5-20251001-v1:0': 'agent-claw-haiku',
}


def get_system_inference_profile_arn(bedrock, model_id: str) -> str:
    """Find the system inference profile ARN for a given cross-region model ID."""
    paginator = bedrock.get_paginator('list_inference_profiles')
    for page in paginator.paginate(typeEquals='SYSTEM_DEFINED'):
        for p in page.get('inferenceProfileSummaries', []):
            if p.get('inferenceProfileId', '') == model_id or \
               any(m.get('modelArn', '').endswith(model_id) for m in p.get('models', [])):
                return p['inferenceProfileArn']
    # Fallback: construct ARN directly (works for cross-region profiles)
    return f'arn:aws:bedrock:{REGION}::foundation-model/{model_id}'


def get_existing_profile(bedrock, name: str) -> dict | None:
    """Return existing application profile by name, or None."""
    paginator = bedrock.get_paginator('list_inference_profiles')
    for page in paginator.paginate(typeEquals='APPLICATION'):
        for p in page.get('inferenceProfileSummaries', []):
            if p.get('inferenceProfileName') == name:
                return p
    return None


def create_or_get_profile(bedrock, name: str, model_id: str, dry_run: bool) -> str:
    """Create application inference profile (idempotent). Returns ARN."""
    existing = get_existing_profile(bedrock, name)
    if existing:
        arn = existing['inferenceProfileArn']
        print(f'  [exists] {name} → {arn}')
        return arn

    source_arn = get_system_inference_profile_arn(bedrock, model_id)
    print(f'  [create] {name}')
    print(f'           source: {source_arn}')

    if dry_run:
        print(f'           [dry-run] skipping create')
        return f'arn:aws:bedrock:{REGION}:{{}account}}:application-inference-profile/{name}-DRY-RUN'

    resp = bedrock.create_inference_profile(
        inferenceProfileName=name,
        description=f'agent-claw {name.split("-")[-1]} model with billing tags',
        modelSource={'copyFrom': source_arn},
        tags=BILLING_TAGS,
    )
    arn = resp['inferenceProfileArn']
    print(f'           → {arn}')
    return arn


def update_ssm(ssm, param: str, value: str, dry_run: bool):
    print(f'  [ssm] {param} = {value[:80]}...' if len(value) > 80 else f'  [ssm] {param} = {value}')
    if not dry_run:
        ssm.put_parameter(Name=param, Value=value, Type='String', Overwrite=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()

    session = boto3.Session(profile_name=PROFILE, region_name=REGION)
    bedrock = session.client('bedrock')
    ssm     = session.client('ssm')

    print('=== Creating inference profiles ===')
    arns = {}
    for name, model_id in PROFILES_TO_CREATE.items():
        arns[name] = create_or_get_profile(bedrock, name, model_id, args.dry_run)

    print('\n=== Updating SSM ===')

    # Main agent model
    update_ssm(ssm, '/agent-claw/model-id', arns['agent-claw-opus'], args.dry_run)

    # Subagents JSON — swap model_id fields
    try:
        resp = ssm.get_parameter(Name='/agent-claw/subagents')
        defs = json.loads(resp['Parameter']['Value'])
    except ClientError as e:
        print(f'  [error] Could not read /agent-claw/subagents: {e}')
        sys.exit(1)

    changed = False
    for agent in defs:
        mid = agent.get('model_id', '')
        profile_name = SUBAGENT_MODEL_MAP.get(mid)
        if profile_name:
            new_arn = arns[profile_name]
            print(f'  [subagent] {agent["name"]}: {mid} → {profile_name}')
            agent['model_id'] = new_arn
            changed = True
        else:
            print(f'  [subagent] {agent["name"]}: {mid} (no mapping, left as-is)')

    if changed:
        update_ssm(ssm, '/agent-claw/subagents', json.dumps(defs, indent=2), args.dry_run)

    print('\n=== Done ===')
    print('Profiles created and SSM updated.')
    print('Redeploy not required — agent reads model IDs from SSM at startup.')
    if args.dry_run:
        print('\n[dry-run mode — no AWS changes were made]')


if __name__ == '__main__':
    main()