# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Audit Trail Skill

Log and query every data access, tool invocation, and configuration change.
HIPAA §164.312(b) audit controls.

Wraps core/audit.py (SQLite, hash-chained) with user-facing tools.
Fallback: file-based JSON audit log when core audit unavailable.

Dependencies: None (sqlite3 is stdlib)
"""

import hashlib
import json
import logging
import os
from datetime import datetime, timedelta
from pathlib import Path


def _get_data_dir():
    """Get data directory (tenant-scoped in Reflection, default in standalone)."""
    try:
        from familiar.core import paths

        return paths.DATA_DIR
    except ImportError:
        return Path.home() / ".familiar" / "data"


try:
    from familiar.core.audit import (
        AuditEvent,
        get_audit_store,
    )

    CORE_AUDIT_AVAILABLE = True
except ImportError:
    CORE_AUDIT_AVAILABLE = False

try:
    from familiar.core.utils import generate_id
except ImportError:
    generate_id = None

logger = logging.getLogger(__name__)

AUDIT_DIR = _get_data_dir() / "audit"
AUDIT_DB = AUDIT_DIR / "audit.db"
AUDIT_JSON_LOG = AUDIT_DIR / "audit_log.jsonl"


def _gen_id():
    if generate_id:
        return generate_id()
    import secrets

    return secrets.token_hex(6)


def _ensure_dir():
    AUDIT_DIR.mkdir(parents=True, exist_ok=True)


# === Fallback File-Based Audit (when core unavailable) ===

_last_hash = None


def _compute_chain_hash(entry_json, prev_hash=None):
    """HMAC-SHA-256 hash chain: HMAC(key, prev_hash || entry) for tamper resistance."""
    import hmac as _hmac

    key = os.environ.get("FAMILIAR_AUDIT_KEY", "").encode("utf-8") or b"familiar-audit-default-key"
    data = (prev_hash or "genesis") + entry_json
    return _hmac.new(key, data.encode("utf-8"), hashlib.sha256).hexdigest()


def _append_audit_entry(entry):
    """Append to JSONL audit log with hash chain."""
    global _last_hash
    _ensure_dir()

    if _last_hash is None:
        # Read last hash from file
        if AUDIT_JSON_LOG.exists():
            try:
                with open(AUDIT_JSON_LOG, "rb") as f:
                    f.seek(0, 2)
                    size = f.tell()
                    if size > 0:
                        f.seek(max(0, size - 4096))
                        lines = f.read().decode("utf-8", errors="replace").strip().split("\n")
                        if lines:
                            last = json.loads(lines[-1])
                            _last_hash = last.get("chain_hash", None)
            except Exception as e:
                logger.debug("Error suppressed: %s", e)
    entry_json = json.dumps(entry, sort_keys=True, default=str)
    chain_hash = _compute_chain_hash(entry_json, _last_hash)
    entry["chain_hash"] = chain_hash
    _last_hash = chain_hash

    with open(AUDIT_JSON_LOG, "a", encoding="utf-8") as f:
        f.write(json.dumps(entry, default=str) + "\n")


def _search_jsonl(
    start=None, end=None, user=None, action=None, resource=None, severity=None, limit=50
):
    """Search the JSONL audit log."""
    if not AUDIT_JSON_LOG.exists():
        return []
    results = []
    with open(AUDIT_JSON_LOG, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
            except json.JSONDecodeError as e:
                logger.debug("JSON parse error suppressed: %s", e)
                continue
            if start and entry.get("timestamp", "") < start:
                continue
            if end and entry.get("timestamp", "") > end:
                continue
            if user and user.lower() not in entry.get("user_id", "").lower():
                continue
            if action and action.lower() != entry.get("action", "").lower():
                continue
            if resource and resource.lower() not in entry.get("resource", "").lower():
                continue
            if severity and severity.lower() != entry.get("severity", "").lower():
                continue
            results.append(entry)
    return results[-limit:]


# === Tool Handlers ===


def log_event(data):
    """Log an audit event (tool execution, data access, config change, auth)."""
    action = data.get("action", "tool_executed")
    user_id = data.get("user_id", "system")
    resource = data.get("resource", "")
    detail = data.get("detail", "")
    severity = data.get("severity", "info")
    metadata = data.get("metadata", {})

    event_id = _gen_id()

    if CORE_AUDIT_AVAILABLE:
        try:
            store = get_audit_store()
            event = AuditEvent(
                event_id=event_id,
                timestamp=datetime.now(),
                user_id=user_id,
                action=action,
                resource=resource,
                detail=detail,
                severity=severity,
                metadata=metadata,
            )
            store.log(event)
            return f"✅ Audit event logged: [{severity}] {action} on {resource} by {user_id} [{event_id}]"
        except Exception as e:
            logger.warning(f"Core audit failed, using fallback: {e}")

    # Fallback
    entry = {
        "event_id": event_id,
        "timestamp": datetime.now().isoformat(),
        "user_id": user_id,
        "action": action,
        "resource": resource,
        "detail": detail,
        "severity": severity,
        "metadata": metadata,
    }
    _append_audit_entry(entry)
    return f"✅ Audit event logged: [{severity}] {action} on {resource} by {user_id} [{event_id}]"


def search_audit_log(data):
    """Search audit events by date range, user, action, resource, severity."""
    start = data.get("start_date")
    end = data.get("end_date")
    user = data.get("user_id")
    action = data.get("action")
    resource = data.get("resource")
    severity = data.get("severity")
    limit = min(data.get("limit", 50), 500)

    if CORE_AUDIT_AVAILABLE:
        try:
            store = get_audit_store()
            start_dt = datetime.fromisoformat(start) if start else None
            end_dt = datetime.fromisoformat(end) if end else None
            events = store.query(
                start=start_dt, end=end_dt, actor_id=user, action=action, limit=limit
            )
            results = [e.to_dict() for e in events]
        except Exception as e:
            logger.warning(f"Core audit query failed: {e}")
            results = _search_jsonl(start, end, user, action, resource, severity, limit)
    else:
        results = _search_jsonl(start, end, user, action, resource, severity, limit)

    if not results:
        return "No audit events found matching criteria."

    lines = [f"📋 Audit Log ({len(results)} events):\n"]
    for e in results:
        ts = str(e.get("timestamp", ""))[:19]
        sev = e.get("severity", "info").upper()
        act = e.get("action", "")
        usr = e.get("user_id", "")
        res = e.get("resource", "")
        det = e.get("detail", "")[:60]
        lines.append(f"  [{sev}] {ts} {usr}: {act} → {res}")
        if det:
            lines.append(f"         {det}")

    return "\n".join(lines)


def compliance_report(data):
    """Generate HIPAA compliance audit report for a date range."""
    start = data.get("start_date", (datetime.now() - timedelta(days=30)).strftime("%Y-%m-%d"))
    end = data.get("end_date", datetime.now().strftime("%Y-%m-%d"))

    all_events = _search_jsonl(start, end, limit=10000) if not CORE_AUDIT_AVAILABLE else []

    if CORE_AUDIT_AVAILABLE:
        try:
            store = get_audit_store()
            start_dt = datetime.fromisoformat(start)
            end_dt = datetime.fromisoformat(end)
            events = store.query(start=start_dt, end=end_dt, limit=10000)
            all_events = [e.to_dict() for e in events]
        except Exception as e:
            logger.debug("Error suppressed: %s", e)
    # Aggregate statistics
    total = len(all_events)
    by_action = {}
    by_user = {}
    by_severity = {}
    phi_accesses = 0
    data_exports = 0
    auth_failures = 0

    for e in all_events:
        act = e.get("action", "unknown")
        usr = e.get("user_id", "unknown")
        sev = e.get("severity", "info")
        by_action[act] = by_action.get(act, 0) + 1
        by_user[usr] = by_user.get(usr, 0) + 1
        by_severity[sev] = by_severity.get(sev, 0) + 1
        if "phi" in act.lower() or "phi" in e.get("detail", "").lower():
            phi_accesses += 1
        if "export" in act.lower():
            data_exports += 1
        if "failed" in act.lower() or "denied" in act.lower():
            auth_failures += 1

    # Chain integrity
    integrity = "✅ Verified" if _verify_chain() else "⚠️ Verification needed"

    lines = [
        "📋 HIPAA Compliance Audit Report",
        f"Period: {start} to {end}",
        f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}",
        f"Chain Integrity: {integrity}",
        "",
        "SUMMARY:",
        f"  Total events: {total}",
        f"  PHI accesses: {phi_accesses}",
        f"  Data exports: {data_exports}",
        f"  Auth failures: {auth_failures}",
        "",
        "BY SEVERITY:",
    ]
    for sev in ["critical", "warning", "info"]:
        if sev in by_severity:
            lines.append(f"  {sev}: {by_severity[sev]}")

    lines.append("\nBY ACTION (top 15):")
    for act, count in sorted(by_action.items(), key=lambda x: -x[1])[:15]:
        lines.append(f"  {act}: {count}")

    lines.append("\nBY USER:")
    for usr, count in sorted(by_user.items(), key=lambda x: -x[1]):
        lines.append(f"  {usr}: {count}")

    # HIPAA checklist
    lines.extend(
        [
            "\nHIPAA CONTROLS CHECKLIST:",
            f"  ✅ §164.312(b) Audit controls: Active ({total} events logged)",
            f"  {'✅' if phi_accesses > 0 else '⚠️'} §164.312(a)(1) Access controls: {'PHI access tracked' if phi_accesses else 'No PHI access detected'}",
            "  ✅ §164.312(c)(2) Integrity mechanism: Hash-chain audit log",
            f"  {'✅' if auth_failures == 0 else '⚠️'} §164.312(d) Authentication: {auth_failures} failures",
        ]
    )

    # Save report
    _ensure_dir()
    report_path = AUDIT_DIR / f"compliance_report_{start}_{end}.txt"
    report_path.write_text("\n".join(lines), encoding="utf-8")
    lines.append(f"\nReport saved: {report_path}")

    return "\n".join(lines)


def _verify_chain():
    """Verify hash chain integrity of JSONL audit log."""
    if not AUDIT_JSON_LOG.exists():
        return True
    prev_hash = None
    try:
        with open(AUDIT_JSON_LOG, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                entry = json.loads(line)
                stored_hash = entry.pop("chain_hash", None)
                entry_json = json.dumps(entry, sort_keys=True, default=str)
                expected = _compute_chain_hash(entry_json, prev_hash)
                if stored_hash and stored_hash != expected:
                    return False
                prev_hash = stored_hash
        return True
    except Exception:
        return False


def data_access_log(data):
    """Log and query data access events for a specific record or data store."""
    action = data.get("action", "query")

    if action == "log":
        return log_event(
            {
                "action": "data_view",
                "user_id": data.get("user_id", "system"),
                "resource": data.get("resource", ""),
                "detail": data.get(
                    "detail",
                    f"Accessed {data.get('record_type', 'record')} {data.get('record_id', '')}",
                ),
                "severity": "info",
                "metadata": {
                    "record_type": data.get("record_type", ""),
                    "record_id": data.get("record_id", ""),
                    "fields_accessed": data.get("fields_accessed", []),
                },
            }
        )

    # Query access history for a resource
    resource = data.get("resource", "")
    record_id = data.get("record_id", "")
    search_term = resource or record_id

    if not search_term:
        return "Please provide a resource name or record_id to query access history."

    events = _search_jsonl(resource=search_term, limit=data.get("limit", 30))

    if not events:
        return f"No access events found for: {search_term}"

    lines = [f"🔒 Data Access Log: {search_term}\n"]
    for e in events:
        ts = str(e.get("timestamp", ""))[:19]
        usr = e.get("user_id", "")
        act = e.get("action", "")
        lines.append(f"  {ts} {usr}: {act}")
        meta = e.get("metadata", {})
        if meta.get("fields_accessed"):
            lines.append(f"    Fields: {', '.join(meta['fields_accessed'])}")

    return "\n".join(lines)


def export_audit_trail(data):
    """Export tamper-evident audit trail as CSV, JSON, or SIEM format."""
    fmt = data.get("format", "json").lower()
    start = data.get("start_date")
    end = data.get("end_date")
    limit = min(data.get("limit", 10000), 50000)

    if CORE_AUDIT_AVAILABLE:
        try:
            store = get_audit_store()
            start_dt = datetime.fromisoformat(start) if start else None
            end_dt = datetime.fromisoformat(end) if end else None
            _ensure_dir()
            ts = datetime.now().strftime("%Y%m%d_%H%M%S")

            if fmt == "csv":
                filepath = AUDIT_DIR / f"audit_export_{ts}.csv"
                store.export_csv(filepath, start=start_dt, end=end_dt)
                return f"✅ Audit trail exported: {filepath}"
            elif fmt == "siem":
                filepath = AUDIT_DIR / f"audit_siem_{ts}.json"
                store.export_siem(filepath, start=start_dt, end=end_dt)
                return f"✅ SIEM export: {filepath}"
            else:
                filepath = AUDIT_DIR / f"audit_export_{ts}.json"
                store.export_json(filepath, start=start_dt, end=end_dt)
                return f"✅ Audit trail exported: {filepath}"
        except Exception as e:
            logger.warning(f"Core audit export failed: {e}")

    # Fallback: export from JSONL
    events = _search_jsonl(start, end, limit=limit)
    if not events:
        return "No audit events to export."

    _ensure_dir()
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")

    if fmt == "csv":
        import csv

        filepath = AUDIT_DIR / f"audit_export_{ts}.csv"
        headers = [
            "event_id",
            "timestamp",
            "user_id",
            "action",
            "resource",
            "detail",
            "severity",
            "chain_hash",
        ]
        with open(filepath, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=headers, extrasaction="ignore")
            writer.writeheader()
            writer.writerows(events)
    else:
        filepath = AUDIT_DIR / f"audit_export_{ts}.json"
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(
                {
                    "events": events,
                    "exported_at": datetime.now().isoformat(),
                    "count": len(events),
                    "chain_verified": _verify_chain(),
                },
                f,
                indent=2,
                default=str,
            )

    # Log the export itself
    log_event(
        {
            "action": "audit_export",
            "user_id": "system",
            "resource": str(filepath),
            "detail": f"Exported {len(events)} events as {fmt}",
            "severity": "info",
        }
    )

    return f"✅ Audit trail exported: {filepath} ({len(events)} events, chain {'verified' if _verify_chain() else 'unverified'})"


# === Tool Definitions ===

TOOLS = [
    {
        "name": "log_event",
        "description": (
            "Log an audit event to the hash-chained audit trail for HIPAA compliance. "
            "Supports tool execution, data access, configuration change, and authentication events. "
            "Events are stored with tamper-evident HMAC-SHA-256 hash chaining. "
            "Use this to manually record significant actions not captured automatically."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "description": "Event action (tool_executed, data_view, config_change, login, etc.)",
                },
                "user_id": {
                    "type": "string",
                    "default": "system",
                    "description": "User who performed the action",
                },
                "resource": {
                    "type": "string",
                    "description": "Resource accessed (file path, record ID, skill name)",
                },
                "detail": {"type": "string", "description": "Additional detail"},
                "severity": {
                    "type": "string",
                    "enum": ["info", "warning", "critical"],
                    "default": "info",
                    "description": "Event severity level",
                },
                "metadata": {"type": "object", "description": "Additional structured metadata"},
            },
            "required": ["action"],
        },
        "input_examples": [
            {
                "action": "data_access",
                "resource": "patient_record_123",
                "user_id": "dr.smith",
                "detail": "Accessed patient chart for appointment",
                "severity": "info",
            },
        ],
        "handler": log_event,
        "category": "audit",
    },
    {
        "name": "search_audit_log",
        "description": (
            "Search audit events by date range, user, action, resource, or severity level. "
            "Queries both the core SQLite audit store and the fallback JSONL log. "
            "Use this to investigate specific activities, review user actions, or check for security incidents. "
            "Returns events with timestamps, user IDs, actions, and details."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "start_date": {"type": "string", "description": "Start date (ISO format)"},
                "end_date": {"type": "string", "description": "End date (ISO format)"},
                "user_id": {
                    "type": "string",
                    "description": "Filter by user who performed the action",
                },
                "action": {
                    "type": "string",
                    "description": "Filter by action type (e.g. tool_executed, data_view, config_change)",
                },
                "resource": {
                    "type": "string",
                    "description": "Filter by resource name or path (substring match)",
                },
                "severity": {
                    "type": "string",
                    "description": "Filter by severity level (info, warning, critical)",
                },
                "limit": {
                    "type": "integer",
                    "default": 50,
                    "description": "Maximum number of results to return",
                },
            },
        },
        "handler": search_audit_log,
        "category": "audit",
    },
    {
        "name": "compliance_report",
        "description": (
            "Generate a HIPAA compliance audit report for a date range with event statistics, PHI access tracking, and controls checklist. "
            "Aggregates events by action, user, and severity, then checks against HIPAA 164.312 requirements. "
            "The report includes hash-chain integrity verification. "
            "Saves the report to the audit directory and returns the full text."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "start_date": {
                    "type": "string",
                    "description": "Period start (default: 30 days ago)",
                },
                "end_date": {"type": "string", "description": "Period end (default: today)"},
            },
        },
        "handler": compliance_report,
        "category": "audit",
    },
    {
        "name": "data_access_log",
        "description": (
            "Log or query data access events for specific records, tracking who viewed what, when, and which fields. "
            "Set action to 'log' to record a new access event, or 'query' to search access history for a resource. "
            "Logs include record type, record ID, and fields accessed for granular HIPAA audit trails. "
            "Use this for patient record access tracking and compliance reporting."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "action": {"type": "string", "enum": ["log", "query"], "default": "query"},
                "user_id": {"type": "string", "description": "User who accessed the data"},
                "resource": {
                    "type": "string",
                    "description": "Resource name or path to log or query",
                },
                "record_type": {
                    "type": "string",
                    "description": "Type of record accessed (e.g. patient, contact, donor)",
                },
                "record_id": {
                    "type": "string",
                    "description": "Unique identifier of the accessed record",
                },
                "fields_accessed": {
                    "type": "array",
                    "items": {"type": "string", "description": "Field name that was accessed"},
                    "description": "List of specific fields that were viewed or modified",
                },
                "detail": {
                    "type": "string",
                    "description": "Additional context about the data access",
                },
                "limit": {
                    "type": "integer",
                    "default": 30,
                    "description": "Maximum number of results to return",
                },
            },
        },
        "handler": data_access_log,
        "category": "audit",
    },
    {
        "name": "export_audit_trail",
        "description": (
            "Export the tamper-evident hash-chained audit trail as CSV, JSON, or SIEM format. "
            "Includes chain integrity verification in the export metadata. "
            "Use this for external compliance reviews, SIEM ingestion, or archival. "
            "The export itself is logged as an audit event."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "format": {
                    "type": "string",
                    "enum": ["csv", "json", "siem"],
                    "default": "json",
                    "description": "Export format",
                },
                "start_date": {
                    "type": "string",
                    "description": "Start date filter in ISO format (e.g. 2026-01-01)",
                },
                "end_date": {
                    "type": "string",
                    "description": "End date filter in ISO format (e.g. 2026-12-31)",
                },
                "limit": {
                    "type": "integer",
                    "default": 10000,
                    "description": "Maximum number of events to export",
                },
            },
        },
        "handler": export_audit_trail,
        "category": "audit",
    },
]
