"""Decision trace helpers for EDA runtime flows."""
from __future__ import annotations

import uuid
from typing import Dict, Optional

from fastapi import Request
from sqlalchemy.ext.asyncio import AsyncSession

from eda_agent.models.persistence import DecisionTrace


class TraceService:
    """Create and persist auditable decision trace entries."""

    @staticmethod
    def identity_from_request(request: Request) -> Dict[str, Optional[str]]:
        headers = request.headers
        return {
            "agent_id": headers.get("x-agent-id"),
            "key_id": headers.get("x-key-id"),
            "request_id": headers.get("x-request-id"),
            "caller_id": headers.get("x-caller-id"),
        }

    async def write_trace(
        self,
        db: AsyncSession,
        *,
        agent_id: uuid.UUID,
        stage: str,
        input_payload: dict,
        norms_applied: list,
        compliance_result: dict,
        value_ranking: list,
        final_action: dict | None,
        result_payload: dict,
        latency_ms: int,
    ) -> DecisionTrace:
        trace = DecisionTrace(
            agent_id=agent_id,
            stage=stage,
            input_payload=input_payload,
            norms_applied=norms_applied,
            compliance_result=compliance_result,
            value_ranking=value_ranking,
            final_action=final_action,
            result_payload=result_payload,
            latency_ms=latency_ms,
        )
        db.add(trace)
        await db.flush()
        return trace
