#!/usr/bin/env python3
"""
计费系统健康检查脚本

功能：
1. 检测假用量问题（配额数据与实际运行Agent不一致）
2. 检查余额不足用户
3. 测试余额不足自动停止功能
4. 生成详细健康报告

使用方法：
在mcp-server容器内运行：

# 完整检查
python check_billing_health.py

# 只检查假用量
python check_billing_health.py --check-quota-only

# 只检查余额
python check_billing_health.py --check-balance-only

# 测试自动停止（需要指定用户ID）
python check_billing_health.py --test-auto-stop <user_id>
"""

import asyncio
import json
import logging
import argparse
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional
from decimal import Decimal

from sqlalchemy import select, func, and_
from sqlalchemy.ext.asyncio import AsyncSession

from database import get_db
from models import (
    Agent, 
    TenantCustomAgentQuota, 
    AgentBillingRecord,
    Balance,
    User
)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class BillingHealthChecker:
    """计费系统健康检查器"""

    def __init__(self):
        self.report_file = Path("/app/logs/billing_health_report.json")
        self.report_file.parent.mkdir(exist_ok=True)

    async def check_quota_consistency(self, db: AsyncSession) -> Dict[str, Any]:
        """
        检查配额一致性（假用量检测）
        
        对比：
        - 配额表中的使用量
        - 实际运行的Agent数量和资源
        """
        logger.info("=" * 60)
        logger.info("🔍 开始检查配额一致性（假用量检测）")
        logger.info("=" * 60)

        # 获取所有配额记录
        result = await db.execute(select(TenantCustomAgentQuota))
        quotas = result.scalars().all()

        inconsistencies = []
        total_fake_cpu = 0.0
        total_fake_memory = 0.0
        total_fake_agents = 0

        for quota in quotas:
            tenant_id = str(quota.tenant_id)

            # 查询配额使用量
            quota_usage = {
                'cpu_used': float(quota.cpu_used or 0),
                'memory_used': float(quota.memory_used or 0),
                'agent_count': quota.agent_count or 0
            }

            # 查询实际运行的Agent
            real_agents_result = await db.execute(
                select(Agent)
                .where(Agent.owner_id == tenant_id)
                .where(Agent.type == 'custom')
                .where(Agent.status == 'active')
            )
            real_agents = real_agents_result.scalars().all()

            # 计算实际使用量
            real_usage = {
                'cpu_used': sum(float(a.cpu or 0) for a in real_agents),
                'memory_used': sum(float(a.memory or 0) for a in real_agents),
                'agent_count': len(real_agents)
            }

            # 计算差异
            cpu_diff = quota_usage['cpu_used'] - real_usage['cpu_used']
            memory_diff = quota_usage['memory_used'] - real_usage['memory_used']
            agent_diff = quota_usage['agent_count'] - real_usage['agent_count']

            # 允许小误差（0.01核，0.01GB）
            has_inconsistency = (
                abs(cpu_diff) > 0.01 or 
                abs(memory_diff) > 0.01 or 
                agent_diff != 0
            )

            if has_inconsistency:
                inconsistency = {
                    'tenant_id': tenant_id,
                    'quota_usage': quota_usage,
                    'real_usage': real_usage,
                    'diff': {
                        'cpu': cpu_diff,
                        'memory': memory_diff,
                        'agent_count': agent_diff
                    },
                    'severity': 'high' if (cpu_diff > 1 or memory_diff > 1 or agent_diff > 2) else 'medium'
                }
                inconsistencies.append(inconsistency)

                total_fake_cpu += cpu_diff
                total_fake_memory += memory_diff
                total_fake_agents += agent_diff

                logger.warning(
                    f"⚠️  检测到配额不一致: 租户={tenant_id}\n"
                    f"    配额显示: CPU={quota_usage['cpu_used']:.2f}核, "
                    f"内存={quota_usage['memory_used']:.2f}GB, Agent={quota_usage['agent_count']}个\n"
                    f"    实际运行: CPU={real_usage['cpu_used']:.2f}核, "
                    f"内存={real_usage['memory_used']:.2f}GB, Agent={real_usage['agent_count']}个\n"
                    f"    假用量: CPU={cpu_diff:.2f}核, "
                    f"内存={memory_diff:.2f}GB, Agent={agent_diff}个"
                )

        summary = {
            'total_tenants': len(quotas),
            'inconsistent_tenants': len(inconsistencies),
            'consistency_rate': (len(quotas) - len(inconsistencies)) / len(quotas) * 100 if quotas else 100,
            'total_fake_cpu': total_fake_cpu,
            'total_fake_memory': total_fake_memory,
            'total_fake_agents': total_fake_agents,
            'details': inconsistencies
        }

        if inconsistencies:
            logger.error(
                f"❌ 发现 {len(inconsistencies)} 个租户存在配额不一致问题！\n"
                f"   假用量汇总: CPU={total_fake_cpu:.2f}核, "
                f"内存={total_fake_memory:.2f}GB, Agent={total_fake_agents}个\n"
                f"   建议运行: python fix_fake_quota.py"
            )
        else:
            logger.info("✅ 所有租户配额数据一致，无假用量问题")

        return summary

    async def check_balance_status(self, db: AsyncSession) -> Dict[str, Any]:
        """
        检查用户余额状态
        
        识别：
        - 余额不足的用户
        - 透支用户
        - 有运行Agent但余额不足的用户（风险）
        """
        logger.info("=" * 60)
        logger.info("💰 开始检查用户余额状态")
        logger.info("=" * 60)

        # 查询所有用户及其余额
        result = await db.execute(
            select(User, Balance)
            .outerjoin(Balance, User.id == Balance.user_id)
        )
        users_data = result.all()

        low_balance_users = []
        overdraft_users = []
        risky_users = []  # 有运行Agent但余额不足

        for user, balance in users_data:
            user_id = str(user.id)
            eu_balance = float(balance.eu_balance) if balance else 0.0
            credit_limit = float(user.credit_limit or 0)
            available = eu_balance + credit_limit

            # 检查是否有运行中的Agent
            running_agents_result = await db.execute(
                select(func.count(AgentBillingRecord.id))
                .where(AgentBillingRecord.user_id == user_id)
                .where(AgentBillingRecord.end_time == None)
            )
            running_agent_count = running_agents_result.scalar() or 0

            user_info = {
                'user_id': user_id,
                'email': user.email,
                'eu_balance': eu_balance,
                'credit_limit': credit_limit,
                'available_balance': available,
                'running_agents': running_agent_count
            }

            # 透支（可用余额为负）
            if available < 0:
                overdraft_users.append(user_info)
                logger.error(
                    f"💥 透支用户: {user.email} (ID: {user_id})\n"
                    f"   账户余额: {eu_balance:.2f} EU\n"
                    f"   授信额度: {credit_limit:.2f} EU\n"
                    f"   可用余额: {available:.2f} EU\n"
                    f"   运行Agent数: {running_agent_count}"
                )

            # 余额不足（低于10 EU）
            elif available < 10:
                low_balance_users.append(user_info)
                logger.warning(
                    f"⚠️  余额不足: {user.email} (ID: {user_id})\n"
                    f"   可用余额: {available:.2f} EU\n"
                    f"   运行Agent数: {running_agent_count}"
                )

            # 风险用户：有运行Agent但余额很低（< 5 EU）
            if running_agent_count > 0 and available < 5:
                risky_users.append(user_info)
                logger.warning(
                    f"🚨 风险用户: {user.email} (ID: {user_id})\n"
                    f"   有 {running_agent_count} 个Agent运行但余额仅剩 {available:.2f} EU"
                )

        summary = {
            'total_users': len(users_data),
            'overdraft_users': len(overdraft_users),
            'low_balance_users': len(low_balance_users),
            'risky_users': len(risky_users),
            'overdraft_details': overdraft_users,
            'low_balance_details': low_balance_users,
            'risky_users_details': risky_users
        }

        logger.info(
            f"\n📊 余额状态汇总:\n"
            f"   总用户数: {len(users_data)}\n"
            f"   透支用户: {len(overdraft_users)}\n"
            f"   余额不足: {len(low_balance_users)}\n"
            f"   风险用户: {len(risky_users)}"
        )

        return summary

    async def test_auto_stop(self, db: AsyncSession, user_id: str) -> Dict[str, Any]:
        """
        测试余额不足自动停止功能
        
        模拟周期计费检测到余额不足时的行为
        """
        logger.info("=" * 60)
        logger.info(f"🧪 测试自动停止功能: user_id={user_id}")
        logger.info("=" * 60)

        from app.billing import get_available_balance
        from app.periodic_billing import stop_user_agents

        # 检查用户余额
        balance, credit_limit, available = await get_available_balance(user_id, db)

        logger.info(
            f"用户余额状态:\n"
            f"   账户余额: {balance:.2f} EU\n"
            f"   授信额度: {credit_limit:.2f} EU\n"
            f"   可用余额: {available:.2f} EU"
        )

        # 检查运行中的Agent
        running_agents_result = await db.execute(
            select(AgentBillingRecord)
            .where(AgentBillingRecord.user_id == user_id)
            .where(AgentBillingRecord.end_time == None)
        )
        running_agents = running_agents_result.scalars().all()

        logger.info(f"运行中的Agent数量: {len(running_agents)}")
        for agent in running_agents:
            logger.info(f"  - {agent.agent_name} (启动时间: {agent.start_time})")

        result = {
            'user_id': user_id,
            'balance': float(balance),
            'credit_limit': float(credit_limit),
            'available': float(available),
            'running_agents_before': len(running_agents),
            'stopped_agents': [],
            'action_taken': 'none'
        }

        # 如果余额不足，执行停止
        if available < 0:
            logger.warning(f"⚠️  可用余额为负 ({available:.2f} EU)，将停止所有Agent...")
            
            stopped = await stop_user_agents(user_id, db)
            result['stopped_agents'] = stopped
            result['action_taken'] = 'stopped'

            logger.info(
                f"✅ 已停止 {len(stopped)} 个Agent:\n" +
                "\n".join(f"  - {name}" for name in stopped)
            )
        else:
            logger.info(f"✅ 余额充足 ({available:.2f} EU)，无需停止Agent")
            result['action_taken'] = 'no_action_needed'

        return result

    async def generate_full_report(self, db: AsyncSession) -> Dict[str, Any]:
        """生成完整的健康检查报告"""
        logger.info("\n" + "=" * 60)
        logger.info("📋 生成计费系统健康检查报告")
        logger.info("=" * 60 + "\n")

        report = {
            'check_time': datetime.utcnow().isoformat(),
            'quota_consistency': await self.check_quota_consistency(db),
            'balance_status': await self.check_balance_status(db)
        }

        # 生成健康评分
        quota_ok = report['quota_consistency']['inconsistent_tenants'] == 0
        no_overdraft = report['balance_status']['overdraft_users'] == 0
        few_risky = report['balance_status']['risky_users'] < 5

        health_score = 0
        if quota_ok:
            health_score += 40
        if no_overdraft:
            health_score += 40
        if few_risky:
            health_score += 20

        report['health_score'] = health_score
        report['health_status'] = (
            'healthy' if health_score >= 90 else
            'warning' if health_score >= 70 else
            'critical'
        )

        # 生成建议
        recommendations = []
        if not quota_ok:
            recommendations.append("运行 fix_fake_quota.py 修复配额不一致问题")
        if not no_overdraft:
            recommendations.append(f"有 {report['balance_status']['overdraft_users']} 个用户透支，建议停止其Agent或充值")
        if not few_risky:
            recommendations.append(f"有 {report['balance_status']['risky_users']} 个风险用户，建议提醒充值")

        report['recommendations'] = recommendations

        # 保存报告
        with open(self.report_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False)

        logger.info(f"\n{'=' * 60}")
        logger.info(f"📊 健康检查完成")
        logger.info(f"{'=' * 60}")
        logger.info(f"健康评分: {health_score}/100 ({report['health_status'].upper()})")
        logger.info(f"报告保存: {self.report_file}")

        if recommendations:
            logger.info(f"\n💡 建议:")
            for i, rec in enumerate(recommendations, 1):
                logger.info(f"   {i}. {rec}")

        return report


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='计费系统健康检查')
    parser.add_argument('--check-quota-only', action='store_true', help='只检查配额一致性')
    parser.add_argument('--check-balance-only', action='store_true', help='只检查余额状态')
    parser.add_argument('--test-auto-stop', type=str, metavar='USER_ID', help='测试自动停止功能（指定用户ID）')
    
    args = parser.parse_args()

    checker = BillingHealthChecker()

    async for db in get_db():
        try:
            if args.check_quota_only:
                await checker.check_quota_consistency(db)
            elif args.check_balance_only:
                await checker.check_balance_status(db)
            elif args.test_auto_stop:
                await checker.test_auto_stop(db, args.test_auto_stop)
            else:
                await checker.generate_full_report(db)

        except Exception as e:
            logger.error(f"执行失败: {str(e)}", exc_info=True)
        finally:
            break


if __name__ == "__main__":
    asyncio.run(main())