# Copyright (c) 2026 MergeGuide, Inc. All rights reserved.
# SPDX-License-Identifier: Proprietary
"""Build artifact provenance tracking for supply chain integrity.

This module implements provenance metadata capture and verification for build
artifacts produced during the PR gate evaluation pipeline. The data model is
inspired by SLSA (Supply-chain Levels for Software Artifacts) provenance
specifications, adapted to MergeGuide's evidence framework.

Provenance records capture:
- What was built (artifact digests, names, types)
- Who/what built it (builder identity, builder version)
- From what source (repository, commit, branch, PR)
- When it was built (timestamps)
- How it was built (build commands, environment)
- Integrity attestation (HMAC signatures, content-addressable hashing)

The provenance record is stored alongside the evidence artifact in S3,
enabling end-to-end traceability from source code to deployed artifact.

Usage:
    # Create provenance during PR gate evaluation
    provenance = BuildProvenance.create(
        repository="org/repo",
        commit_sha="abc123",
        branch="feature/new-endpoint",
        pr_number=42,
        builder_id="mergeguide-pr-gate",
    )

    # Add artifacts produced during evaluation
    provenance.add_artifact(
        name="evaluation-evidence.json",
        digest="sha256:abc123...",
        artifact_type="evidence",
    )

    # Finalize and serialize
    record = provenance.finalize()
    json_str = record.to_json()
"""

from __future__ import annotations

import hashlib
import json
import uuid
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from enum import StrEnum
from typing import Any


class ArtifactType(StrEnum):
    """Types of build artifacts tracked for provenance."""

    EVIDENCE = "evidence"
    SBOM = "sbom"
    SCAN_RESULT = "scan_result"
    POLICY_OUTPUT = "policy_output"
    SIGNATURE_REPORT = "signature_report"
    BUILD_LOG = "build_log"
    ATTESTATION = "attestation"


class ProvenanceStatus(StrEnum):
    """Status of provenance record collection."""

    PENDING = "pending"
    COMPLETE = "complete"
    FAILED = "failed"


@dataclass
class ProvenanceArtifact:
    """A single artifact tracked in the provenance record.

    Attributes:
        name: Human-readable artifact name (e.g., "evaluation-evidence.json").
        digest: Content-addressable digest in "algorithm:hex" format.
        artifact_type: Classification of the artifact.
        size_bytes: Size of the artifact in bytes (0 if unknown).
        storage_uri: URI where the artifact is stored (S3 key, file path).
        metadata: Additional key-value metadata for the artifact.
    """

    name: str
    digest: str
    artifact_type: ArtifactType
    size_bytes: int = 0
    storage_uri: str = ""
    metadata: dict[str, str] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "digest": self.digest,
            "artifact_type": self.artifact_type.value,
            "size_bytes": self.size_bytes,
            "storage_uri": self.storage_uri,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ProvenanceArtifact:
        return cls(
            name=data["name"],
            digest=data["digest"],
            artifact_type=ArtifactType(data.get("artifact_type", "evidence")),
            size_bytes=data.get("size_bytes", 0),
            storage_uri=data.get("storage_uri", ""),
            metadata=data.get("metadata", {}),
        )


@dataclass
class BuildSource:
    """Source code reference for the build.

    Attributes:
        repository: Full repository name (owner/repo).
        commit_sha: Git commit SHA that was evaluated.
        branch: Branch name.
        pr_number: Pull request number (0 if not from a PR).
        ref: Git ref (branch, tag, or SHA).
        diff_hash: SHA-256 hash of the diff content that was evaluated.
    """

    repository: str
    commit_sha: str
    branch: str = ""
    pr_number: int = 0
    ref: str = ""
    diff_hash: str = ""

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> BuildSource:
        return cls(**data)


@dataclass
class BuilderInfo:
    """Identity and version of the builder that produced artifacts.

    Attributes:
        builder_id: Unique identifier for the builder (e.g., "mergeguide-pr-gate").
        version: Builder version string.
        environment: Deployment environment (dev, test, prod).
        region: Cloud region where the build ran.
    """

    builder_id: str = "mergeguide-pr-gate"
    version: str = "1.0.0"
    environment: str = ""
    region: str = ""

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> BuilderInfo:
        return cls(**data)


@dataclass
class ProvenanceRecord:
    """Complete provenance record for a PR gate evaluation.

    This is the top-level data structure that captures full provenance
    metadata for a single PR gate run. It links source code, builder
    identity, produced artifacts, and timing information into a single
    verifiable record.

    Attributes:
        provenance_id: Unique identifier for this provenance record.
        schema_version: Version of the provenance schema.
        evaluation_id: ID of the corresponding evidence evaluation.
        status: Collection status (pending, complete, failed).
        created_at: ISO 8601 timestamp when collection started.
        completed_at: ISO 8601 timestamp when collection finished.
        source: Source code reference.
        builder: Builder identity and version.
        artifacts: List of artifacts produced during evaluation.
        digest_chain: Ordered list of content-addressable hashes for all artifacts.
        metadata: Additional key-value metadata.
    """

    provenance_id: str = field(default_factory=lambda: f"prov-{uuid.uuid4().hex[:12]}")
    schema_version: str = "1.0"
    evaluation_id: str = ""
    status: ProvenanceStatus = ProvenanceStatus.PENDING
    created_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
    completed_at: str = ""
    source: BuildSource = field(default_factory=lambda: BuildSource(repository="", commit_sha=""))
    builder: BuilderInfo = field(default_factory=BuilderInfo)
    artifacts: list[ProvenanceArtifact] = field(default_factory=list)
    digest_chain: list[str] = field(default_factory=list)
    metadata: dict[str, str] = field(default_factory=dict)

    @property
    def artifact_count(self) -> int:
        return len(self.artifacts)

    @property
    def total_artifact_bytes(self) -> int:
        return sum(a.size_bytes for a in self.artifacts)

    @property
    def is_complete(self) -> bool:
        return self.status == ProvenanceStatus.COMPLETE

    def add_artifact(
        self,
        name: str,
        digest: str,
        artifact_type: ArtifactType | str,
        size_bytes: int = 0,
        storage_uri: str = "",
        metadata: dict[str, str] | None = None,
    ) -> ProvenanceArtifact:
        """Add an artifact to the provenance record.

        Args:
            name: Artifact name.
            digest: Content digest in "algorithm:hex" format.
            artifact_type: Type classification.
            size_bytes: Size in bytes.
            storage_uri: Storage location URI.
            metadata: Additional metadata.

        Returns:
            The created ProvenanceArtifact.
        """
        if isinstance(artifact_type, str):
            artifact_type = ArtifactType(artifact_type)

        artifact = ProvenanceArtifact(
            name=name,
            digest=digest,
            artifact_type=artifact_type,
            size_bytes=size_bytes,
            storage_uri=storage_uri,
            metadata=metadata or {},
        )
        self.artifacts.append(artifact)
        self.digest_chain.append(digest)
        return artifact

    def finalize(self) -> ProvenanceRecord:
        """Mark the provenance record as complete.

        Sets the completed_at timestamp and status to COMPLETE.
        Computes a final chain digest over all artifact digests.

        Returns:
            Self for chaining.
        """
        self.completed_at = datetime.now(UTC).isoformat()
        self.status = ProvenanceStatus.COMPLETE
        return self

    def fail(self, reason: str = "") -> ProvenanceRecord:
        """Mark the provenance record as failed.

        Args:
            reason: Optional failure reason.

        Returns:
            Self for chaining.
        """
        self.completed_at = datetime.now(UTC).isoformat()
        self.status = ProvenanceStatus.FAILED
        if reason:
            self.metadata["failure_reason"] = reason
        return self

    def compute_chain_digest(self) -> str:
        """Compute a SHA-256 digest over the ordered artifact digests.

        Creates a summary hash by concatenating all content-addressable
        artifact digests and hashing the result. This enables integrity
        verification of the complete artifact set but is NOT a
        cryptographic hash chain (H(n) = hash(H(n-1) + data_n)).

        Returns:
            SHA-256 hex digest of the concatenated artifact digests.
        """
        if not self.digest_chain:
            return ""
        combined = "\n".join(self.digest_chain)
        return f"sha256:{hashlib.sha256(combined.encode()).hexdigest()}"

    def to_dict(self) -> dict[str, Any]:
        return {
            "provenance_id": self.provenance_id,
            "schema_version": self.schema_version,
            "evaluation_id": self.evaluation_id,
            "status": self.status.value,
            "created_at": self.created_at,
            "completed_at": self.completed_at,
            "source": self.source.to_dict(),
            "builder": self.builder.to_dict(),
            "artifacts": [a.to_dict() for a in self.artifacts],
            "digest_chain": self.digest_chain,
            "chain_digest": self.compute_chain_digest(),
            "artifact_count": self.artifact_count,
            "total_artifact_bytes": self.total_artifact_bytes,
            "metadata": self.metadata,
        }

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent, sort_keys=True)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ProvenanceRecord:
        source = BuildSource.from_dict(data.get("source", {}))
        builder = BuilderInfo.from_dict(data.get("builder", {}))
        artifacts = [ProvenanceArtifact.from_dict(a) for a in data.get("artifacts", [])]

        return cls(
            provenance_id=data.get("provenance_id", ""),
            schema_version=data.get("schema_version", "1.0"),
            evaluation_id=data.get("evaluation_id", ""),
            status=ProvenanceStatus(data.get("status", "pending")),
            created_at=data.get("created_at", ""),
            completed_at=data.get("completed_at", ""),
            source=source,
            builder=builder,
            artifacts=artifacts,
            digest_chain=data.get("digest_chain", []),
            metadata=data.get("metadata", {}),
        )

    @classmethod
    def from_json(cls, json_str: str) -> ProvenanceRecord:
        return cls.from_dict(json.loads(json_str))


class BuildProvenance:
    """Factory for creating and managing provenance records during PR gate runs.

    This class provides a high-level interface for creating provenance records
    as part of the PR gate evaluation pipeline. It handles the lifecycle of
    provenance collection from initialization through finalization.

    Usage:
        provenance = BuildProvenance.create(
            repository="org/repo",
            commit_sha="abc123",
            branch="main",
            pr_number=42,
        )

        # Add evidence artifact
        provenance.add_artifact(
            name="evidence.json",
            digest=compute_sha256(content),
            artifact_type=ArtifactType.EVIDENCE,
            size_bytes=len(content),
            storage_uri="s3://bucket/key",
        )

        # Finalize
        record = provenance.finalize()
    """

    def __init__(self, record: ProvenanceRecord) -> None:
        self._record = record

    @classmethod
    def create(
        cls,
        repository: str,
        commit_sha: str,
        branch: str = "",
        pr_number: int = 0,
        builder_id: str = "mergeguide-pr-gate",
        builder_version: str = "1.0.0",
        environment: str = "",
        region: str = "",
        evaluation_id: str = "",
        diff_content: str = "",
    ) -> BuildProvenance:
        """Create a new provenance tracking session.

        Args:
            repository: Full repository name (owner/repo).
            commit_sha: Git commit SHA being evaluated.
            branch: Branch name.
            pr_number: PR number.
            builder_id: Builder identifier.
            builder_version: Builder version.
            environment: Deployment environment.
            region: Cloud region.
            evaluation_id: Corresponding evaluation ID.
            diff_content: Diff content for hashing.

        Returns:
            BuildProvenance instance ready for artifact tracking.
        """
        diff_hash = ""
        if diff_content:
            diff_hash = f"sha256:{hashlib.sha256(diff_content.encode()).hexdigest()}"

        source = BuildSource(
            repository=repository,
            commit_sha=commit_sha,
            branch=branch,
            pr_number=pr_number,
            ref=commit_sha,
            diff_hash=diff_hash,
        )

        builder = BuilderInfo(
            builder_id=builder_id,
            version=builder_version,
            environment=environment,
            region=region,
        )

        record = ProvenanceRecord(
            evaluation_id=evaluation_id,
            source=source,
            builder=builder,
        )

        return cls(record)

    @property
    def record(self) -> ProvenanceRecord:
        return self._record

    @property
    def provenance_id(self) -> str:
        return self._record.provenance_id

    def add_artifact(
        self,
        name: str,
        digest: str,
        artifact_type: ArtifactType | str,
        size_bytes: int = 0,
        storage_uri: str = "",
        metadata: dict[str, str] | None = None,
    ) -> ProvenanceArtifact:
        """Add an artifact to the provenance record.

        Delegates to ProvenanceRecord.add_artifact.
        """
        return self._record.add_artifact(
            name=name,
            digest=digest,
            artifact_type=artifact_type,
            size_bytes=size_bytes,
            storage_uri=storage_uri,
            metadata=metadata,
        )

    def add_evidence_artifact(
        self,
        evidence_json: str,
        evidence_key: str,
    ) -> ProvenanceArtifact:
        """Convenience method for adding an evidence artifact.

        Args:
            evidence_json: Serialized evidence JSON content.
            evidence_key: S3 key or storage path.

        Returns:
            The created ProvenanceArtifact.
        """
        digest = f"sha256:{hashlib.sha256(evidence_json.encode()).hexdigest()}"
        return self.add_artifact(
            name="evaluation-evidence.json",
            digest=digest,
            artifact_type=ArtifactType.EVIDENCE,
            size_bytes=len(evidence_json.encode()),
            storage_uri=evidence_key,
        )

    def add_scan_result(
        self,
        scan_type: str,
        result_json: str,
        storage_uri: str = "",
    ) -> ProvenanceArtifact:
        """Convenience method for adding a scan result artifact.

        Args:
            scan_type: Type of scan (e.g., "build", "iac", "cicd").
            result_json: Serialized scan result.
            storage_uri: Storage location.

        Returns:
            The created ProvenanceArtifact.
        """
        digest = f"sha256:{hashlib.sha256(result_json.encode()).hexdigest()}"
        return self.add_artifact(
            name=f"{scan_type}-scan-result.json",
            digest=digest,
            artifact_type=ArtifactType.SCAN_RESULT,
            size_bytes=len(result_json.encode()),
            storage_uri=storage_uri,
            metadata={"scan_type": scan_type},
        )

    def add_signature_report(
        self,
        report_json: str,
        storage_uri: str = "",
    ) -> ProvenanceArtifact:
        """Convenience method for adding a signature verification report.

        Args:
            report_json: Serialized signature scan result.
            storage_uri: Storage location.

        Returns:
            The created ProvenanceArtifact.
        """
        digest = f"sha256:{hashlib.sha256(report_json.encode()).hexdigest()}"
        return self.add_artifact(
            name="commit-signature-report.json",
            digest=digest,
            artifact_type=ArtifactType.SIGNATURE_REPORT,
            size_bytes=len(report_json.encode()),
            storage_uri=storage_uri,
        )

    def finalize(self) -> ProvenanceRecord:
        """Finalize the provenance record.

        Returns:
            The finalized ProvenanceRecord.
        """
        return self._record.finalize()

    def fail(self, reason: str = "") -> ProvenanceRecord:
        """Mark provenance collection as failed.

        Args:
            reason: Failure reason.

        Returns:
            The failed ProvenanceRecord.
        """
        return self._record.fail(reason)


def compute_artifact_digest(content: str | bytes) -> str:
    """Compute SHA-256 digest for artifact content.

    Args:
        content: String or bytes content to hash.

    Returns:
        Digest in "sha256:hex" format.
    """
    if isinstance(content, str):
        content = content.encode()
    return f"sha256:{hashlib.sha256(content).hexdigest()}"
