from datetime import UTC, datetime
from decimal import Decimal
from uuid import uuid4

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.contracts import AgentResult, DomainEvent
from app.core.errors import RunLeaseLostError
from app.model.audit import InteractionAudit
from app.model.conversation import ConversationMessage
from app.model.platform import AgentRun, DomainEventOutbox, RequestIdempotency


class AgentPersistenceService:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def complete_run(
        self, run_id: str, result: AgentResult, memory_extraction_requested: bool = True,
        *, worker_id: str | None = None,
    ) -> int:
        now = datetime.now(UTC).replace(tzinfo=None)
        async with self.session.begin():
            run = await self.session.scalar(
                select(AgentRun).where(AgentRun.run_id == run_id).with_for_update()
            )
            if run is None:
                raise ValueError("run not found")
            if result.run_id != run_id:
                raise ValueError("result belongs to another run")
            if run.status == "succeeded" and run.result_message_id is not None:
                return run.result_message_id
            if worker_id is not None and (
                run.status != "running" or run.worker_id != worker_id
                or run.locked_until is None or run.locked_until <= now
            ):
                raise RunLeaseLostError("运行租约失效或已取消")
            if run.status not in {"queued", "running"}:
                raise RunLeaseLostError("不能覆盖运行终态")
            message = ConversationMessage(
                session_id=run.session_id, customer_id=run.user_id, portal="agent",
                role="assistant", content=result.result.text,
                trace_id=run.trace_id, created_at=now,
                intent=result.result.intent.intent if result.result.intent else None,
                confidence=(Decimal(str(result.result.intent.confidence))
                            if result.result.intent else None),
                source_references=[ref.model_dump(mode="json")
                                   for ref in result.result.source_references],
                tool_calls={"calls": [call.model_dump(mode="json")
                                      for call in result.result.tool_calls]},
            )
            self.session.add(message)
            await self.session.flush()
            run.result_message_id = message.id
            run.status = "succeeded"
            run.result_version = 1
            run.completed_at = now
            run.updated_at = now
            run.locked_until = None
            self.session.add(InteractionAudit(
                actor_type="agent", actor_id=run.user_id, target_customer_id=run.user_id,
                session_id=run.session_id, portal="agent", action_type="agent.run_completed",
                detail={"run_id": run_id, "result_message_id": message.id}, created_at=now,
            ))
            await self.session.execute(
                update(RequestIdempotency)
                .where(RequestIdempotency.id == run.idempotency_id)
                .values(status="completed", result_message_id=message.id, updated_at=now)
            )
            events = [DomainEvent(
                event_id=str(uuid4()), event_type="agent.run_completed", aggregate_type="agent_run",
                aggregate_id=run_id, trace_id=run.trace_id,
                payload={"run_id": run_id}, occurred_at=now,
            )]
            if memory_extraction_requested:
                events.append(DomainEvent(
                    event_id=str(uuid4()), event_type="memory.extraction_requested",
                    aggregate_type="agent_run", aggregate_id=run_id, trace_id=run.trace_id,
                    payload={"run_id": run_id, "message_id": message.id,
                             "customer_id": run.user_id}, occurred_at=now,
                ))
            for event in events:
                self.session.add(DomainEventOutbox(
                    event_id=event.event_id, event_type=event.event_type,
                    aggregate_type=event.aggregate_type, aggregate_id=event.aggregate_id,
                    trace_id=event.trace_id, payload=event.payload, occurred_at=now,
                    created_at=now, updated_at=now,
                ))
            return message.id