"""Trusted runtime adapter: publish Agent files, sign the result, call ARR."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import stat
import time
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Protocol, Sequence, Tuple
from urllib.parse import urlparse

from arr_ingestion.contracts import OPAQUE_ID_RE, PROCESSOR_VERSION_RE, SHA256_RE, XLSX_MIME
from arr_processing.config import ResultVerificationConfig
from arr_processing.contracts import ProcessingResult
from arr_processing.errors import ProcessingError
from arr_storage.aliyun_oss_v2 import AliyunOssConfig, AliyunOssV2Client
from arr_storage.exchange import OutputExchangeConfig, OutputExchangePublisher


FROZEN_FIELDS = {
    "contract_version",
    "job_id",
    "source_file_id",
    "status",
    "business_date",
    "processor_result",
    "files",
}
PROCESSOR_RESULT_FIELDS = {
    "version",
    "status",
    "business_date",
    "message",
    "metrics",
    "outputs",
    "errors",
}
FILE_FIELDS = {"filename", "local_path"}
FILE_ROLES = {"daily_report", "result_json", "structured_result", "exception_report"}
ROLE_MIME = {
    "daily_report": XLSX_MIME,
    "result_json": "application/json",
    "structured_result_json": "application/json",
    "exception_report": XLSX_MIME,
}
RESULT_ROLE = {
    "daily_report": "daily_report",
    "result_json": "result_json",
    "structured_result": "structured_result_json",
    "exception_report": "exception_report",
}
SAFE_FILENAME_RE = re.compile(r"^[^/\\\x00-\x1f\x7f]{1,255}$")


class RuntimeWritebackError(RuntimeError):
    def __init__(self, code: str, safe_message: str, *, retryable: bool = False) -> None:
        super().__init__(safe_message)
        self.code = code
        self.safe_message = safe_message
        self.retryable = retryable


class OutputPublisher(Protocol):
    def publish(
        self,
        *,
        file_handle: str,
        source: Path,
        mime_type: str,
        metadata: Mapping[str, str],
    ) -> Any:
        ...


class HttpClient(Protocol):
    def post(self, url: str, **kwargs: Any) -> Any:
        ...


def _strict_json(raw: bytes) -> Mapping[str, Any]:
    def pairs(values: Sequence[Tuple[str, Any]]) -> Dict[str, Any]:
        output: Dict[str, Any] = {}
        for key, value in values:
            if key in output:
                raise ValueError("duplicate JSON key")
            output[key] = value
        return output

    try:
        value = json.loads(raw.decode("utf-8"), object_pairs_hook=pairs)
    except (UnicodeError, ValueError, json.JSONDecodeError):
        raise RuntimeWritebackError(
            "RUNTIME_RESULT_INVALID", "Agent result is not strict JSON"
        ) from None
    if not isinstance(value, Mapping):
        raise RuntimeWritebackError(
            "RUNTIME_RESULT_INVALID", "Agent result must be a JSON object"
        )
    return value


def _opaque(value: Any, label: str) -> str:
    if not isinstance(value, str) or not OPAQUE_ID_RE.fullmatch(value):
        raise RuntimeWritebackError("RUNTIME_RESULT_INVALID", f"{label} is invalid")
    return value


@dataclass(frozen=True)
class FrozenFile:
    filename: str
    path: Path


@dataclass(frozen=True)
class FrozenAgentResult:
    job_id: str
    source_file_id: str
    status: str
    business_date: Optional[date]
    files: Mapping[str, Optional[FrozenFile]]

    @classmethod
    def parse(cls, raw: bytes, *, output_root: Path) -> "FrozenAgentResult":
        values = _strict_json(raw)
        if set(values) != FROZEN_FIELDS or values.get("contract_version") != "arr-opera-daily-agent-1":
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "Agent result contract is invalid"
            )
        job_id = _opaque(values.get("job_id"), "job identifier")
        source_file_id = _opaque(values.get("source_file_id"), "source file identifier")
        status = values.get("status")
        if status not in {"success", "failed"}:
            raise RuntimeWritebackError("RUNTIME_RESULT_INVALID", "Agent status is invalid")
        raw_date = values.get("business_date")
        try:
            if raw_date == "":
                business_date = None
            elif isinstance(raw_date, str):
                business_date = date.fromisoformat(raw_date)
            else:
                raise ValueError
        except ValueError:
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "Agent business date is invalid"
            ) from None
        processor_business_date = business_date.isoformat() if business_date else None
        processor = values.get("processor_result")
        if not isinstance(processor, Mapping) or set(processor) != PROCESSOR_RESULT_FIELDS:
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "processor result contract is invalid"
            )
        if (
            processor.get("version") != "3.0"
            or processor.get("status") != status
            or processor.get("business_date") != processor_business_date
            or not isinstance(processor.get("message"), str)
            or not isinstance(processor.get("metrics"), Mapping)
            or not isinstance(processor.get("outputs"), Mapping)
            or not isinstance(processor.get("errors"), list)
        ):
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "processor result identity is invalid"
            )
        raw_files = values.get("files")
        if not isinstance(raw_files, Mapping) or set(raw_files) != FILE_ROLES:
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "Agent output file contract is invalid"
            )
        root = output_root.resolve(strict=True)
        if not root.is_dir():
            raise RuntimeWritebackError(
                "RUNTIME_OUTPUT_ROOT_INVALID", "Agent output root is unavailable"
            )
        files: Dict[str, Optional[FrozenFile]] = {}
        for role in sorted(FILE_ROLES):
            item = raw_files.get(role)
            if item is None:
                files[role] = None
                continue
            if not isinstance(item, Mapping) or set(item) != FILE_FIELDS:
                raise RuntimeWritebackError(
                    "RUNTIME_RESULT_INVALID", f"{role} file contract is invalid"
                )
            filename = item.get("filename")
            local_path = item.get("local_path")
            if (
                not isinstance(filename, str)
                or SAFE_FILENAME_RE.fullmatch(filename) is None
                or not isinstance(local_path, str)
                or not Path(local_path).is_absolute()
            ):
                raise RuntimeWritebackError(
                    "RUNTIME_RESULT_INVALID", f"{role} file identity is invalid"
                )
            path = Path(local_path)
            try:
                details = path.lstat()
                resolved = path.resolve(strict=True)
                resolved.relative_to(root)
            except (OSError, ValueError):
                raise RuntimeWritebackError(
                    "RUNTIME_OUTPUT_INVALID", f"{role} output is outside the approved root"
                ) from None
            if path.is_symlink() or not stat.S_ISREG(details.st_mode) or resolved.name != filename:
                raise RuntimeWritebackError(
                    "RUNTIME_OUTPUT_INVALID", f"{role} output is not an approved file"
                )
            expected_extension = ".json" if role in {"result_json", "structured_result"} else ".xlsx"
            if not filename.lower().endswith(expected_extension):
                raise RuntimeWritebackError(
                    "RUNTIME_OUTPUT_INVALID", f"{role} output type is invalid"
                )
            files[role] = FrozenFile(filename, resolved)

        if files["result_json"] is None or files["structured_result"] is None:
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "required Agent result files are missing"
            )
        if status == "success":
            valid_shape = (
                business_date is not None
                and files["daily_report"] is not None
                and files["exception_report"] is None
            )
        else:
            valid_shape = files["daily_report"] is None and files["exception_report"] is not None
        if not valid_shape:
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "Agent result file shape is invalid"
            )
        return cls(job_id, source_file_id, str(status), business_date, files)


@dataclass
class AgentRuntimeWritebackAdapter:
    publisher: OutputPublisher
    signer: Any
    http_client: HttpClient
    callback_url: str
    processor_version: str
    rule_set_sha256: str
    max_http_attempts: int = 3
    retry_delay_seconds: float = 0.5

    def __post_init__(self) -> None:
        parsed = urlparse(self.callback_url)
        localhost = parsed.hostname in {"127.0.0.1", "localhost", "::1"}
        if (
            parsed.scheme not in ({"http", "https"} if localhost else {"https"})
            or not parsed.netloc
            or parsed.query
            or parsed.fragment
        ):
            raise ValueError("ARR Agent callback URL must be an HTTPS URL")
        if not PROCESSOR_VERSION_RE.fullmatch(self.processor_version):
            raise ValueError("processor version is invalid")
        if not SHA256_RE.fullmatch(self.rule_set_sha256):
            raise ValueError("rule set identity is invalid")
        if self.max_http_attempts < 1 or self.retry_delay_seconds < 0:
            raise ValueError("runtime writeback retry policy is invalid")

    def deliver(
        self,
        raw_frozen_result: bytes,
        *,
        output_root: Path,
        attempt_no: int,
        remote_run_id: str,
    ) -> Mapping[str, Any]:
        frozen = FrozenAgentResult.parse(raw_frozen_result, output_root=output_root)
        remote_run_id = _opaque(remote_run_id, "remote run identifier")
        if not isinstance(attempt_no, int) or isinstance(attempt_no, bool) or not 1 <= attempt_no <= 9999:
            raise RuntimeWritebackError(
                "RUNTIME_RESULT_INVALID", "attempt number is invalid"
            )
        artifacts: Dict[str, Any] = {
            "daily_report": None,
            "result_json": None,
            "structured_result_json": None,
            "exception_report": None,
        }
        for frozen_role in sorted(FILE_ROLES):
            item = frozen.files[frozen_role]
            if item is None:
                continue
            role = RESULT_ROLE[frozen_role]
            digest, byte_size = self._identity(item.path)
            file_handle = self._file_handle(
                frozen.job_id,
                attempt_no,
                remote_run_id,
                role,
                digest,
            )
            mime_type = ROLE_MIME[role]
            metadata = {
                "arr-exchange-schema": "1.0",
                "arr-file-handle": file_handle,
                "arr-job-id": frozen.job_id,
                "arr-attempt-no": str(attempt_no),
                "arr-role": role,
                "arr-sha256": digest,
                "arr-byte-size": str(byte_size),
                "arr-mime-type": mime_type,
            }
            try:
                self.publisher.publish(
                    file_handle=file_handle,
                    source=item.path,
                    mime_type=mime_type,
                    metadata=metadata,
                )
            except Exception as error:
                if isinstance(error, RuntimeWritebackError):
                    raise
                raise RuntimeWritebackError(
                    "RUNTIME_OUTPUT_PUBLISH_FAILED",
                    "Agent output could not be published",
                    retryable=True,
                ) from None
            artifacts[role] = {
                "file_handle": file_handle,
                "original_filename": item.filename,
                "sha256": digest,
                "byte_size": byte_size,
                "mime_type": mime_type,
            }
        result = ProcessingResult.from_dict(
            {
                "contract_version": "arr-opera-daily-result-1",
                "delivery_id": self._delivery_id(frozen.job_id, attempt_no, remote_run_id),
                "job_id": frozen.job_id,
                "attempt_no": attempt_no,
                "remote_run_id": remote_run_id,
                "status": frozen.status,
                "business_date": frozen.business_date.isoformat() if frozen.business_date else None,
                "processor_version": self.processor_version,
                "rule_set_sha256": self.rule_set_sha256,
                "result_schema_version": "3.0",
                "artifacts": artifacts,
            }
        )
        signed = self.signer.sign(result)
        return self._post_same_bytes(signed)

    def _post_same_bytes(self, signed: bytes) -> Mapping[str, Any]:
        for attempt in range(1, self.max_http_attempts + 1):
            try:
                response = self.http_client.post(
                    self.callback_url,
                    content=signed,
                    headers={
                        "Content-Type": "application/json",
                        "Accept": "application/json",
                        "User-Agent": "arr-agent-runtime-writeback/1.0",
                    },
                )
            except Exception:
                if attempt >= self.max_http_attempts:
                    raise RuntimeWritebackError(
                        "RUNTIME_CALLBACK_UNAVAILABLE",
                        "ARR callback did not respond",
                        retryable=True,
                    ) from None
                time.sleep(self.retry_delay_seconds * attempt)
                continue
            status = int(getattr(response, "status_code", 0) or 0)
            if status == 200:
                try:
                    payload = response.json()
                except Exception:
                    raise RuntimeWritebackError(
                        "RUNTIME_CALLBACK_PROTOCOL_INVALID",
                        "ARR callback returned an invalid response",
                    ) from None
                if (
                    not isinstance(payload, Mapping)
                    or payload.get("ok") is not True
                    or not isinstance(payload.get("data"), Mapping)
                ):
                    raise RuntimeWritebackError(
                        "RUNTIME_CALLBACK_PROTOCOL_INVALID",
                        "ARR callback returned an invalid response",
                    )
                return dict(payload["data"])
            if status in {408, 429} or status >= 500:
                if attempt < self.max_http_attempts:
                    time.sleep(self.retry_delay_seconds * attempt)
                    continue
                raise RuntimeWritebackError(
                    "RUNTIME_CALLBACK_UNAVAILABLE",
                    "ARR callback is temporarily unavailable",
                    retryable=True,
                )
            raise RuntimeWritebackError(
                "RUNTIME_CALLBACK_REJECTED",
                "ARR rejected the Agent result",
            )
        raise AssertionError("unreachable retry loop")

    @staticmethod
    def _identity(path: Path) -> Tuple[str, int]:
        digest = hashlib.sha256()
        size = 0
        with path.open("rb") as source:
            for chunk in iter(lambda: source.read(1024 * 1024), b""):
                size += len(chunk)
                digest.update(chunk)
        return digest.hexdigest(), size

    @staticmethod
    def _file_handle(
        job_id: str,
        attempt_no: int,
        remote_run_id: str,
        role: str,
        sha256: str,
    ) -> str:
        value = "\x1f".join((job_id, str(attempt_no), remote_run_id, role, sha256))
        return "arrout_" + hashlib.sha256(value.encode("utf-8")).hexdigest()[:48]

    @staticmethod
    def _delivery_id(job_id: str, attempt_no: int, remote_run_id: str) -> str:
        value = "\x1f".join((job_id, str(attempt_no), remote_run_id, "delivery"))
        return "arrdel_" + hashlib.sha256(value.encode("utf-8")).hexdigest()[:48]


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="arr-agent-writeback")
    parser.add_argument("--frozen-result", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--attempt-no", type=int, required=True)
    parser.add_argument("--remote-run-id", required=True)
    parser.add_argument("--processor-version", required=True)
    parser.add_argument("--rule-set-sha256", required=True)
    parser.add_argument(
        "--callback-url",
        default=os.environ.get("ARR_AGENT_CALLBACK_URL", ""),
    )
    return parser


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = _parser().parse_args(argv)
    if not args.callback_url:
        print(json.dumps({"ok": False, "error": {"code": "CALLBACK_CONFIG_MISSING"}}))
        return 2
    try:
        import httpx

        oss_client = AliyunOssV2Client(AliyunOssConfig.from_environment())
        oss_client.assert_immutable_writes_supported()
        publisher = OutputExchangePublisher(
            oss_client,
            OutputExchangeConfig.from_environment(),
        )
        signing = ResultVerificationConfig.from_environment().signer()
        with httpx.Client(timeout=httpx.Timeout(30.0, connect=10.0)) as http_client:
            receipt = AgentRuntimeWritebackAdapter(
                publisher=publisher,
                signer=signing,
                http_client=http_client,
                callback_url=args.callback_url,
                processor_version=args.processor_version,
                rule_set_sha256=args.rule_set_sha256,
            ).deliver(
                args.frozen_result.read_bytes(),
                output_root=args.output_root,
                attempt_no=args.attempt_no,
                remote_run_id=args.remote_run_id,
            )
        oss_client.close()
    except (OSError, ValueError, ProcessingError, RuntimeWritebackError) as error:
        code = getattr(error, "code", "RUNTIME_WRITEBACK_FAILED")
        print(json.dumps({"ok": False, "error": {"code": code}}))
        return 3
    print(json.dumps({"ok": True, "data": receipt}, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
