"""Langfuse集成 - 可观测性追踪"""
from typing import Optional, Dict, Any
from datetime import datetime
import httpx
from src.config import get_settings


class LangfuseClient:
    """Langfuse客户端 - 追踪LLM调用和任务执行"""

    def __init__(self):
        settings = get_settings()
        self.public_key = getattr(settings, 'LANGFUSE_PUBLIC_KEY', '')
        self.secret_key = getattr(settings, 'LANGFUSE_SECRET_KEY', '')
        self.host = getattr(settings, 'LANGFUSE_HOST', 'https://cloud.langfuse.com')
        self.enabled = bool(self.public_key and self.secret_key)

    async def trace_task(
        self,
        task_id: str,
        user_id: str,
        query: str,
        complexity_score: float,
        metadata: Optional[Dict] = None
    ) -> Optional[str]:
        """追踪任务创建"""
        if not self.enabled:
            return None

        trace_data = {
            'id': task_id,
            'name': 'ai-researcher-task',
            'userId': user_id,
            'input': {'query': query, 'complexity_score': complexity_score},
            'metadata': metadata or {},
            'timestamp': datetime.utcnow().isoformat(),
        }

        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{self.host}/api/public/traces",
                    headers=self._get_headers(),
                    json=trace_data
                )
                if response.status_code == 200:
                    return response.json().get('id')
        except Exception as e:
            print(f"Langfuse trace error: {e}")

        return None

    async def trace_generation(
        self,
        task_id: str,
        model: str,
        prompt: str,
        completion: str,
        tokens_used: int,
        cost: Optional[float] = None,
    ) -> Optional[str]:
        """追踪LLM生成调用"""
        if not self.enabled:
            return None

        generation_data = {
            'traceId': task_id,
            'name': f'llm-{model}',
            'model': model,
            'modelParameters': {},
            'input': prompt,
            'output': completion,
            'usage': {
                'input': tokens_used // 2,
                'output': tokens_used // 2,
                'total': tokens_used,
            },
            'cost': cost,
            'timestamp': datetime.utcnow().isoformat(),
        }

        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{self.host}/api/public/generations",
                    headers=self._get_headers(),
                    json=generation_data
                )
                if response.status_code == 200:
                    return response.json().get('id')
        except Exception as e:
            print(f"Langfuse generation error: {e}")

        return None

    async def score_trace(
        self,
        trace_id: str,
        name: str,
        value: float,
        comment: Optional[str] = None,
    ):
        """为追踪添加评分"""
        if not self.enabled:
            return

        score_data = {
            'traceId': trace_id,
            'name': name,
            'value': value,
            'comment': comment,
        }

        try:
            async with httpx.AsyncClient() as client:
                await client.post(
                    f"{self.host}/api/public/scores",
                    headers=self._get_headers(),
                    json=score_data
                )
        except Exception as e:
            print(f"Langfuse score error: {e}")

    def _get_headers(self) -> Dict[str, str]:
        """获取请求头"""
        import base64
        auth = base64.b64encode(f"{self.public_key}:{self.secret_key}".encode()).decode()
        return {
            'Authorization': f'Basic {auth}',
            'Content-Type': 'application/json',
        }