"""Privacy-minimized persistence and projection for SuperAgent trace SSE events."""

from __future__ import annotations

import json
import os
import re
import threading
from datetime import datetime, timezone
from hashlib import sha256
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Protocol, Union

from agent_integration.client import OpenAgentEvent
from arr_ingestion.contracts import OPAQUE_ID_RE


TRACE_RECORD_VERSION = "arr-agent-trace-event-1"
MAX_TRACE_EVENTS = 10_000
MAX_TRACE_FILE_BYTES = 8 * 1024 * 1024
MAX_TRACE_SUMMARY_CHARS = 1_200

_LEVELS = frozenset({"info", "success", "warning", "error"})
_CODE_RE = re.compile(r"^[A-Z][A-Z0-9_]{0,95}$")
_SAFE_DETAIL_KEYS = frozenset(
    {
        "attempt_no",
        "remote_run_id",
        "remote_status",
        "trace_event",
        "invalid_records",
        "capture_error",
    }
)
_SENSITIVE_ASSIGNMENT_RE = re.compile(
    r"(?i)\b(submission[_ -]?grant|access[_ -]?key(?:[_ -]?(?:id|secret))?|"
    r"api[_ -]?key|secret|token|password|authorization|cookie)\b"
    r"\s*[:=]\s*(?:[\"']?)[^\s,;}\]]+"
)
_BEARER_RE = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]+")
_URL_RE = re.compile(r"(?i)\b(?:https?|oss)://[^\s<>{}\[\]\"']+")
_PATH_RE = re.compile(
    r"(?<![A-Za-z0-9._-])/(?:Users|home|private|tmp|var|opt|etc)(?:/[^\s,;}\]]+)+"
)
_EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b")
_LONG_NUMBER_RE = re.compile(r"(?<!\d)(?:\+?\d[\d ()-]{7,}\d)(?!\d)")
_XML_FRAGMENT_RE = re.compile(r"<[^>\r\n]{1,240}>")


class AgentTraceStoreError(RuntimeError):
    """Raised when the local trace store cannot safely persist a record."""


class AgentTraceSink(Protocol):
    def append(self, job_id: str, record: Mapping[str, Any]) -> None:
        ...


class AgentTraceReader(Protocol):
    def read(self, job_id: str) -> List[Dict[str, Any]]:
        ...


def _utc_now() -> datetime:
    return datetime.now(timezone.utc)


def _timestamp(value: Any, *, fallback: Optional[datetime] = None) -> str:
    if isinstance(value, datetime):
        parsed = value
    elif isinstance(value, str):
        try:
            parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
        except ValueError:
            parsed = fallback or _utc_now()
    else:
        parsed = fallback or _utc_now()
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc).isoformat()


def sanitize_trace_summary(value: Any) -> Optional[str]:
    """Keep the server's operational summary while removing common secret/PII forms."""

    if not isinstance(value, str):
        return None
    text = " ".join(value.replace("\x00", " ").split())
    if not text:
        return None
    text = _BEARER_RE.sub("Bearer [REDACTED]", text)
    text = _SENSITIVE_ASSIGNMENT_RE.sub(lambda match: f"{match.group(1)}=[REDACTED]", text)
    text = _URL_RE.sub("[URL]", text)
    text = _PATH_RE.sub("[PATH]", text)
    text = _EMAIL_RE.sub("[EMAIL]", text)
    text = _LONG_NUMBER_RE.sub("[NUMBER]", text)
    text = _XML_FRAGMENT_RE.sub("[XML]", text)
    return text[:MAX_TRACE_SUMMARY_CHARS]


def trace_run_id(event: OpenAgentEvent) -> Optional[str]:
    """Return the run id only from the documented trace run-start event."""

    payload = event.data
    if event.event != "trace" or not isinstance(payload, Mapping):
        return None
    if payload.get("event") != "run.started":
        return None
    run_id = payload.get("run_id")
    if not isinstance(run_id, str) or OPAQUE_ID_RE.fullmatch(run_id) is None:
        return None
    return run_id


def project_trace_event(
    event: OpenAgentEvent,
    *,
    attempt_no: int,
    sequence: int,
    now: Optional[datetime] = None,
) -> Optional[Dict[str, Any]]:
    """Project one raw SSE event into the strict local diagnostic contract."""

    fallback = now or _utc_now()
    if event.event == "end":
        return _record(
            sequence=sequence,
            timestamp=_timestamp(None, fallback=fallback),
            level="info",
            code="AGENT_TRACE_STREAM_ENDED",
            title="SuperAgent trace 流已结束",
            message="SuperAgent 已关闭本次运行的 trace 连接。",
            details={"attempt_no": attempt_no, "trace_event": "end"},
        )
    if event.event != "trace" or not isinstance(event.data, Mapping):
        return None

    payload = event.data
    name = payload.get("event")
    if not isinstance(name, str) or not name:
        return None
    if name in {"message.delta", "message.final"}:
        # Assistant text can contain result payloads, grants, object routes or PII.
        return None

    nested = payload.get("data")
    summary = sanitize_trace_summary(
        nested.get("summary") if isinstance(nested, Mapping) else None
    )
    run_id = payload.get("run_id")
    if not isinstance(run_id, str) or OPAQUE_ID_RE.fullmatch(run_id) is None:
        run_id = None
    status = payload.get("status")
    if not isinstance(status, str) or not status or len(status) > 32:
        status = None
    details: Dict[str, Any] = {
        "attempt_no": attempt_no,
        "trace_event": name,
        "remote_run_id": run_id,
        "remote_status": status,
    }
    timestamp = _timestamp(payload.get("ts"), fallback=fallback)

    if name == "run.started":
        return _record(
            sequence=sequence,
            timestamp=timestamp,
            level="info",
            code="AGENT_RUN_STARTED",
            title="SuperAgent 运行已启动",
            message="已建立包含任务 trace 的 SuperAgent 运行流。",
            details=details,
        )
    if name == "task.updated":
        if summary is None:
            return None
        return _record(
            sequence=sequence,
            timestamp=timestamp,
            level="info",
            code="AGENT_TASK_UPDATED",
            title="Agent 任务状态更新",
            message=summary,
            details=details,
        )
    if name == "step.updated":
        if summary is None:
            return None
        return _record(
            sequence=sequence,
            timestamp=timestamp,
            level="info",
            code="AGENT_STEP_UPDATED",
            title="Agent 执行步骤更新",
            message=summary,
            details=details,
        )
    if name == "run.completed":
        succeeded = status in {"success", "completed", "succeeded"}
        return _record(
            sequence=sequence,
            timestamp=timestamp,
            level="success" if succeeded else "warning",
            code="AGENT_RUN_COMPLETED",
            title="SuperAgent 运行已结束",
            message=summary or "SuperAgent 已报告运行终态；业务成功仍以 ARR 数据库提交为准。",
            details=details,
        )
    if name in {"run.failed", "error"}:
        return _record(
            sequence=sequence,
            timestamp=timestamp,
            level="error",
            code="AGENT_RUN_FAILED",
            title="SuperAgent 运行失败",
            message=summary or "SuperAgent trace 报告运行失败。",
            details=details,
        )
    if name in {"run.cancelled", "run.canceled"}:
        return _record(
            sequence=sequence,
            timestamp=timestamp,
            level="warning",
            code="AGENT_RUN_CANCELLED",
            title="SuperAgent 运行已取消",
            message=summary or "SuperAgent trace 报告运行已取消。",
            details=details,
        )
    if summary is None:
        return None
    return _record(
        sequence=sequence,
        timestamp=timestamp,
        level="info",
        code="AGENT_TRACE_EVENT",
        title="Agent trace 事件",
        message=summary,
        details=details,
    )


def capture_failure_record(
    *,
    attempt_no: int,
    sequence: int,
    error: BaseException,
    now: Optional[datetime] = None,
) -> Dict[str, Any]:
    """Create a safe local event without serializing the exception message."""

    error_name = type(error).__name__
    if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,95}", error_name):
        error_name = "Exception"
    return _record(
        sequence=sequence,
        timestamp=_timestamp(now or _utc_now()),
        level="warning",
        code="AGENT_TRACE_CAPTURE_FAILED",
        title="SuperAgent trace 采集中断",
        message="远端运行可能仍在继续；当前日志连接未能完整消费。",
        details={
            "attempt_no": attempt_no,
            "trace_event": "capture.failed",
            "capture_error": error_name,
        },
    )


def _record(
    *,
    sequence: int,
    timestamp: str,
    level: str,
    code: str,
    title: str,
    message: str,
    details: Optional[Mapping[str, Any]] = None,
) -> Dict[str, Any]:
    value: Dict[str, Any] = {
        "version": TRACE_RECORD_VERSION,
        "id": f"agent.{sequence:06d}.{code.lower()}",
        "timestamp": timestamp,
        "stage": "agent",
        "level": level,
        "code": code,
        "title": title,
        "message": message,
    }
    safe_details = _validated_details(details or {})
    if safe_details:
        value["details"] = safe_details
    return value


def _validated_details(value: Mapping[str, Any]) -> Dict[str, Any]:
    output: Dict[str, Any] = {}
    for key, item in value.items():
        if key not in _SAFE_DETAIL_KEYS or item is None or item == "":
            continue
        if key in {"attempt_no", "invalid_records"}:
            if isinstance(item, int) and not isinstance(item, bool) and item >= 0:
                output[key] = item
        elif key == "remote_run_id":
            if isinstance(item, str) and OPAQUE_ID_RE.fullmatch(item) is not None:
                output[key] = item
        elif key == "capture_error":
            if isinstance(item, str) and re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,95}", item):
                output[key] = item
        elif isinstance(item, str) and re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", item):
            output[key] = item
    return output


def validate_trace_record(record: Mapping[str, Any]) -> Dict[str, Any]:
    required = {
        "version",
        "id",
        "timestamp",
        "stage",
        "level",
        "code",
        "title",
        "message",
    }
    optional = {"details"}
    if not isinstance(record, Mapping) or not required <= set(record) <= required | optional:
        raise AgentTraceStoreError("trace record field set is invalid")
    if record.get("version") != TRACE_RECORD_VERSION or record.get("stage") != "agent":
        raise AgentTraceStoreError("trace record version or stage is invalid")
    if record.get("level") not in _LEVELS:
        raise AgentTraceStoreError("trace record level is invalid")
    code = record.get("code")
    if not isinstance(code, str) or _CODE_RE.fullmatch(code) is None:
        raise AgentTraceStoreError("trace record code is invalid")
    timestamp = _bounded_text(record.get("timestamp"), 64, "timestamp")
    try:
        parsed_timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
    except ValueError as error:
        raise AgentTraceStoreError("trace record timestamp is invalid") from error
    if parsed_timestamp.tzinfo is None:
        raise AgentTraceStoreError("trace record timestamp is invalid")
    title = sanitize_trace_summary(_bounded_text(record.get("title"), 160, "title"))
    message = sanitize_trace_summary(
        _bounded_text(record.get("message"), MAX_TRACE_SUMMARY_CHARS, "message")
    )
    if title is None or message is None:
        raise AgentTraceStoreError("trace record text is invalid")
    output: Dict[str, Any] = {
        "version": TRACE_RECORD_VERSION,
        "id": _bounded_text(record.get("id"), 180, "id"),
        "timestamp": parsed_timestamp.astimezone(timezone.utc).isoformat(),
        "stage": "agent",
        "level": str(record["level"]),
        "code": code,
        "title": title,
        "message": message,
    }
    details = record.get("details")
    if details is not None:
        if not isinstance(details, Mapping):
            raise AgentTraceStoreError("trace record details are invalid")
        safe_details = _validated_details(details)
        if safe_details:
            output["details"] = safe_details
    return output


def _bounded_text(value: Any, maximum: int, label: str) -> str:
    if not isinstance(value, str) or not value or len(value) > maximum:
        raise AgentTraceStoreError(f"trace record {label} is invalid")
    return value


class JsonlAgentTraceStore:
    """Append-only local trace store; raw SSE payloads are never written."""

    def __init__(self, root: Union[str, Path]) -> None:
        self.root = Path(root)
        self._lock = threading.RLock()

    def append(self, job_id: str, record: Mapping[str, Any]) -> None:
        path = self._path(job_id)
        value = validate_trace_record(record)
        line = (
            json.dumps(
                value,
                ensure_ascii=False,
                allow_nan=False,
                sort_keys=True,
                separators=(",", ":"),
            )
            + "\n"
        ).encode("utf-8")
        try:
            with self._lock:
                if self.root.is_symlink() or (
                    self.root.exists() and not self.root.is_dir()
                ):
                    raise AgentTraceStoreError("agent trace root is invalid")
                self.root.mkdir(mode=0o700, parents=True, exist_ok=True)
                flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND
                if hasattr(os, "O_NOFOLLOW"):
                    flags |= os.O_NOFOLLOW
                descriptor = os.open(path, flags, 0o600)
                try:
                    remaining = memoryview(line)
                    while remaining:
                        written = os.write(descriptor, remaining)
                        if written <= 0:
                            raise OSError("trace append made no progress")
                        remaining = remaining[written:]
                finally:
                    os.close(descriptor)
        except (OSError, ValueError) as error:
            raise AgentTraceStoreError("cannot append agent trace record") from error

    def read(self, job_id: str) -> List[Dict[str, Any]]:
        path = self._path(job_id)
        try:
            with self._lock:
                if self.root.is_symlink():
                    raise AgentTraceStoreError("agent trace root is invalid")
                if not path.exists():
                    return []
                if path.is_symlink() or not path.is_file():
                    raise AgentTraceStoreError("agent trace path is invalid")
                metadata = path.stat()
                if metadata.st_size > MAX_TRACE_FILE_BYTES:
                    return [self._read_warning(metadata.st_mtime, 1, "file.too_large")]
                modified_at = metadata.st_mtime
                lines = path.read_text(encoding="utf-8").splitlines()
        except AgentTraceStoreError:
            raise
        except (OSError, UnicodeError) as error:
            raise AgentTraceStoreError("cannot read agent trace records") from error

        records: List[Dict[str, Any]] = []
        invalid = 0
        for line in lines[:MAX_TRACE_EVENTS]:
            try:
                payload = json.loads(line)
                records.append(validate_trace_record(payload))
            except (ValueError, TypeError, AgentTraceStoreError):
                invalid += 1
        if len(lines) > MAX_TRACE_EVENTS:
            invalid += len(lines) - MAX_TRACE_EVENTS
        if invalid:
            records.append(self._read_warning(modified_at, invalid, "records.invalid"))
        return records

    def _path(self, job_id: str) -> Path:
        if not isinstance(job_id, str) or OPAQUE_ID_RE.fullmatch(job_id) is None:
            raise AgentTraceStoreError("job id is invalid")
        return self.root / f"{sha256(job_id.encode('utf-8')).hexdigest()}.jsonl"

    @staticmethod
    def _read_warning(modified_at: float, count: int, reason: str) -> Dict[str, Any]:
        return _record(
            sequence=MAX_TRACE_EVENTS + 1,
            timestamp=datetime.fromtimestamp(modified_at, tz=timezone.utc).isoformat(),
            level="warning",
            code="AGENT_TRACE_RECORD_INVALID",
            title="本地 Agent trace 记录不完整",
            message="部分本地 trace 记录无法安全读取，已从任务日志中忽略。",
            details={
                "trace_event": reason,
                "invalid_records": count,
            },
        )