"""Async client for Alignment Service API."""

from __future__ import annotations

from pathlib import Path
from typing import Any, BinaryIO, TypeVar
from uuid import UUID

from pydantic import BaseModel

from .schemas import (
    AddFindingRequest,
    AlignmentCommandsResponse,
    AlignmentCloseRequest,
    AlignmentFindingsResponse,
    BucketCommandsResponse,
    BucketFindingsResponse,
    MapFindingsRequest,
    AlignmentCreateRequest,
    AlignmentExecuteRequest,
    AlignmentExecuteResponse,
    AlignmentResponse,
    AlignmentTypeResponse,
    AlignmentUpdateRequest,
    BucketCreateRequest,
    BucketDetailResponse,
    BucketEventResponse,
    BucketResponse,
    BucketStatusUpdateRequest,
    CommandCreateRequest,
    CommandEventResponse,
    CommandFindingsMapCreateRequest,
    CommandFindingsMapResponse,
    CommandGraphResponse,
    CommandGroupEventResponse,
    CommandGroupCreateRequest,
    CommandGroupResponse,
    CommandGroupStatusUpdateRequest,
    CommandGroupUpdateRequest,
    CommandResponse,
    DependencyGraphResponse,
    FindingDependencyGraphRequest,
    FindingMatchQueryCreateRequest,
    FindingMatchQueryResponse,
    FindingMatchQueryUpdateRequest,
    FindingResponse,
    ImpactAssessmentWithDependenciesResponse,
    QueryProviderResponse,
    RetireFindingRequest,
    UpdateCommandFindingsMappingRequest,
    UpdateFindingRequest,
    PatchFindingAdditionalDataRequest,
)
from .schemas.command import PushCommandGroupRequest, UpdateCommandRequest
from datasdr_base_api_client import BaseHTTPClient, JSON_TYPE, TokenProvider

SchemaT = TypeVar("SchemaT", bound=BaseModel)


class AlignmentServiceClient(BaseHTTPClient):
    """Typed async API client for Alignment Service."""

    def __init__(
        self,
        base_url: str,
        token: str | None = None,
        timeout: int = 30,
        token_provider: TokenProvider | None = None,
        auth_retry_attempts: int = 3,
        api_prefix: str = "/v1",
    ):
        super().__init__(
            base_url=base_url,
            token=token,
            timeout=timeout,
            token_provider=token_provider,
            auth_retry_attempts=auth_retry_attempts,
        )
        self.api_prefix = api_prefix.rstrip("/")

    def _path(self, endpoint: str) -> str:
        endpoint_path = endpoint if endpoint.startswith("/") else f"/{endpoint}"
        return f"{self.api_prefix}{endpoint_path}"

    @staticmethod
    def _dump(model: BaseModel) -> dict[str, Any]:
        return model.model_dump(mode="json", by_alias=True, exclude_none=True)

    def _parse(self, payload: JSON_TYPE, schema: type[SchemaT]) -> SchemaT:
        if payload is None:
            raise ValueError(f"Expected {schema.__name__} payload but got no content")
        return schema.model_validate(payload)

    def _parse_list(self, payload: JSON_TYPE, schema: type[SchemaT]) -> list[SchemaT]:
        if not payload:
            return []
        if not isinstance(payload, list):
            raise ValueError(f"Expected list payload for {schema.__name__}")
        return [schema.model_validate(item) for item in payload]

    async def get_alignment_types(self) -> list[AlignmentTypeResponse]:
        response = await self.get(self._path("/alignments/types"))
        return self._parse_list(response, AlignmentTypeResponse)

    async def get_alignment_type(self, type_id: int) -> AlignmentTypeResponse:
        response = await self.get(self._path(f"/alignments/types/{type_id}"))
        return self._parse(response, AlignmentTypeResponse)

    async def get_alignment_type_by_name(self, type_name: str) -> AlignmentTypeResponse:
        response = await self.get(self._path(f"/alignments/types/name/{type_name}"))
        return self._parse(response, AlignmentTypeResponse)

    async def create_alignment(
        self, request: AlignmentCreateRequest
    ) -> AlignmentResponse:
        response = await self.post(
            self._path("/alignments"),
            json=self._dump(request),
        )
        return self._parse(response, AlignmentResponse)

    async def upload_alignment_source_file(
        self,
        *,
        name: str,
        file: bytes | bytearray | BinaryIO | Path | str,
        filename: str | None = None,
        content_type: str | None = None,
        description: str | None = None,
        from_version: str | None = None,
        to_version: str | None = None,
        alignment_type_id: int | None = None,
        alignment_type_name: str | None = None,
        alignment_id: UUID | str | None = None,
    ) -> JSON_TYPE:
        """Upload a source file via the canonical alignment creation flow."""

        file_bytes: bytes
        target_filename = filename

        if isinstance(file, (bytes, bytearray)):
            file_bytes = bytes(file)
        elif isinstance(file, (str, Path)):
            path = Path(file)
            file_bytes = path.read_bytes()
            target_filename = target_filename or path.name
        elif hasattr(file, "read"):
            raw = file.read()
            if isinstance(raw, str):
                file_bytes = raw.encode("utf-8")
            else:
                file_bytes = bytes(raw)
            target_filename = target_filename or getattr(file, "name", None)
        else:
            raise TypeError("file must be bytes, Path/str, or a file-like object")

        upload_filename = target_filename or "alignment-source-file"
        upload_content_type = content_type or "application/octet-stream"

        files = {
            "file": (
                upload_filename,
                file_bytes,
                upload_content_type,
            )
        }

        form_data: dict[str, str] = {
            "name": name,
        }
        if description is not None:
            form_data["description"] = description
        if from_version is not None:
            form_data["from_version"] = from_version
        if to_version is not None:
            form_data["to_version"] = to_version
        if alignment_type_id is not None:
            form_data["alignment_type_id"] = str(alignment_type_id)
        if alignment_type_name is not None:
            form_data["alignment_type_name"] = alignment_type_name
        if alignment_id is not None:
            form_data["alignment_id"] = str(alignment_id)

        return await self._request(
            "POST",
            self._path("/post-processing/alignments/source-file"),
            files=files,
            data=form_data,
        )

    async def get_alignments(
        self, type_id: int | None = None
    ) -> list[AlignmentResponse]:
        params = {"typeId": type_id} if type_id is not None else None
        response = await self.get(self._path("/alignments"), params=params)
        return self._parse_list(response, AlignmentResponse)

    async def get_alignment(self, alignment_id: UUID) -> AlignmentResponse:
        response = await self.get(self._path(f"/alignments/{alignment_id}"))
        return self._parse(response, AlignmentResponse)

    async def get_alignment_by_name(self, name: str) -> AlignmentResponse:
        response = await self.get(self._path(f"/alignments/name/{name}"))
        return self._parse(response, AlignmentResponse)

    async def get_alignments_by_type(self, type_id: int) -> list[AlignmentResponse]:
        response = await self.get(self._path(f"/alignments/type/{type_id}/alignments"))
        return self._parse_list(response, AlignmentResponse)

    async def update_alignment(
        self,
        alignment_id: UUID,
        request: AlignmentUpdateRequest,
    ) -> AlignmentResponse:
        response = await self.patch(
            self._path(f"/alignments/{alignment_id}"),
            json=self._dump(request),
        )
        return self._parse(response, AlignmentResponse)

    async def close_alignment(
        self,
        alignment_id: UUID,
        request: AlignmentCloseRequest | None = None,
    ) -> AlignmentResponse:
        response = await self.post(
            self._path(f"/alignments/{alignment_id}/close"),
            json=self._dump(request or AlignmentCloseRequest()),
        )
        return self._parse(response, AlignmentResponse)

    async def delete_alignment(self, alignment_id: UUID) -> None:
        await self.delete(self._path(f"/alignments/{alignment_id}"))

    async def get_alignment_findings(self, alignment_id: UUID) -> list[FindingResponse]:
        response = await self.get(self._path(f"/alignments/{alignment_id}/findings"))
        return self._parse_list(response, FindingResponse)

    async def get_alignment_buckets(self, alignment_id: UUID) -> list[BucketResponse]:
        response = await self.get(self._path(f"/alignments/{alignment_id}/buckets"))
        return self._parse_list(response, BucketResponse)

    async def execute_alignment(
        self,
        alignment_id: UUID,
        request: AlignmentExecuteRequest,
    ) -> AlignmentExecuteResponse:
        response = await self.post(
            self._path(f"/alignments/{alignment_id}/execute"),
            json=self._dump(request),
        )
        return self._parse(response, AlignmentExecuteResponse)

    async def execute_alignment_bucket(
        self,
        alignment_id: UUID,
        bucket_id: UUID,
        request: AlignmentExecuteRequest,
    ) -> AlignmentExecuteResponse:
        response = await self.post(
            self._path(f"/alignments/{alignment_id}/buckets/{bucket_id}/execute"),
            json=self._dump(request),
        )
        return self._parse(response, AlignmentExecuteResponse)

    async def create_bucket(self, request: BucketCreateRequest) -> BucketResponse:
        response = await self.post(self._path("/buckets"), json=self._dump(request))
        return self._parse(response, BucketResponse)

    async def get_bucket(self, bucket_id: UUID) -> BucketResponse:
        response = await self.get(self._path(f"/buckets/{bucket_id}"))
        return self._parse(response, BucketResponse)

    async def get_bucket_detail(self, bucket_id: UUID) -> BucketDetailResponse:
        response = await self.get(self._path(f"/buckets/{bucket_id}/detail"))
        return self._parse(response, BucketDetailResponse)

    async def get_buckets_by_alignment(
        self, alignment_id: UUID
    ) -> list[BucketResponse]:
        response = await self.get(self._path(f"/buckets/alignment/{alignment_id}"))
        return self._parse_list(response, BucketResponse)

    async def get_buckets_by_status(self, status: str) -> list[BucketResponse]:
        response = await self.get(self._path(f"/buckets/status/{status}"))
        return self._parse_list(response, BucketResponse)

    async def get_buckets_by_type(self, bucket_type: str) -> list[BucketResponse]:
        response = await self.get(self._path(f"/buckets/type/{bucket_type}"))
        return self._parse_list(response, BucketResponse)

    async def update_bucket_status(
        self,
        bucket_id: UUID,
        request: BucketStatusUpdateRequest,
    ) -> BucketResponse:
        response = await self.patch(
            self._path(f"/buckets/{bucket_id}/status"),
            json=self._dump(request),
        )
        return self._parse(response, BucketResponse)

    async def delete_bucket(self, bucket_id: UUID) -> None:
        await self.delete(self._path(f"/buckets/{bucket_id}"))

    async def get_bucket_events(self, bucket_id: UUID) -> list[BucketEventResponse]:
        response = await self.get(self._path(f"/buckets/{bucket_id}/events"))
        return self._parse_list(response, BucketEventResponse)

    async def list_query_providers(self) -> list[QueryProviderResponse]:
        response = await self.get(self._path("/findings/match-queries/providers"))
        return self._parse_list(response, QueryProviderResponse)

    async def create_add_finding(self, request: AddFindingRequest) -> FindingResponse:
        response = await self.post(
            self._path("/findings/add"),
            json=self._dump(request),
        )
        return self._parse(response, FindingResponse)

    async def create_add_finding_bulk(
        self, request: list[AddFindingRequest]
    ) -> list[FindingResponse]:
        response = await self.post(
            self._path("/findings/add/bulk"),
            json=[self._dump(item) for item in request],
        )
        return self._parse_list(response, FindingResponse)

    async def create_retire_finding(
        self, request: RetireFindingRequest
    ) -> FindingResponse:
        response = await self.post(
            self._path("/findings/retire"),
            json=self._dump(request),
        )
        return self._parse(response, FindingResponse)

    async def create_retire_finding_bulk(
        self, request: list[RetireFindingRequest]
    ) -> list[FindingResponse]:
        response = await self.post(
            self._path("/findings/retire/bulk"),
            json=[self._dump(item) for item in request],
        )
        return self._parse_list(response, FindingResponse)

    async def create_update_finding(
        self, request: UpdateFindingRequest
    ) -> FindingResponse:
        response = await self.post(
            self._path("/findings/update"),
            json=self._dump(request),
        )
        return self._parse(response, FindingResponse)

    async def create_update_finding_bulk(
        self, request: list[UpdateFindingRequest]
    ) -> list[FindingResponse]:
        response = await self.post(
            self._path("/findings/update/bulk"),
            json=[self._dump(item) for item in request],
        )
        return self._parse_list(response, FindingResponse)

    async def update_finding_additional_data(
        self,
        finding_id: UUID,
        request: PatchFindingAdditionalDataRequest,
    ) -> FindingResponse:
        response = await self.patch(
            self._path(f"/findings/{finding_id}/additional-data"),
            json=self._dump(request),
        )
        return self._parse(response, FindingResponse)

    async def create_finding_match_query(
        self,
        finding_id: UUID,
        request: FindingMatchQueryCreateRequest,
        execute_query: bool = False,
    ) -> FindingMatchQueryResponse:
        response = await self._request(
            "POST",
            self._path(f"/findings/{finding_id}/match-queries"),
            json=self._dump(request),
            params={"executeQuery": execute_query},
        )
        return self._parse(response, FindingMatchQueryResponse)

    async def update_finding_match_query(
        self,
        match_query_id: int,
        request: FindingMatchQueryUpdateRequest,
    ) -> FindingMatchQueryResponse:
        response = await self.patch(
            self._path(f"/findings/match-queries/{match_query_id}"),
            json=self._dump(request),
        )
        return self._parse(response, FindingMatchQueryResponse)

    async def get_finding_match_queries(
        self, finding_id: UUID
    ) -> list[FindingMatchQueryResponse]:
        response = await self.get(self._path(f"/findings/{finding_id}/match-queries"))
        return self._parse_list(response, FindingMatchQueryResponse)

    async def get_finding_match_query(
        self, match_query_id: int
    ) -> FindingMatchQueryResponse:
        response = await self.get(
            self._path(f"/findings/match-queries/{match_query_id}")
        )
        return self._parse(response, FindingMatchQueryResponse)

    async def get_finding(self, finding_id: UUID) -> FindingResponse:
        response = await self.get(self._path(f"/findings/{finding_id}"))
        return self._parse(response, FindingResponse)

    async def move_finding(
        self, finding_id: UUID, new_bucket_id: UUID
    ) -> FindingResponse:
        response = await self._request(
            "PATCH",
            self._path(f"/findings/{finding_id}/bucket"),
            json={},
            params={"newBucketId": str(new_bucket_id)},
        )
        return self._parse(response, FindingResponse)

    async def get_findings_by_bucket(
        self,
        bucket_id: UUID,
        finding_type: str | None = None,
    ) -> list[FindingResponse]:
        params = {"finding_type": finding_type} if finding_type else None
        response = await self.get(
            self._path(f"/findings/bucket/{bucket_id}"),
            params=params,
        )
        return self._parse_list(response, FindingResponse)

    async def get_findings_by_alignment(
        self, alignment_id: UUID
    ) -> list[FindingResponse]:
        response = await self.get(self._path(f"/findings/alignment/{alignment_id}"))
        return self._parse_list(response, FindingResponse)

    async def delete_finding(self, finding_id: UUID) -> None:
        await self.delete(self._path(f"/findings/{finding_id}"))

    async def create_command_group(
        self, request: CommandGroupCreateRequest
    ) -> CommandGroupResponse:
        response = await self.post(
            self._path("/commands/groups"),
            json=self._dump(request),
        )
        return self._parse(response, CommandGroupResponse)

    async def update_command_group(
        self,
        group_id: UUID,
        request: CommandGroupUpdateRequest,
    ) -> CommandGroupResponse:
        response = await self.patch(
            self._path(f"/commands/groups/{group_id}"),
            json=self._dump(request),
        )
        return self._parse(response, CommandGroupResponse)

    async def get_command_group(self, group_id: UUID) -> CommandGroupResponse:
        response = await self.get(self._path(f"/commands/groups/{group_id}"))
        return self._parse(response, CommandGroupResponse)

    async def get_command_groups_by_bucket(
        self, bucket_id: UUID
    ) -> list[CommandGroupResponse]:
        response = await self.get(self._path(f"/commands/groups/bucket/{bucket_id}"))
        return self._parse_list(response, CommandGroupResponse)

    async def get_command_groups_by_task(
        self, task_id: UUID
    ) -> list[CommandGroupResponse]:
        response = await self.get(self._path(f"/commands/groups/task/{task_id}"))
        return self._parse_list(response, CommandGroupResponse)

    async def get_command_groups_by_status(
        self, status: str
    ) -> list[CommandGroupResponse]:
        response = await self.get(self._path(f"/commands/groups/status/{status}"))
        return self._parse_list(response, CommandGroupResponse)

    async def update_command_group_status(
        self,
        group_id: UUID,
        request: CommandGroupStatusUpdateRequest,
    ) -> CommandGroupResponse:
        response = await self.patch(
            self._path(f"/commands/groups/{group_id}/status"),
            json=self._dump(request),
        )
        return self._parse(response, CommandGroupResponse)

    async def remove_task_association_from_command_group(
        self, group_id: UUID
    ) -> CommandGroupResponse:
        response = await self.delete(self._path(f"/commands/groups/{group_id}/task"))
        return self._parse(response, CommandGroupResponse)

    async def push_command_group_to_task(
        self,
        group_id: UUID,
        request: PushCommandGroupRequest,
    ) -> CommandGroupResponse:
        response = await self.post(
            self._path(f"/commands/groups/{group_id}/task"),
            json=self._dump(request),
        )
        return self._parse(response, CommandGroupResponse)

    async def delete_command_group(self, group_id: UUID) -> None:
        await self.delete(self._path(f"/commands/groups/{group_id}"))

    async def create_command(
        self,
        group_id: UUID,
        request: CommandCreateRequest,
    ) -> CommandResponse:
        response = await self.post(
            self._path(f"/commands/groups/{group_id}/commands"),
            json=self._dump(request),
        )
        return self._parse(response, CommandResponse)

    async def update_command(
        self,
        command_id: UUID,
        request: UpdateCommandRequest,
    ) -> CommandResponse:
        response = await self.patch(
            self._path(f"/commands/{command_id}"),
            json=self._dump(request),
        )
        return self._parse(response, CommandResponse)

    async def get_commands_by_group(self, group_id: UUID) -> list[CommandResponse]:
        response = await self.get(self._path(f"/commands/groups/{group_id}/commands"))
        return self._parse_list(response, CommandResponse)

    async def get_command(self, command_id: UUID) -> CommandResponse:
        response = await self.get(self._path(f"/commands/{command_id}"))
        return self._parse(response, CommandResponse)

    async def create_command_findings_mapping(
        self,
        request: CommandFindingsMapCreateRequest,
    ) -> CommandFindingsMapResponse:
        response = await self.post(
            self._path("/commands/mappings"),
            json=self._dump(request),
        )
        return self._parse(response, CommandFindingsMapResponse)

    async def update_command_findings_mappings(
        self,
        command_id: UUID,
        request: UpdateCommandFindingsMappingRequest,
    ) -> list[CommandFindingsMapResponse]:
        response = await self.put(
            self._path(f"/commands/mappings/command/{command_id}"),
            json=self._dump(request),
        )
        return self._parse_list(response, CommandFindingsMapResponse)

    async def get_mappings_by_command(
        self, command_id: UUID
    ) -> list[CommandFindingsMapResponse]:
        response = await self.get(
            self._path(f"/commands/mappings/command/{command_id}")
        )
        return self._parse_list(response, CommandFindingsMapResponse)

    async def get_mappings_by_command_group(
        self, group_id: UUID
    ) -> list[CommandFindingsMapResponse]:
        response = await self.get(self._path(f"/commands/mappings/group/{group_id}"))
        return self._parse_list(response, CommandFindingsMapResponse)

    async def get_mappings_by_finding(
        self, finding_id: UUID
    ) -> list[CommandFindingsMapResponse]:
        response = await self.get(
            self._path(f"/commands/mappings/finding/{finding_id}")
        )
        return self._parse_list(response, CommandFindingsMapResponse)

    async def get_command_group_events(
        self, group_id: UUID
    ) -> list[CommandGroupEventResponse]:
        response = await self.get(self._path(f"/commands/groups/{group_id}/events"))
        return self._parse_list(response, CommandGroupEventResponse)

    async def get_command_events(self, command_id: UUID) -> list[CommandEventResponse]:
        response = await self.get(self._path(f"/commands/{command_id}/events"))
        return self._parse_list(response, CommandEventResponse)

    async def get_impact_assessment(
        self, assessment_id: int
    ) -> ImpactAssessmentWithDependenciesResponse:
        response = await self.get(self._path(f"/impact-assessments/{assessment_id}"))
        return self._parse(response, ImpactAssessmentWithDependenciesResponse)

    async def get_impact_assessment_by_finding(
        self, finding_id: UUID
    ) -> ImpactAssessmentWithDependenciesResponse:
        response = await self.get(
            self._path(f"/impact-assessments/finding/{finding_id}")
        )
        return self._parse(response, ImpactAssessmentWithDependenciesResponse)

    async def get_finding_dependency_graph(
        self, finding_id: UUID
    ) -> DependencyGraphResponse:
        response = await self.get(
            self._path(f"/impact-assessments/finding/{finding_id}/dependency-graph")
        )
        return self._parse(response, DependencyGraphResponse)

    async def get_bucket_dependency_graph_filtered(
        self,
        bucket_id: UUID,
        request: FindingDependencyGraphRequest,
    ) -> DependencyGraphResponse:
        response = await self.post(
            self._path(f"/impact-assessments/bucket/{bucket_id}/dependency-graph"),
            json=self._dump(request),
        )
        return self._parse(response, DependencyGraphResponse)

    async def get_command_graph(self, command_id: UUID) -> CommandGraphResponse:
        response = await self.get(
            self._path(f"/impact-assessments/command/{command_id}/command-graph")
        )
        return self._parse(response, CommandGraphResponse)

    async def get_command_group_graph(
        self, command_group_id: UUID
    ) -> CommandGraphResponse:
        response = await self.get(
            self._path(
                f"/impact-assessments/command-group/{command_group_id}/command-graph"
            )
        )
        return self._parse(response, CommandGraphResponse)

    async def get_background_task(self, job_id: str) -> dict[str, Any]:
        response = await self.get(self._path(f"/bgtasks/{job_id}"))
        if not isinstance(response, dict):
            raise ValueError("Expected dictionary payload for background task status")
        return response

    # ------------------------------------------------------------------
    # Generic post-processing endpoints
    # ------------------------------------------------------------------

    async def map_findings_bucket(
        self, bucket_id: UUID, request: MapFindingsRequest
    ) -> BucketFindingsResponse:
        """Map findings to DSG Core item IDs for a single bucket."""
        response = await self.post(
            self._path(f"/post-processing/buckets/{bucket_id}/map-findings"),
            json=self._dump(request),
        )
        return self._parse(response, BucketFindingsResponse)

    async def store_finding_dependencies_bucket(
        self, bucket_id: UUID
    ) -> BucketFindingsResponse:
        """Store finding dependencies from DSG Core for a single bucket."""
        response = await self.post(
            self._path(f"/post-processing/buckets/{bucket_id}/store-dependencies"),
            json={},
        )
        return self._parse(response, BucketFindingsResponse)

    async def dryrun_commands_bucket(self, bucket_id: UUID) -> BucketCommandsResponse:
        """Dry-run commands for a single bucket."""
        response = await self.post(
            self._path(f"/post-processing/buckets/{bucket_id}/dryrun-commands"),
            json={},
        )
        return self._parse(response, BucketCommandsResponse)

    async def map_findings_alignment(
        self, alignment_id: UUID, request: MapFindingsRequest
    ) -> AlignmentFindingsResponse:
        """Map findings to DSG Core item IDs for an entire alignment."""
        response = await self.post(
            self._path(f"/post-processing/alignments/{alignment_id}/map-findings"),
            json=self._dump(request),
        )
        return self._parse(response, AlignmentFindingsResponse)

    async def store_finding_dependencies_alignment(
        self, alignment_id: UUID
    ) -> AlignmentFindingsResponse:
        """Store finding dependencies from DSG Core for an entire alignment."""
        response = await self.post(
            self._path(
                f"/post-processing/alignments/{alignment_id}/store-dependencies"
            ),
            json={},
        )
        return self._parse(response, AlignmentFindingsResponse)

    async def dryrun_commands_alignment(
        self, alignment_id: UUID
    ) -> AlignmentCommandsResponse:
        """Dry-run commands for all buckets in an entire alignment."""
        response = await self.post(
            self._path(f"/post-processing/alignments/{alignment_id}/dryrun-commands"),
            json={},
        )
        return self._parse(response, AlignmentCommandsResponse)
