"""SOC2 Trust Services Criteria control mapping.

Maps AttestMCP audit evidence to specific SOC2 controls.
This module generates the structured evidence an auditor needs
to assess whether AI agent memory operations meet compliance
requirements.

Covers:
  CC6.1 — Logical access security (identity on every record)
  CC6.3 — Access authorization (policy decisions)
  CC7.2 — System monitoring (continuous audit trail)
  CC8.1 — Change management (tamper-evident chain)
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from attestmcp.models import (
    AuditRecord,
    ChainVerification,
    IdentitySource,
    PolicyDecision,
    TrustLevel,
)

if TYPE_CHECKING:
    from datetime import datetime


@dataclass(frozen=True)
class ControlEvidence:
    """Evidence that a specific SOC2 control is satisfied."""

    control_id: str
    control_name: str
    description: str
    evidence_points: list[str] = field(default_factory=list)
    satisfied: bool = True
    gaps: list[str] = field(default_factory=list)


@dataclass(frozen=True)
class SOC2Report:
    """Complete SOC2 evidence report for a given audit period."""

    period_start: datetime | None
    period_end: datetime | None
    total_records: int
    chain_verification: ChainVerification
    controls: list[ControlEvidence]
    summary_stats: dict[str, int | str]


def map_controls(
    records: list[AuditRecord],
    chain_verification: ChainVerification,
    *,
    period_start: datetime | None = None,
    period_end: datetime | None = None,
) -> SOC2Report:
    """Map audit records to SOC2 Trust Services Criteria.

    Args:
        records: Audit records from the evidence period.
        chain_verification: Result of verify_chain() for the period.
        period_start: Start of the evidence window.
        period_end: End of the evidence window.

    Returns:
        Complete SOC2 evidence report with control mappings.
    """
    stats = _compute_stats(records)

    controls = [
        _map_cc61(records, stats),
        _map_cc63(records, stats),
        _map_cc72(records, stats),
        _map_cc81(chain_verification, stats),
    ]

    return SOC2Report(
        period_start=period_start,
        period_end=period_end,
        total_records=len(records),
        chain_verification=chain_verification,
        controls=controls,
        summary_stats=stats,
    )


def _compute_stats(records: list[AuditRecord]) -> dict[str, int | str]:
    """Compute summary statistics from audit records."""
    if not records:
        return {
            "total": 0,
            "with_identity": 0,
            "anonymous": 0,
            "verified_identity": 0,
            "claimed_identity": 0,
            "tools_call": 0,
            "unique_users": 0,
            "unique_agents": 0,
            "unique_sessions": 0,
            "policy_allow": 0,
            "policy_deny": 0,
            "errors": 0,
            "avg_duration_ms": "0",
        }

    with_identity = sum(
        1
        for r in records
        if r.agent_identity.source != IdentitySource.ANONYMOUS
    )
    verified = sum(
        1
        for r in records
        if r.agent_identity.trust_level == TrustLevel.VERIFIED
    )
    claimed = sum(
        1
        for r in records
        if r.agent_identity.trust_level == TrustLevel.CLAIMED
    )
    tools_call = sum(1 for r in records if r.method == "tools/call")
    unique_users = len({
        r.agent_identity.user
        for r in records
        if r.agent_identity.user
    })
    unique_agents = len({
        r.agent_identity.agent_name
        for r in records
        if r.agent_identity.agent_name
    })
    unique_sessions = len({
        r.session_id for r in records if r.session_id
    })
    policy_allow = sum(
        1 for r in records if r.policy_decision == PolicyDecision.ALLOW
    )
    policy_deny = sum(
        1 for r in records if r.policy_decision == PolicyDecision.DENY
    )
    errors = sum(1 for r in records if r.error is not None)
    durations = [r.duration_ms for r in records if r.duration_ms > 0]
    avg_dur = (
        str(round(sum(durations) / len(durations)))
        if durations
        else "0"
    )

    return {
        "total": len(records),
        "with_identity": with_identity,
        "anonymous": len(records) - with_identity,
        "verified_identity": verified,
        "claimed_identity": claimed,
        "tools_call": tools_call,
        "unique_users": unique_users,
        "unique_agents": unique_agents,
        "unique_sessions": unique_sessions,
        "policy_allow": policy_allow,
        "policy_deny": policy_deny,
        "errors": errors,
        "avg_duration_ms": avg_dur,
    }


def _map_cc61(
    records: list[AuditRecord], stats: dict[str, int | str]
) -> ControlEvidence:
    """CC6.1 — Logical and Physical Access Controls.

    The entity implements logical access security software,
    infrastructure, and architectures over protected information
    assets to protect them from security events.
    """
    evidence_points: list[str] = []
    gaps: list[str] = []

    total = stats["total"]
    if isinstance(total, int) and total > 0:
        with_id = stats["with_identity"]
        assert isinstance(with_id, int)
        pct = round(with_id / total * 100, 1)
        evidence_points.append(
            f"{with_id}/{total} operations ({pct}%) include agent identity"
        )

    verified = stats["verified_identity"]
    if isinstance(verified, int) and verified > 0:
        evidence_points.append(
            f"{verified} operations have cryptographically verified identity (JWKS)"
        )

    users = stats["unique_users"]
    agents = stats["unique_agents"]
    if isinstance(users, int) and users > 0:
        evidence_points.append(f"{users} unique user(s) identified")
    if isinstance(agents, int) and agents > 0:
        evidence_points.append(f"{agents} unique agent(s) identified")

    anon = stats["anonymous"]
    if isinstance(anon, int) and isinstance(total, int) and total > 0 and anon / total > 0.2:
            gaps.append(
                f"{anon}/{total} operations are anonymous — "
                "consider requiring authentication"
            )

    satisfied = len(gaps) == 0 and len(evidence_points) > 0

    return ControlEvidence(
        control_id="CC6.1",
        control_name="Logical and Physical Access Controls",
        description=(
            "Every AI agent memory operation is logged with "
            "the requesting identity. Identity is extracted from "
            "JWT bearer tokens and recorded as verified (JWKS-validated) "
            "or claimed (decoded without signature verification)."
        ),
        evidence_points=evidence_points,
        satisfied=satisfied,
        gaps=gaps,
    )


def _map_cc63(
    records: list[AuditRecord], stats: dict[str, int | str]
) -> ControlEvidence:
    """CC6.3 — Access Authorization.

    The entity authorizes, modifies, or removes access to data,
    software, functions, and other protected information assets
    based on roles.
    """
    evidence_points: list[str] = []
    gaps: list[str] = []

    allow = stats["policy_allow"]
    deny = stats["policy_deny"]
    if isinstance(allow, int):
        evidence_points.append(
            f"{allow} operations evaluated as ALLOW by policy engine"
        )
    if isinstance(deny, int) and deny > 0:
        evidence_points.append(
            f"{deny} operations evaluated as DENY by policy engine"
        )

    # Check for role-based access patterns
    roles_seen = {
        r.agent_identity.role
        for r in records
        if r.agent_identity.role
    }
    if roles_seen:
        evidence_points.append(
            f"Role-based access observed: {', '.join(sorted(roles_seen))}"
        )
    else:
        gaps.append(
            "No role claims detected in agent tokens — "
            "role-based access control not yet in effect"
        )

    return ControlEvidence(
        control_id="CC6.3",
        control_name="Access Authorization",
        description=(
            "Policy decisions (allow/deny/escalate/redact) are recorded "
            "for every memory operation. Role claims from agent tokens "
            "are captured to support role-based access reviews."
        ),
        evidence_points=evidence_points,
        satisfied=True,
        gaps=gaps,
    )


def _map_cc72(
    records: list[AuditRecord], stats: dict[str, int | str]
) -> ControlEvidence:
    """CC7.2 — System Monitoring.

    The entity monitors system components and the operation of
    those components for anomalies that are indicative of malicious
    acts, natural disasters, and errors affecting the entity's
    ability to meet its objectives.
    """
    evidence_points: list[str] = []
    gaps: list[str] = []

    total = stats["total"]
    if isinstance(total, int) and total > 0:
        evidence_points.append(
            f"Continuous monitoring active: {total} operations recorded"
        )

    tools = stats["tools_call"]
    if isinstance(tools, int) and tools > 0:
        evidence_points.append(
            f"{tools} tool invocations captured with request/response hashes"
        )

    sessions = stats["unique_sessions"]
    if isinstance(sessions, int) and sessions > 0:
        evidence_points.append(
            f"{sessions} unique MCP sessions tracked"
        )

    errors = stats["errors"]
    if isinstance(errors, int) and errors > 0:
        evidence_points.append(
            f"{errors} error(s) recorded with upstream failure details"
        )

    avg_ms = stats["avg_duration_ms"]
    if avg_ms != "0":
        evidence_points.append(
            f"Average proxy latency: {avg_ms}ms"
        )

    if isinstance(total, int) and total == 0:
        gaps.append("No operations recorded in this period")

    return ControlEvidence(
        control_id="CC7.2",
        control_name="System Monitoring",
        description=(
            "Every MCP memory operation is intercepted and logged "
            "in real-time. Request parameter hashes and response "
            "hashes provide non-repudiation. Latency and error "
            "metrics enable anomaly detection."
        ),
        evidence_points=evidence_points,
        satisfied=isinstance(total, int) and total > 0,
        gaps=gaps,
    )


def _map_cc81(
    chain_verification: ChainVerification,
    stats: dict[str, int | str],
) -> ControlEvidence:
    """CC8.1 — Change Management.

    The entity authorizes, designs, develops or acquires, configures,
    documents, tests, approves, and implements changes to
    infrastructure, data, software, and procedures.
    """
    evidence_points: list[str] = []
    gaps: list[str] = []

    if chain_verification.valid:
        evidence_points.append(
            f"Hash chain integrity verified: {chain_verification.records_checked} "
            f"records, zero breaks detected"
        )
        evidence_points.append(
            "SHA-256 chain links each record to its predecessor — "
            "any modification, deletion, or reordering is detectable"
        )
    else:
        gaps.append(
            f"CHAIN INTEGRITY FAILURE at record "
            f"{chain_verification.break_at_record_id}"
        )
        gaps.append(
            f"Expected hash: {chain_verification.expected_hash}"
        )
        gaps.append(
            f"Actual hash: {chain_verification.actual_hash}"
        )

    evidence_points.append(
        "Audit records are append-only — no UPDATE or DELETE "
        "operations are permitted on the evidence store"
    )

    return ControlEvidence(
        control_id="CC8.1",
        control_name="Change Management",
        description=(
            "Audit records are stored in a tamper-evident hash chain. "
            "Each record's chain_hash = SHA-256(previous_hash + canonical_json(record)). "
            "Chain verification runs automatically and can be triggered on demand."
        ),
        evidence_points=evidence_points,
        satisfied=chain_verification.valid,
        gaps=gaps,
    )
