from __future__ import annotations

from .shared import *

class ModelCapabilityCostTrackingMixin:
    async def track_cost_and_roi(
        self,
        truth_node_id: int,
        period: str = "weekly",
        tenant_id: int = 1,
    ) -> Dict[str, Any]:
        """算力成本跟踪与 ROI 核算

        Args:
            truth_node_id: 宪章（需求）ID
            period: 周期类型（daily, weekly, monthly, project）
            tenant_id: 租户 ID

        Returns:
            包含成本明细、ROI 分析等信息的字典
        """
        # 获取相关任务和模型绑定
        tasks = await self._get_truth_node_tasks(truth_node_id, tenant_id)
        bindings = await self._get_truth_node_bindings(truth_node_id, tenant_id)

        if not tasks:
            return {
                "total_tokens": 0,
                "total_cost_usd": 0.0,
                "task_breakdown": [],
                "roi_analysis": {},
                "cost_variance": 0.0,
            }

        # 汇总 token 使用量和成本
        total_tokens = 0
        total_cost = 0.0
        task_costs = []

        for task in tasks:
            # 从 AI 调用日志统计 token
            task_tokens = await self._get_task_token_usage(task.id)
            task_cost = await self._calculate_task_cost(task, bindings)

            total_tokens += task_tokens
            total_cost += task_cost
            task_costs.append({
                "task_id": task.id,
                "task_name": task.title,
                "tokens": task_tokens,
                "cost_usd": round(task_cost, 4),
            })

        # 获取或创建成本跟踪记录
        cost_tracking = await self._get_or_create_cost_tracking(
            truth_node_id=truth_node_id,
            period=period,
            tenant_id=tenant_id,
            bindings=bindings,
        )

        # 更新实际数据
        cost_tracking.actual_tokens = total_tokens
        cost_tracking.actual_cost_usd = round(total_cost, 4)

        # AI 分析 ROI
        truth_node = await self._get_truth_node(truth_node_id, tenant_id)
        roi_analysis = await self._analyze_cost_roi(
            total_cost=total_cost,
            total_tokens=total_tokens,
            task_count=len(tasks),
            truth_node_content=str(truth_node.content or truth_node.title or "") if truth_node else "",
            tenant_id=tenant_id,
        )

        cost_tracking.roi_score = roi_analysis.get("roi_score")
        cost_tracking.roi_analysis = roi_analysis

        # 生成优化建议
        optimization = await self._generate_cost_optimization(
            truth_node_id=truth_node_id,
            tenant_id=tenant_id,
            total_cost=total_cost,
        )
        cost_tracking.optimization_suggestions = optimization

        await self.db.commit()
        await self.db.refresh(cost_tracking)

        # 计算成本偏差
        cost_variance = total_cost - (cost_tracking.estimated_cost_usd or 0)

        return {
            "total_tokens": total_tokens,
            "total_cost_usd": round(total_cost, 4),
            "task_breakdown": task_costs,
            "roi_analysis": roi_analysis,
            "cost_variance": round(cost_variance, 4),
            "optimization_suggestions": optimization,
        }

    async def _get_truth_node_tasks(
        self, truth_node_id: int, tenant_id: int
    ) -> List[Task]:
        """获取宪章下的所有任务"""
        result = await self.db.execute(
            select(Task).where(
                and_(
                    Task.truth_node_id == truth_node_id,
                    Task.tenant_id == tenant_id,
                )
            )
        )
        return list(result.scalars().all())

    async def _get_truth_node_bindings(
        self, truth_node_id: int, tenant_id: int
    ) -> List[ModelRequirementBinding]:
        """获取宪章的模型绑定"""
        result = await self.db.execute(
            select(ModelRequirementBinding).where(
                and_(
                    ModelRequirementBinding.truth_node_id == truth_node_id,
                    ModelRequirementBinding.tenant_id == tenant_id,
                )
            )
        )
        return list(result.scalars().all())

    async def _get_task_token_usage(self, task_id: int) -> int:
        """从 AI 调用日志统计任务的 token 使用量"""
        # 注意：需要根据实际日志记录方式调整
        # 这里假设日志中有 task_id 或相关标识
        result = await self.db.execute(
            select(func.sum(AICallLog.total_tokens)).where(
                AICallLog.function_name.contains(f"task_{task_id}")
            )
        )
        total = result.scalar() or 0
        return int(total)

    async def _calculate_task_cost(
        self, task: Task, bindings: List[ModelRequirementBinding]
    ) -> float:
        """计算任务成本"""
        # 如果有绑定记录，使用绑定的预估成本
        if bindings:
            return sum(float(b.estimated_cost or 0) for b in bindings) / max(len(bindings), 1)
        return 0.0

    async def _get_or_create_cost_tracking(
        self,
        truth_node_id: int,
        period: str,
        tenant_id: int,
        bindings: List[ModelRequirementBinding],
    ) -> CostTracking:
        """获取或创建成本跟踪记录"""
        # 查找现有记录
        result = await self.db.execute(
            select(CostTracking).where(
                and_(
                    CostTracking.truth_node_id == truth_node_id,
                    CostTracking.period == period,
                    CostTracking.tenant_id == tenant_id,
                )
            ).order_by(desc(CostTracking.created_at))
        )
        tracking = result.scalar_one_or_none()

        if not tracking:
            # 创建新记录
            tracking = CostTracking(
                tenant_id=tenant_id,
                truth_node_id=truth_node_id,
                model_version_id=bindings[0].model_version_id if bindings else None,
                period=period,
                estimated_tokens=0,
                actual_tokens=0,
                estimated_cost_usd=sum(float(b.estimated_cost or 0) for b in bindings),
                actual_cost_usd=0.0,
                roi_score=None,
                roi_analysis={},
                optimization_suggestions={},
            )
            self.db.add(tracking)
            await self.db.commit()
            await self.db.refresh(tracking)

        return tracking

    async def _analyze_cost_roi(
        self,
        total_cost: float,
        total_tokens: int,
        task_count: int,
        truth_node_content: str,
        tenant_id: int,
    ) -> Dict[str, Any]:
        """AI 分析成本效益"""
        from ...config.ai_prompts import PROMPT_ANALYZE_COST_ROI

        prompt = PROMPT_ANALYZE_COST_ROI.format(
            total_cost=f"${total_cost:.2f}",
            total_tokens=f"{total_tokens:,}",
            task_count=task_count,
            truth_node_content=truth_node_content[:2000],  # 限制长度
        )

        try:
            text = await self.ai_service._invoke_with_retry(
                prompt,
                db=self.db,
                tenant_id=tenant_id,
                function_name="analyze_cost_roi",
            )
            return self.ai_service._parse_json(text)
        except Exception:
            # 失败时返回默认分析
            return {
                "roi_score": 5.0,
                "cost_efficiency": "fair",
                "cost_per_task": round(total_cost / max(task_count, 1), 2),
                "optimization_suggestions": [],
                "benchmark_comparison": "无法分析",
            }

    async def _generate_cost_optimization(
        self, truth_node_id: int, tenant_id: int, total_cost: float
    ) -> Dict[str, Any]:
        """生成成本优化建议"""
        # 简单的优化建议逻辑
        suggestions = []

        # 检查是否可以使用更便宜的模型
        bindings = await self._get_truth_node_bindings(truth_node_id, tenant_id)
        if bindings:
            for binding in bindings:
                model = await self._get_model_version(binding.model_version_id)
                if model and model.cost_per_1k_tokens and model.cost_per_1k_tokens > 0.001:
                    suggestions.append({
                        "type": "model_downgrade",
                        "priority": "medium",
                        "description": f"考虑切换到更便宜的模型替代 {model.name}",
                        "estimated_savings_percent": 30,
                    })

        # 检查是否可以启用缓存
        suggestions.append({
            "type": "cache_optimization",
            "priority": "high",
            "description": "启用 Redis 缓存可减少 20-40% 重复调用",
            "estimated_savings_percent": 30,
        })

        return {
            "suggestions": suggestions,
            "total_potential_savings_percent": 30 if suggestions else 0,
        }