"""Data flow mapper.

Generates a visual map of where personal data flows from your AI agent
to external providers. Useful for NDPA cross-border transfer documentation
and GDPR records of processing.
"""

from .audit import AuditLog


_PROVIDER_COUNTRIES = {
    "openai": "United States",
    "anthropic": "United States",
    "google": "United States / EU",
    "azure": "Configurable (check region)",
    "aws": "Configurable (check region)",
    "custom": "Unknown — verify",
}


def generate_map(audit_log: AuditLog) -> str:
    """Generate a Markdown data flow map from audit log data.

    Returns:
        Markdown string showing data flows.
    """
    entries = audit_log.query()
    if not entries:
        return "No data flows recorded. Run your agent with agent-shield to collect data."

    # Aggregate flows
    flows: dict[str, dict] = {}
    for e in entries:
        provider = e.get("provider", "unknown")
        if provider not in flows:
            flows[provider] = {
                "models": set(),
                "purposes": set(),
                "pii_types": set(),
                "call_count": 0,
                "total_tokens": 0,
                "users": set(),
            }
        flows[provider]["models"].add(e.get("model", "unknown"))
        flows[provider]["purposes"].add(e.get("purpose", "unknown"))
        for pii_type in e.get("pii_detected", {}):
            flows[provider]["pii_types"].add(pii_type)
        flows[provider]["call_count"] += 1
        flows[provider]["total_tokens"] += e.get("tokens_in", 0) + e.get("tokens_out", 0)
        uid = e.get("user_id")
        if uid:
            flows[provider]["users"].add(uid)

    # Build markdown
    lines = [
        "# Data Flow Map",
        "",
        "Generated from agent-shield audit log.",
        "",
        "## Summary",
        "",
        f"- **Total API calls:** {sum(f['call_count'] for f in flows.values())}",
        f"- **External providers:** {len(flows)}",
        f"- **Unique users:** {len(set().union(*(f['users'] for f in flows.values())))}",
        "",
        "## Flows",
        "",
        "| Provider | Country | Models | Purposes | PII Types Sent | Calls | Tokens |",
        "|----------|---------|--------|----------|----------------|-------|--------|",
    ]

    for provider, data in sorted(flows.items(), key=lambda x: -x[1]["call_count"]):
        country = _PROVIDER_COUNTRIES.get(provider, "Unknown")
        models = ", ".join(sorted(data["models"]))
        purposes = ", ".join(sorted(data["purposes"]))
        pii_types = ", ".join(sorted(data["pii_types"])) if data["pii_types"] else "None detected"
        lines.append(
            f"| {provider} | {country} | {models} | {purposes} | "
            f"{pii_types} | {data['call_count']} | {data['total_tokens']:,} |"
        )

    lines.extend([
        "",
        "## Cross-Border Transfer Documentation",
        "",
    ])

    for provider, data in sorted(flows.items()):
        country = _PROVIDER_COUNTRIES.get(provider, "Unknown")
        if country in ("Unknown — verify", "Configurable (check region)"):
            status = "VERIFY — country not confirmed"
        elif "United States" in country:
            status = "REQUIRES TRANSFER SAFEGUARDS"
        else:
            status = "ASSESS — check adequacy status"

        lines.extend([
            f"### {provider}",
            "",
            f"- **Destination country:** {country}",
            f"- **Transfer status:** {status}",
            f"- **Data types transferred:** {', '.join(sorted(data['pii_types'])) if data['pii_types'] else 'No PII detected (verify manually)'}",
            f"- **Volume:** {data['call_count']} calls, {data['total_tokens']:,} tokens",
            f"- **Purpose:** {', '.join(sorted(data['purposes']))}",
            "- **DPA in place:** [ ] Yes / [ ] No — **ACTION REQUIRED if No**",
            "- **Transfer mechanism:** [ ] SCCs / [ ] Adequacy / [ ] Consent / [ ] Other",
            "",
        ])

    lines.extend([
        "---",
        "",
        "*Generated by [agent-shield](https://github.com/Thezenmonster/agent-shield).",
        "For cross-border transfer compliance guidance, see",
        "[Cross-Border Data Transfers from Nigeria](https://januscompliance.co.uk/blog/cross-border-data-transfers-nigeria-how-to-comply).*",
    ])

    return "\n".join(lines)


def generate_mermaid(audit_log: AuditLog, app_name: str = "Your App") -> str:
    """Generate a Mermaid diagram of data flows.

    Returns:
        Mermaid diagram string (paste into any Mermaid renderer).
    """
    entries = audit_log.query()
    if not entries:
        return "graph LR\n    A[No data flows recorded]"

    providers = set()
    pii_by_provider: dict[str, set] = {}
    for e in entries:
        p = e.get("provider", "unknown")
        providers.add(p)
        if p not in pii_by_provider:
            pii_by_provider[p] = set()
        for pii_type in e.get("pii_detected", {}):
            pii_by_provider[p].add(pii_type)

    lines = ["graph LR", f"    APP[{app_name}]"]
    for i, provider in enumerate(sorted(providers)):
        node_id = f"P{i}"
        country = _PROVIDER_COUNTRIES.get(provider, "?")
        pii = ", ".join(sorted(pii_by_provider.get(provider, set()))) or "data"
        lines.append(f"    {node_id}[{provider}<br/>{country}]")
        lines.append(f"    APP -->|{pii}| {node_id}")

    return "\n".join(lines)
