"""Shared SDK context and helper operations."""

from __future__ import annotations

from contextlib import contextmanager
from pathlib import Path
from typing import Callable
from uuid import uuid4

from ..editor import edit_file
from ..exceptions import (
    ActiveWorkflowError,
    InvalidBranchError,
    InvalidMetabranchError,
)
from ..git import GitClient
from ..models import (
    CONFIG_FILE_NAME,
    INITIALIZE_CONFIG_COMMIT_SUBJECT,
    METADATA_COMMIT_TRAILER,
    MetabranchDefinition,
    UPDATE_CONFIG_COMMIT_SUBJECT,
    WorkflowJournal,
)
from .workspace import LocalState, TrackedConfig


class _Context:
    """Shared mutable context for the split SDK clients."""

    def __init__(
        self,
        git_client: GitClient | None = None,
        tracked_config: TrackedConfig | None = None,
        local_state: LocalState | None = None,
        reconciliation_editor: Callable[[str, Path], str] | None = None,
    ) -> None:
        self.git_client = git_client or GitClient()
        self.tracked_config = tracked_config or TrackedConfig(self.git_client)
        self.local_state = local_state or LocalState(self.git_client)
        self.reconciliation_editor = reconciliation_editor or edit_file
        self.progress_reporter: Callable[[str], None] | None = None

    def load_definition(self) -> MetabranchDefinition:
        """Return the current metabranch definition."""
        if not self.tracked_config.exists():
            raise InvalidMetabranchError()
        return self.tracked_config.load()

    def save_definition(
        self,
        definition: MetabranchDefinition,
        *,
        initialize: bool = False,
    ) -> bool:
        """Persist and auto-commit the tracked metabranch definition."""
        self.tracked_config.save(definition)
        self.local_state.register(definition)
        return self.git_client.commit_only_path(
            CONFIG_FILE_NAME,
            subject=(
                INITIALIZE_CONFIG_COMMIT_SUBJECT
                if initialize
                else UPDATE_CONFIG_COMMIT_SUBJECT
            ),
            body=METADATA_COMMIT_TRAILER,
        )

    def current_workflow_snapshot(self, journal: WorkflowJournal) -> WorkflowJournal:
        """Return the latest persisted snapshot for a workflow journal."""
        return self.local_state.load_workflow(journal.workflow_id) or journal

    def with_validated_branch_updates(
        self,
        definition: MetabranchDefinition,
        *,
        added: tuple[str, ...] = (),
        removed: tuple[str, ...] = (),
    ) -> MetabranchDefinition:
        """Validate branch updates against Git before persisting them."""
        net_additions = tuple(branch for branch in added if branch not in set(removed))
        for branch_name in net_additions:
            if not self.git_client.branch_exists(branch_name):
                raise InvalidBranchError()
        return definition.update_branches(added=added, removed=removed)

    def ensure_no_active_workflow(self) -> None:
        """Reject operations when a workflow is already in progress."""
        journal = self.local_state.get_active_workflow()
        if journal is not None and journal.status not in {"completed", "aborted"}:
            raise ActiveWorkflowError(
                "Another metabranch workflow is in progress; resume or abort it first."
            )

    def start_workflow(
        self,
        *,
        workflow_type: str,
        definition: MetabranchDefinition,
        args: dict[str, object],
    ) -> WorkflowJournal:
        """Persist the initial workflow journal and acquire the repo lock."""
        checkpoints = {
            "return_branch": self.git_client.current_branch_name(),
            f"ref:{definition.name}": self.git_client.branch_head(definition.name),
        }
        for branch_name in definition.branches:
            checkpoints[f"ref:{branch_name}"] = self.git_client.branch_head(branch_name)

        journal = WorkflowJournal(
            workflow_id=uuid4().hex,
            workflow_type=workflow_type,
            metabranch_name=definition.name,
            status="running",
            current_step="preflight",
            args=args,
            checkpoints=checkpoints,
        )
        self.local_state.acquire_repo_lock(
            owner=journal.workflow_id,
            workflow_id=journal.workflow_id,
            operation=workflow_type,
            metabranch_name=definition.name,
        )
        self.local_state.save_workflow(journal)
        self.local_state.set_active(definition.name, workflow_id=journal.workflow_id)
        return journal

    def save_workflow_step(
        self, journal: WorkflowJournal, step: str
    ) -> WorkflowJournal:
        """Persist a new running step for the workflow."""
        updated = journal.with_updates(status="running", current_step=step)
        self.local_state.save_workflow(updated)
        return updated

    def restore_journal_checkpoint(self, journal: WorkflowJournal) -> None:
        """Restore refs and branch position from a saved workflow checkpoint."""
        self.git_client.abort_in_progress_operation()
        ref_names = sorted(
            (
                key.removeprefix("ref:")
                for key in journal.checkpoints
                if key.startswith("ref:")
            ),
            key=lambda name: (name == journal.metabranch_name, name),
        )
        for branch_name in ref_names:
            self.git_client.checkout(branch_name)
            self.git_client.reset_hard(journal.checkpoints[f"ref:{branch_name}"])

        return_branch = journal.checkpoints.get(
            "return_branch", journal.metabranch_name
        )
        self.git_client.checkout(return_branch)

    @contextmanager
    def operation_lock(self, operation: str, metabranch_name: str | None = None):
        """Acquire and release a short-lived repo mutation lock."""
        owner = f"{operation}:{uuid4().hex}"
        self.local_state.acquire_repo_lock(
            owner=owner,
            operation=operation,
            metabranch_name=metabranch_name,
        )
        try:
            yield
        finally:
            self.local_state.release_repo_lock(owner)
