【AWS】Kiroのコストの管理コストはいくらかかるのか？

"""
Kiro コスト管理自動レポート生成 Lambda関数
- S3上のKiro日次レポート（CLI/IDE）を読み込み、ユーザ別に集計
- プラン料金マスタと突き合わせ、利用率・超過・余剰を算出
- 合計料金をプランマスタから計算（ユーザ数×月額 + 超過分）
- HTMLレポートを生成し、S3保存 + SES送信（管理者 + Teamsチャネル）
"""

import os
import csv
import io
import json
import boto3
from datetime import date, timedelta, datetime
from collections import defaultdict


# ─── プラン料金マスタ（コード内定義） ───
# 出典: https://kiro.dev/pricing/ (2026年8月時点)
PLAN_MASTER = {
    'FREE':     {'plan_limit': 50,    'monthly_cost_usd': 0,   'overage_cost_per_credit_usd': 0},
    'PRO':      {'plan_limit': 1000,  'monthly_cost_usd': 20,  'overage_cost_per_credit_usd': 0.04},
    'PRO_PLUS': {'plan_limit': 2000,  'monthly_cost_usd': 40,  'overage_cost_per_credit_usd': 0.04},
    'PRO_MAX':  {'plan_limit': 5000,  'monthly_cost_usd': 100, 'overage_cost_per_credit_usd': 0.04},
    'POWER':    {'plan_limit': 10000, 'monthly_cost_usd': 200, 'overage_cost_per_credit_usd': 0.04},
}

# ─── 損益分岐点（固定値、事前計算） ───
# 各プランにおいて「上位プランに切り替えた方が安くなる利用クレジット数」
# 計算式: plan_limit + (上位プラン月額 - 現プラン月額) / overage_cost_per_credit
# 例: PRO → PRO_PLUS: 1000 + (40-20)/0.04 = 1,500 credits
BREAKEVEN_THRESHOLDS = {
    'PRO':      {'next_tier': 'PRO_PLUS', 'threshold': 1000 + (40 - 20) / 0.04},    # 1,500
    'PRO_PLUS': {'next_tier': 'PRO_MAX',  'threshold': 2000 + (100 - 40) / 0.04},   # 3,500
    'PRO_MAX':  {'next_tier': 'POWER',    'threshold': 5000 + (200 - 100) / 0.04},   # 7,500
}

SOURCE_BUCKET = os.environ['SOURCE_BUCKET']
SOURCE_PREFIX = os.environ['SOURCE_PREFIX']
ACCOUNT_ID = os.environ['ACCOUNT_ID']
OUTPUT_BUCKET = os.environ.get('OUTPUT_BUCKET', 'tkn-kiro-cost-report').strip()
SES_SENDER = os.environ['SES_SENDER']
SES_RECIPIENTS = os.environ['SES_RECIPIENTS'].split(',')
TEAMS_CHANNEL_EMAIL = os.environ.get('TEAMS_CHANNEL_EMAIL', '')
SURPLUS_THRESHOLD_PERCENT = float(os.environ.get('SURPLUS_THRESHOLD_PERCENT', '30'))

s3 = boto3.client('s3')
ses = boto3.client('ses', region_name='ap-northeast-1')  # SESは東京リージョンを使用


def lambda_handler(event, context):
    report_type = event.get('report_type', 'mid_month')  # 'mid_month' or 'end_month'
    today = date.today()

    # ─── 対象期間の決定 ───
    start_date = today.replace(day=1)
    end_date = today - timedelta(days=1)

    # ─── プラン料金マスタ（コード内定義済み） ───
    plan_master = PLAN_MASTER

    # ─── 日次CSVの読み込み・集計 ───
    user_data = aggregate_daily_reports(start_date, end_date)

    # ─── ユーザごとの計算 ───
    user_reports = calculate_user_metrics(user_data, plan_master)

    # ─── 合計料金の計算（プランマスタから算出） ───
    kiro_cost = calculate_total_cost(user_reports)

    # ─── 月末レポートの追加判定 ───
    breakeven_users = []
    surplus_users = []
    if report_type == 'end_month':
        breakeven_users = find_breakeven_exceeded_users(user_reports, plan_master)
        surplus_users = find_surplus_users(user_reports)

    # ─── HTMLレポート生成 ───
    html = generate_html_report(
        report_type=report_type,
        user_reports=user_reports,
        kiro_cost=kiro_cost,
        start_date=start_date,
        end_date=end_date,
        breakeven_users=breakeven_users,
        surplus_users=surplus_users,
    )

    # ─── S3保存 ───
    report_key = f"reports/{today.strftime('%Y/%m')}/{report_type}_{today.isoformat()}.html"
    s3.put_object(
        Bucket=OUTPUT_BUCKET,
        Key=report_key,
        Body=html.encode('utf-8'),
        ContentType='text/html; charset=utf-8',
    )

    # ─── メール送信：管理者宛（HTMLフルレポート） ───
    subject = build_subject(report_type, today)
    if SES_RECIPIENTS:
        send_html_email(subject, html, SES_RECIPIENTS)

    # ─── 署名付きURL生成 + Teams通知 ───
    if TEAMS_CHANNEL_EMAIL:
        report_url = s3.generate_presigned_url(
            'get_object',
            Params={'Bucket': OUTPUT_BUCKET, 'Key': report_key},
            ExpiresIn=21600,  # 6時間（Lambda IAMロールの一時認証情報の上限）
        )
        teams_summary = build_teams_summary(
            report_type, user_reports, kiro_cost, start_date, end_date,
            breakeven_users, surplus_users, report_url
        )
        send_text_email(subject, teams_summary, [TEAMS_CHANNEL_EMAIL])

    return {
        'statusCode': 200,
        'body': json.dumps({
            'report_type': report_type,
            'users_count': len(user_reports),
            'kiro_cost_usd': kiro_cost,
            's3_key': report_key,
        })
    }


# ═══════════════════════════════════════════════════════════════
# データ取得・集計
# ═══════════════════════════════════════════════════════════════

def aggregate_daily_reports(start_date, end_date):
    """
    期間内の全日付のCSV（CLI + IDE）を読み込み、UserId単位で合算する。
    日次データは増分値のため、全日分を加算する。
    """
    user_data = defaultdict(lambda: {
        'Credits_Used': 0.0,
        'Overage_Credits_Used': 0.0,
        'Total_Messages': 0,
        'User_Email': '',
        'Subscription_Tier': '',
        'Client_Types': set(),
    })

    current = start_date
    while current <= end_date:
        date_str = current.strftime('%Y%m%d')
        date_path = current.strftime('%Y/%m/%d')

        for client_type in ['CLI', 'IDE']:
            key = (
                f"{SOURCE_PREFIX}/{date_path}/00/"
                f"KIRO_{client_type}_{ACCOUNT_ID}_user_report_{date_str}0000.csv"
            )
            try:
                obj = s3.get_object(Bucket=SOURCE_BUCKET, Key=key)
                content = obj['Body'].read().decode('utf-8')
                reader = csv.DictReader(io.StringIO(content))
                for row in reader:
                    uid = row['UserId'].strip('"')
                    user_data[uid]['Credits_Used'] += float(row['Credits_Used'])
                    user_data[uid]['Overage_Credits_Used'] += float(row['Overage_Credits_Used'])
                    user_data[uid]['Total_Messages'] += int(row['Total_Messages'])
                    user_data[uid]['User_Email'] = row['User_Email'].strip('"')
                    user_data[uid]['Subscription_Tier'] = row['Subscription_Tier']
                    user_data[uid]['Client_Types'].add(client_type)
            except s3.exceptions.NoSuchKey:
                # その日のファイルが存在しない場合はスキップ
                pass
            except Exception as e:
                print(f"Warning: Failed to read {key}: {e}")

        current += timedelta(days=1)

    return user_data


def calculate_user_metrics(user_data, plan_master):
    """ユーザごとの利用率・超過・余剰を計算"""
    reports = []
    for uid, data in user_data.items():
        tier = data['Subscription_Tier']
        plan_info = plan_master.get(tier, plan_master.get('PRO'))  # フォールバック

        plan_limit = plan_info['plan_limit']
        credits_used = data['Credits_Used']
        overage_used = data['Overage_Credits_Used']

        usage_rate = (credits_used / plan_limit * 100) if plan_limit > 0 else 0
        surplus = max(0, plan_limit - credits_used)
        overage_cost = overage_used * plan_info['overage_cost_per_credit_usd']

        reports.append({
            'user_id': uid,
            'user_email': data['User_Email'],
            'tier': tier,
            'plan_limit': plan_limit,
            'credits_used': round(credits_used, 2),
            'overage_credits': round(overage_used, 2),
            'overage_cost_usd': round(overage_cost, 2),
            'usage_rate': round(usage_rate, 1),
            'surplus_credits': round(surplus, 2),
            'total_messages': data['Total_Messages'],
            'client_types': ', '.join(sorted(data['Client_Types'])),
        })

    # 利用率の高い順にソート
    reports.sort(key=lambda x: x['usage_rate'], reverse=True)
    return reports


def calculate_total_cost(user_reports):
    """プランマスタからKiro合計料金を算出（ユーザ数×月額 + 超過コスト合計）"""
    total = 0.0
    for user in user_reports:
        tier = user['tier']
        plan_info = PLAN_MASTER.get(tier, PLAN_MASTER.get('PRO'))
        total += plan_info['monthly_cost_usd']
        total += user['overage_cost_usd']
    return round(total, 2)


# ═══════════════════════════════════════════════════════════════
# 月末レポート判定ロジック
# ═══════════════════════════════════════════════════════════════

def find_breakeven_exceeded_users(user_reports, plan_master):
    """
    損益分岐点超過ユーザ：利用クレジットが閾値を超えているユーザ。
    閾値 = plan_limit + (上位プラン月額 - 現プラン月額) / $0.04
    """
    results = []

    for user in user_reports:
        current_tier = user['tier']
        if current_tier not in BREAKEVEN_THRESHOLDS:
            continue

        threshold_info = BREAKEVEN_THRESHOLDS[current_tier]
        threshold = threshold_info['threshold']
        next_tier = threshold_info['next_tier']

        if user['credits_used'] >= threshold:
            current_cost = plan_master[current_tier]['monthly_cost_usd']
            next_cost = plan_master[next_tier]['monthly_cost_usd']
            actual_total = current_cost + user['overage_cost_usd']

            results.append({
                **user,
                'recommended_tier': next_tier,
                'breakeven_threshold': threshold,
                'current_monthly_total': round(actual_total, 2),
                'recommended_monthly_cost': next_cost,
                'estimated_savings': round(actual_total - next_cost, 2),
            })

    return results


def find_surplus_users(user_reports):
    """
    余剰多数ユーザ：利用率が閾値未満で、FREEプラン以外のユーザ。
    推奨アクション = プランダウングレードまたは継続意思確認
    """
    tier_order = ['FREE', 'PRO', 'PRO_PLUS', 'PRO_MAX', 'POWER']
    results = []

    for user in user_reports:
        if user['tier'] == 'FREE':
            continue
        if user['usage_rate'] < SURPLUS_THRESHOLD_PERCENT:
            # 下位プランの推奨
            current_idx = tier_order.index(user['tier']) if user['tier'] in tier_order else 1
            recommended_tier = tier_order[max(0, current_idx - 1)]

            results.append({
                **user,
                'recommended_tier': recommended_tier,
            })

    return results


# ═══════════════════════════════════════════════════════════════
# HTMLレポート生成
# ═══════════════════════════════════════════════════════════════

def generate_html_report(report_type, user_reports, kiro_cost, start_date, end_date,
                         breakeven_users, surplus_users):
    """HTMLレポートを生成"""
    title = '中間レポート' if report_type == 'mid_month' else '月末レポート'
    period = f"{start_date.strftime('%Y/%m/%d')} ～ {end_date.strftime('%Y/%m/%d')}"

    # 横棒グラフ
    chart_rows = ''
    for u in user_reports:
        bar_width = min(u['usage_rate'], 100)
        bar_color = '#e53e3e' if u['usage_rate'] >= 100 else '#dd6b20' if u['usage_rate'] >= 80 else '#4361ee'
        rate_color = get_rate_color(u['usage_rate'])
        chart_rows += f"""        <div class="bar-row">
            <span class="bar-label">{u['user_email']}</span>
            <div class="bar-bg">
                <div class="bar-fill" style="width:{bar_width}%;background:{bar_color};"></div>
                <div class="bar-text"><span>{u['credits_used']:,.1f} / {u['plan_limit']:,.0f}</span></div>
            </div>
            <span class="bar-rate" style="color:{rate_color};">{u['usage_rate']:.1f}%</span>
        </div>\n"""

    # ユーザ一覧テーブル
    user_rows = ''
    for u in user_reports:
        rate_color = get_rate_color(u['usage_rate'])
        user_rows += f"""
        <tr>
            <td>{u['user_email']}</td>
            <td>{u['tier']}</td>
            <td class="num">{u['credits_used']:,.2f} / {u['plan_limit']:,.0f}</td>
            <td class="num" style="color:{rate_color};font-weight:bold;">{u['usage_rate']:.1f}%</td>
            <td class="num">{u['overage_credits']:,.2f}</td>
            <td class="num">{u['surplus_credits']:,.2f}</td>
            <td class="num">{u['total_messages']:,}</td>
            <td>{u['client_types']}</td>
        </tr>"""

    # 損益分岐点超過ユーザテーブル（月末のみ）
    breakeven_section = ''
    if report_type == 'end_month' and breakeven_users:
        breakeven_rows = ''
        for u in breakeven_users:
            breakeven_rows += f"""
            <tr>
                <td>{u['user_email']}</td>
                <td>{u['tier']}</td>
                <td class="num">${u['current_monthly_total']:,.2f}</td>
                <td><strong>{u['recommended_tier']}</strong></td>
                <td class="num">${u['recommended_monthly_cost']:,.2f}</td>
                <td class="num" style="color:green;">${u['estimated_savings']:,.2f}</td>
            </tr>"""

        breakeven_section = f"""
        <h2>⚠️ 損益分岐点超過ユーザ（推奨：プランアップグレード）</h2>
        <p>超過コストが上位プランとの差額を超えているため、アップグレードにより月額削減が見込めるユーザです。</p>
        <table>
            <thead>
                <tr>
                    <th>ユーザ</th>
                    <th>現プラン</th>
                    <th>現月額コスト（実質）</th>
                    <th>推奨プラン</th>
                    <th>推奨プラン月額</th>
                    <th>推定削減額/月</th>
                </tr>
            </thead>
            <tbody>{breakeven_rows}</tbody>
        </table>"""

    # 余剰多数ユーザテーブル（月末のみ）
    surplus_section = ''
    if report_type == 'end_month' and surplus_users:
        surplus_rows = ''
        for u in surplus_users:
            surplus_rows += f"""
            <tr>
                <td>{u['user_email']}</td>
                <td>{u['tier']}</td>
                <td class="num">{u['usage_rate']:.1f}%</td>
                <td class="num">{u['surplus_credits']:,.2f}</td>
                <td>{u['recommended_tier']}</td>
            </tr>"""

        surplus_section = f"""
        <h2>💤 余剰多数ユーザ（プラン継続意思確認）</h2>
        <p>利用率が {SURPLUS_THRESHOLD_PERCENT:.0f}% 未満のユーザです。ダウングレードまたは継続意思の確認を推奨します。</p>
        <table>
            <thead>
                <tr>
                    <th>ユーザ</th>
                    <th>現プラン</th>
                    <th>利用率</th>
                    <th>余剰クレジット</th>
                    <th>推奨プラン</th>
                </tr>
            </thead>
            <tbody>{surplus_rows}</tbody>
        </table>"""

    html = f"""<!DOCTYPE html>
<html lang="ja">
<head>
    <meta charset="UTF-8">
    <title>Kiro {title} - {end_date.strftime('%Y年%m月')}</title>
    <style>
        body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; margin: 20px; color: #333; }}
        h1 {{ color: #1a1a2e; border-bottom: 3px solid #4361ee; padding-bottom: 10px; }}
        h2 {{ color: #2d3748; margin-top: 30px; }}
        .summary {{ background: #f7f9fc; border-left: 4px solid #4361ee; padding: 15px 20px; margin: 20px 0; border-radius: 4px; }}
        .summary .cost {{ font-size: 1.5em; font-weight: bold; color: #4361ee; }}
        table {{ border-collapse: collapse; width: 100%; margin: 15px 0; font-size: 0.9em; }}
        th, td {{ border: 1px solid #e2e8f0; padding: 8px 12px; text-align: left; }}
        th {{ background: #4361ee; color: white; }}
        tr:nth-child(even) {{ background: #f8fafc; }}
        tr:hover {{ background: #edf2f7; }}
        .num {{ text-align: right; font-variant-numeric: tabular-nums; }}
        .bar-row {{ display: flex; align-items: center; margin-bottom: 8px; }}
        .bar-label {{ width: 180px; font-size: 0.85em; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }}
        .bar-bg {{ flex: 1; margin: 0 10px; background: #e2e8f0; border-radius: 4px; height: 24px; position: relative; }}
        .bar-fill {{ height: 100%; border-radius: 4px; min-width: 2px; }}
        .bar-text {{ position: absolute; top: 0; left: 0; height: 100%; width: 100%; display: flex; align-items: center; padding-left: 8px; }}
        .bar-text span {{ font-size: 0.75em; color: #2d3748; font-weight: bold; }}
        .bar-rate {{ width: 50px; text-align: right; font-size: 0.85em; font-weight: bold; }}
        .footer {{ margin-top: 30px; padding-top: 15px; border-top: 1px solid #e2e8f0; color: #718096; font-size: 0.85em; }}
    </style>
</head>
<body>
    <h1>🔍 Kiro コスト管理 {title}</h1>
    <p>対象期間：{period}</p>

    <div class="summary">
        <p>Kiro 合計料金（概算）：<span class="cost">${kiro_cost:,.2f}</span></p>
        <p>内訳：ユーザ月額合計 + 超過クレジット料金（$0.04/credit）</p>
        <p>アクティブユーザ数：{len(user_reports)} 名</p>
    </div>

    <h2>📈 ユーザ別利用クレジット</h2>
    <div style="margin: 15px 0;">
{chart_rows}
    </div>

    <h2>📊 ユーザ別利用状況</h2>
    <table>
        <thead>
            <tr>
                <th>ユーザ</th>
                <th>プラン</th>
                <th>利用クレジット / 上限</th>
                <th>利用率</th>
                <th>超過クレジット</th>
                <th>余剰クレジット</th>
                <th>メッセージ数</th>
                <th>利用環境</th>
            </tr>
        </thead>
        <tbody>{user_rows}</tbody>
    </table>

    {breakeven_section}
    {surplus_section}

    <div class="footer">
        <p>このレポートは自動生成されています。プラン料金は <a href="https://kiro.dev/pricing/">kiro.dev/pricing</a> に基づきます。</p>
        <p>生成日時：{datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')} UTC</p>
    </div>
</body>
</html>"""

    return html


def get_rate_color(rate):
    """利用率に応じた色を返す"""
    if rate >= 100:
        return '#e53e3e'  # 赤（超過）
    elif rate >= 80:
        return '#dd6b20'  # オレンジ（高利用）
    elif rate < 30:
        return '#718096'  # グレー（低利用）
    else:
        return '#2d3748'  # 通常


# ═══════════════════════════════════════════════════════════════
# メール送信
# ═══════════════════════════════════════════════════════════════

def build_subject(report_type, today):
    """メール件名を生成"""
    type_label = '中間' if report_type == 'mid_month' else '月末'
    return f"[Kiro] {today.strftime('%Y年%m月')} コスト管理{type_label}レポート"


def build_teams_summary(report_type, user_reports, kiro_cost, start_date, end_date,
                        breakeven_users, surplus_users, report_url):
    """Teams用テキストサマリを生成"""
    type_label = '中間' if report_type == 'mid_month' else '月末'
    period = f"{start_date.strftime('%Y/%m/%d')} - {end_date.strftime('%Y/%m/%d')}"

    lines = [
        f"対象期間: {period}",
        "",
        f"■ 合計料金: ${kiro_cost:,.2f}（{len(user_reports)}名）",
        "",
        "■ 利用率上位:",
    ]

    # 上位5名を表示
    for u in user_reports[:5]:
        mark = "🔴" if u['usage_rate'] >= 100 else "🟠" if u['usage_rate'] >= 80 else "🔵"
        lines.append(f"  {mark} {u['user_email']}: {u['usage_rate']:.1f}% ({u['credits_used']:,.0f}/{u['plan_limit']:,.0f})")

    if len(user_reports) > 5:
        lines.append(f"  ... 他 {len(user_reports) - 5} 名")

    # 月末レポートの追加情報
    if breakeven_users:
        lines.append("")
        lines.append(f"⚠️ 損益分岐点超過: {len(breakeven_users)}名（アップグレード推奨）")
        for u in breakeven_users[:3]:
            lines.append(f"  - {u['user_email']} ({u['tier']} → {u['recommended_tier']}, 推定削減 ${u['estimated_savings']:,.2f}/月)")

    if surplus_users:
        lines.append("")
        lines.append(f"💤 余剰多数: {len(surplus_users)}名（利用率{SURPLUS_THRESHOLD_PERCENT:.0f}%未満）")
        for u in surplus_users[:3]:
            lines.append(f"  - {u['user_email']} ({u['tier']}, 利用率{u['usage_rate']:.1f}%)")
        if len(surplus_users) > 3:
            lines.append(f"  ... 他 {len(surplus_users) - 3} 名")

    lines.append("")
    lines.append(f'📎 <a href="{report_url}">詳細レポートを開く（HTML）</a>')
    lines.append(f"   ※ リンク有効期限: 生成から約6時間")
    lines.append("")
    lines.append(f"📂 S3保存先: s3://{OUTPUT_BUCKET}/reports/{date.today().strftime('%Y/%m')}/{report_type}_{date.today().isoformat()}.html")

    return "\n".join(lines)


def send_html_email(subject, html_body, recipients):
    """SESでHTMLメールを送信（管理者宛）"""
    try:
        ses.send_email(
            Source=SES_SENDER,
            Destination={'ToAddresses': recipients},
            Message={
                'Subject': {'Data': subject, 'Charset': 'UTF-8'},
                'Body': {
                    'Html': {'Data': html_body, 'Charset': 'UTF-8'},
                },
            },
        )
        print(f"HTML email sent to {len(recipients)} recipients")
    except Exception as e:
        print(f"Error sending HTML email: {e}")
        raise


def send_text_email(subject, text_body, recipients):
    """SESでHTML形式（テキスト内容）のメールを送信（Teams宛、URLリンク化のため）"""
    # TeamsはテキストメールだとURLをリンク化しないため、
    # 本文をHTML形式で送信する（改行をbrに変換）
    # ※ build_teams_summary内で既にaタグを含む場合はそのまま維持
    html_body = text_body.replace('\n', '<br>\n')
    html_body = f"""<html><body style="font-family:sans-serif;font-size:14px;">{html_body}</body></html>"""

    try:
        ses.send_email(
            Source=SES_SENDER,
            Destination={'ToAddresses': recipients},
            Message={
                'Subject': {'Data': subject, 'Charset': 'UTF-8'},
                'Body': {
                    'Html': {'Data': html_body, 'Charset': 'UTF-8'},
                },
            },
        )
        print(f"HTML email sent to {len(recipients)} recipients (Teams)")
    except Exception as e:
        print(f"Error sending email to Teams: {e}")
        raise