from __future__ import annotations

from typing import Any

from .locks import WorkspaceLockManager
from .models import QueryRequest, QueryResponse, SyncRequest, SyncResponse
from .qmd_client import QMDClient
from .sync_service import SyncService


class QueryService:
    def __init__(
        self,
        *,
        sync_service: SyncService,
        qmd_client: QMDClient,
        lock_manager: WorkspaceLockManager,
    ) -> None:
        self.sync_service = sync_service
        self.qmd_client = qmd_client
        self.lock_manager = lock_manager

    def handle_query(self, request: QueryRequest) -> QueryResponse:
        resolved = self.sync_service.workspace_manager.resolve_workspace(
            branch=request.branch,
            memory_profile=request.memory_profile,
        )

        with self.lock_manager.acquire(resolved.workspace_name) as lock_stats:
            sync_meta = self.sync_service.sync_for_query(
                branch=resolved.branch,
                memory_profile=None,
                require_latest=request.require_latest,
            )

            results, qmd_command = self.qmd_client.run_query(
                index_name=sync_meta.qmd_index,
                collection_name=sync_meta.qmd_collection,
                query_type=request.query_type,
                query=request.query,
                n=request.n,
            )

            debug: dict[str, Any] | None = None
            if request.debug:
                debug = {
                    "lock_wait_ms": lock_stats.wait_ms,
                    "workspace_changed": sync_meta.changed,
                    "previous_commit": sync_meta.previous_commit,
                    "qmd_index": sync_meta.qmd_index,
                    "qmd_command": qmd_command,
                    "update_ran": sync_meta.update_ran,
                    "embed_ran": sync_meta.embed_ran,
                    "embed_error": sync_meta.embed_error,
                }

            return QueryResponse(
                ok=True,
                branch=sync_meta.branch,
                resolved_workspace=sync_meta.workspace_relative_path,
                commit_hash=sync_meta.commit_hash,
                synced_at=sync_meta.synced_at,
                query_type=request.query_type,
                results=results,
                qmd_collection=sync_meta.qmd_collection,
                debug=debug,
            )

    def handle_sync(self, request: SyncRequest) -> SyncResponse:
        resolved = self.sync_service.workspace_manager.resolve_workspace(
            branch=request.branch,
            memory_profile=request.memory_profile,
        )

        with self.lock_manager.acquire(resolved.workspace_name) as lock_stats:
            sync_meta = self.sync_service.sync_explicit(
                branch=resolved.branch,
                memory_profile=None,
                require_latest=request.require_latest,
            )
            debug = {
                "lock_wait_ms": lock_stats.wait_ms,
                "workspace_changed": sync_meta.changed,
                "previous_commit": sync_meta.previous_commit,
                "qmd_index": sync_meta.qmd_index,
                "update_ran": sync_meta.update_ran,
                "embed_ran": sync_meta.embed_ran,
                "embed_error": sync_meta.embed_error,
            }
            return SyncResponse(
                ok=True,
                branch=sync_meta.branch,
                resolved_workspace=sync_meta.workspace_relative_path,
                commit_hash=sync_meta.commit_hash,
                synced_at=sync_meta.synced_at,
                qmd_collection=sync_meta.qmd_collection,
                debug=debug,
            )