"""Codegen domain - data access layer."""
from __future__ import annotations
from typing import Any, Optional
from uuid import UUID

_PENDING_USER_INPUTS_UNSET = object()
_INTERRUPTION_PAYLOAD_UNSET = object()

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import defer

from src.domain.codegen.models import CodegenRun
from src.domain.codegen.run_lifecycle import ACTIVE_RUN_RESERVING_STATUSES


class CodegenRunRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    @property
    def db(self) -> AsyncSession:
        """The session, for collaborators that must share this transaction (capability approvals)."""
        return self._db

    async def commit(self) -> None:
        """Commit a small callback update before awaiting external I/O."""
        await self._db.commit()

    async def get_project_discovery(self, project_id: UUID) -> dict | None:
        """Read the run's Core composition input before dispatching install work."""
        from src.domain.projects.models import Project

        return await self._db.scalar(
            select(Project.discovery_json).where(
                Project.id == project_id, Project.deleted_at.is_(None)
            )
        )

    async def lock_project_for_run_creation(self, project_id: UUID) -> None:
        from src.domain.projects.models import Project
        await self._db.execute(select(Project.id).where(Project.id == project_id).with_for_update())

    async def get_by_id(self, run_id: UUID, *, for_update: bool = False) -> Optional[CodegenRun]:
        """``for_update`` locks the row until the transaction ends: a read-then-write decision on
        the run (e.g. retry) then cannot interleave with the same decision made concurrently."""
        stmt = select(CodegenRun).where(CodegenRun.id == run_id)
        if for_update:
            stmt = stmt.with_for_update()
        result = await self._db.execute(stmt)
        return result.scalar_one_or_none()

    async def get_active_reserving_run(self, project_id: UUID) -> Optional[CodegenRun]:
        """Latest run whose status still reserves the project chat slot (see run_lifecycle)."""
        result = await self._db.execute(
            select(CodegenRun)
            .where(CodegenRun.project_id == project_id)
            .where(CodegenRun.status.in_(tuple(ACTIVE_RUN_RESERVING_STATUSES)))
            .order_by(CodegenRun.created_at.desc())
            .limit(1)
        )
        return result.scalar_one_or_none()

    async def list_for_device(self, device_id: UUID, statuses: Optional[list[str]] = None) -> list[CodegenRun]:
        """Runs assigned to a device (execution_target='local'), optionally filtered by status.

        Used by ``GET /api/devices/{device_id}/runs`` — the desktop app polls this after
        heartbeat/reconnect to pick up work it may have missed over the WS channel.
        """
        stmt = (
            select(CodegenRun)
            .options(defer(CodegenRun.agent_transcript))
            .where(CodegenRun.device_id == device_id)
            .order_by(CodegenRun.created_at.desc())
        )
        if statuses:
            stmt = stmt.where(CodegenRun.status.in_(tuple(statuses)))
        result = await self._db.execute(stmt)
        return list(result.scalars().all())

    async def try_claim_for_device(self, run_id: UUID, device_id: UUID) -> Optional[CodegenRun]:
        """Atomically take a local queued run. Returns None if another claimer won."""
        from datetime import datetime, timezone

        from sqlalchemy import update

        now = datetime.now(timezone.utc)
        result = await self._db.execute(
            update(CodegenRun)
            .where(CodegenRun.id == run_id)
            .where(CodegenRun.device_id == device_id)
            .where(CodegenRun.execution_target == "local")
            .where(CodegenRun.status.in_(("queued", "pending")))
            .values(status="running", claimed_at=now, started_at=now)
        )
        if result.rowcount != 1:
            return None
        await self._db.flush()
        return await self.get_by_id(run_id)

    async def list_for_project(self, project_id: UUID) -> list[CodegenRun]:
        # Defer large agent_transcript blob — list views use CodegenRunSummary; fetch full run for GET /runs/{id}.
        result = await self._db.execute(
            select(CodegenRun)
            .options(defer(CodegenRun.agent_transcript))
            .where(CodegenRun.project_id == project_id)
            .order_by(CodegenRun.created_at.desc())
        )
        return list(result.scalars().all())

    async def create(
        self,
        project_id: UUID,
        prompt_id: Optional[UUID] = None,
        created_by: Optional[UUID] = None,
        run_mode: str = "agent",
        selected_elements: Optional[list] = None,
        parent_plan_id: Optional[UUID] = None,
        execution_target: str = "cloud",
        device_id: Optional[UUID] = None,
    ) -> CodegenRun:
        run = CodegenRun(
            project_id=project_id,
            prompt_id=prompt_id,
            status="pending",
            created_by=created_by,
            run_mode=run_mode,
            selected_elements=selected_elements,
            parent_plan_id=parent_plan_id,
            execution_target=execution_target,
            device_id=device_id,
        )
        self._db.add(run)
        await self._db.flush()
        await self._db.refresh(run)
        return run

    async def update_status(
        self,
        run_id: UUID,
        status: str,
        plan: Optional[dict] = None,
        plan_document: Optional[str] = None,
        result_commit_sha: Optional[str] = None,
        result_pr_url: Optional[str] = None,
        result_summary: Optional[str] = None,
        revision_id: Optional[UUID] = None,
        error_message: Optional[str] = None,
        tokens_used: Optional[int] = None,
        agent_transcript: Optional[str] = None,
        pending_user_inputs: Any = _PENDING_USER_INPUTS_UNSET,
        interruption_payload: Any = _INTERRUPTION_PAYLOAD_UNSET,
    ) -> Optional[CodegenRun]:
        run = await self.get_by_id(run_id)
        if not run:
            return None

        from datetime import datetime, timezone

        run.status = status
        if status == "queued":
            run.completed_at = None
        if plan is not None:
            run.plan = plan
        if plan_document is not None:
            run.plan_document = plan_document
        if result_commit_sha is not None:
            run.result_commit_sha = result_commit_sha
        if result_pr_url is not None:
            run.result_pr_url = result_pr_url
        if result_summary is not None:
            run.result_summary = result_summary
        if revision_id is not None:
            run.revision_id = revision_id
        if error_message is not None:
            run.error_message = error_message
        if tokens_used is not None:
            run.tokens_used = tokens_used
        if agent_transcript is not None:
            run.agent_transcript = agent_transcript
        if pending_user_inputs is not _PENDING_USER_INPUTS_UNSET:
            run.pending_user_inputs = pending_user_inputs
        if interruption_payload is not _INTERRUPTION_PAYLOAD_UNSET:
            run.interruption_payload = interruption_payload

        if status == "running" and run.started_at is None:
            run.started_at = datetime.now(timezone.utc)
        if status in (
            "completed",
            "completed_with_limitations",
            "completed_with_l3",
            "blocked",
            "failed",
            "canceled",
            "plan_ready",
            "awaiting_user_input",
            "user_input_required",
            "clarification_required",
        ):
            run.completed_at = datetime.now(timezone.utc)
        if status in (
            "completed",
            "completed_with_limitations",
            "completed_with_l3",
            "blocked",
            "awaiting_user_input",
            "user_input_required",
            "clarification_required",
        ):
            run.worker_checkpoint = None
            run.checkpoint_payload = None
        if status in ("completed", "completed_with_limitations", "completed_with_l3", "blocked", "failed", "canceled"):
            run.interruption_payload = None

        await self._db.flush()
        return run

    async def set_pending_user_inputs(
        self, run_id: UUID, payload: Optional[dict]
    ) -> Optional[CodegenRun]:
        """Set or clear ``pending_user_inputs`` without changing status or timestamps."""
        run = await self.get_by_id(run_id)
        if not run:
            return None
        run.pending_user_inputs = payload
        await self._db.flush()
        return run

    async def cancel_run(
        self, run_id: UUID, *, error_message: Optional[str] = None
    ) -> Optional[CodegenRun]:
        """Cancel any run that still holds the project slot (see ACTIVE_RUN_RESERVING_STATUSES)."""
        run = await self.get_by_id(run_id)
        if not run:
            return None
        if run.status not in ACTIVE_RUN_RESERVING_STATUSES:
            return None
        return await self.update_status(
            run_id,
            "canceled",
            error_message=error_message if error_message is not None else "Canceled",
            pending_user_inputs=None,
        )

    async def refresh_started_at(self, run_id: UUID) -> None:
        """Reset ``started_at`` when reclaiming a stale ``running`` run."""
        from datetime import datetime, timezone

        run = await self.get_by_id(run_id)
        if run:
            run.started_at = datetime.now(timezone.utc)
            await self._db.flush()

    async def update_checkpoint(
        self,
        run_id: UUID,
        worker_checkpoint: Optional[str],
        checkpoint_payload: Optional[dict] = None,
    ) -> Optional[CodegenRun]:
        run = await self.get_by_id(run_id)
        if not run:
            return None
        run.worker_checkpoint = worker_checkpoint
        if checkpoint_payload is not None:
            run.checkpoint_payload = checkpoint_payload
        await self._db.flush()
        return run

    async def reset_run_for_agent_execution(
        self, run_id: UUID, plan: dict
    ) -> Optional[CodegenRun]:
        """Same-row transition plan_ready → agent queue (no child run)."""
        run = await self.get_by_id(run_id)
        if not run:
            return None
        run.run_mode = "agent"
        run.parent_plan_id = None
        run.status = "queued"
        run.plan = plan
        run.completed_at = None
        run.worker_checkpoint = None
        run.checkpoint_payload = None
        run.interruption_payload = None
        await self._db.flush()
        await self._db.refresh(run)
        return run

    async def reset_run_from_failure(self, run_id: UUID) -> Optional[CodegenRun]:
        """Retry: reuse row, clear terminal failure state, re-queue."""
        run = await self.get_by_id(run_id)
        if not run:
            return None
        run.status = "queued"
        run.error_message = None
        run.completed_at = None
        run.worker_checkpoint = None
        run.checkpoint_payload = None
        run.started_at = None
        await self._db.flush()
        await self._db.refresh(run)
        return run

    async def requeue_with_plan(self, run_id: UUID, plan: dict) -> Optional[CodegenRun]:
        """Resume same run with a new plan payload (e.g. after user input)."""
        run = await self.get_by_id(run_id)
        if not run:
            return None
        run.status = "queued"
        run.plan = plan
        run.completed_at = None
        run.worker_checkpoint = None
        run.checkpoint_payload = None
        run.pending_user_inputs = None
        run.interruption_payload = None
        await self._db.flush()
        await self._db.refresh(run)
        return run