#!/usr/bin/env python3
"""Check the typed response-reference and raw proxy forwarding boundaries."""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PROXY = ROOT / "src" / "appview" / "proxy.rs"


class CheckError(ValueError):
    pass


def block_from(source: str, marker: str, next_marker: str) -> str:
    start = source.find(marker)
    if start < 0:
        raise CheckError(f"missing source marker: {marker}")
    end = source.find(next_marker, start + len(marker))
    if end < 0:
        raise CheckError(f"missing source marker: {next_marker}")
    return source[start:end]


def parser_region(source: str) -> str:
    return block_from(source, "struct RecordRef", "/// Read-your-writes override")


def repo_get_record_region(source: str) -> str:
    return block_from(source, "pub(super) async fn repo_get_record", "pub(super) async fn repo_list_records")


def forwarding_region(source: str) -> str:
    return block_from(source, "pub(super) fn forward_xrpc", "async fn actor_did")


def raw_response_region(source: str) -> str:
    return block_from(source, "fn raw_response", "async fn actor_did")


def proxy_response_parser_has_borrowed_typed_envelopes(source: str) -> None:
    region = parser_region(source)
    required = (
        "struct RecordRef<'a>", "uri: AtUri<&'a str>", "cid: Cid<&'a str>",
        "struct RecordCid<'a>", "fn parse_record_ref<'a>", "fn parse_record_cid<'a>",
        "expected_did: &Did", "expected_collection: &Nsid", "Rkey<&'a str>",
        "out.uri.authority()", "out.uri.collection()?", "out.uri.rkey()?",
        "out.cid.is_valid()",
    )
    for needle in required:
        if needle not in region:
            raise CheckError(f"borrowed response parser is missing {needle}")
    for needle in ("serde_json::from_slice::<Value>", ".get(\"cid\")", "Cid::new_owned", ".rsplit('/')"):
        if needle in region:
            raise CheckError(f"response parser retains forbidden extraction: {needle}")


def proxy_get_record_uses_borrowed_cid_parser(source: str) -> None:
    region = repo_get_record_region(source)
    if "parse_record_cid(resp.buffer())" not in region:
        raise CheckError("repo_get_record does not use parse_record_cid")
    for needle in ("serde_json::from_slice::<Value>", ".get(\"cid\")"):
        if needle in region:
            raise CheckError(f"repo_get_record retains manual CID extraction: {needle}")


def proxy_forwarding_contract_is_unchanged(source: str) -> None:
    region = forwarding_region(source)
    for needle in ("Ok(resp) => raw_response(resp.status(), resp.buffer().to_vec())", "ClientErrorKind::Http { status }", "raw_response(*status, body)"):
        if needle not in region:
            raise CheckError(f"forward_xrpc no longer preserves {needle}")
    if 'header::CONTENT_TYPE, "application/json"' not in raw_response_region(source):
        raise CheckError("raw_response no longer synthesizes application/json")


def run_checks(source: str) -> None:
    proxy_response_parser_has_borrowed_typed_envelopes(source)
    proxy_get_record_uses_borrowed_cid_parser(source)
    proxy_forwarding_contract_is_unchanged(source)


def positive_fixture() -> str:
    return '''
    pub(super) async fn repo_get_record() {
        let pds_cid = parse_record_cid(resp.buffer()).map(|cid| cid.as_str().to_owned());
    }
    pub(super) async fn repo_list_records() {}
    pub(super) fn forward_xrpc<Resp: XrpcResp>(result: Result<XrpcResponse<Resp>, ClientError>) -> Response {
        match result {
            Ok(resp) => raw_response(resp.status(), resp.buffer().to_vec()),
            Err(err) => match err.kind() {
                ClientErrorKind::Http { status } => raw_response(*status, body),
            },
        }
    }
    fn raw_response(status: StatusCode, body: Vec<u8>) -> Response {
        Response::builder().header(header::CONTENT_TYPE, "application/json").body(body).unwrap()
    }
    async fn actor_did() {}
    #[derive(serde::Deserialize)]
    struct RecordRef<'a> { uri: AtUri<&'a str>, cid: Cid<&'a str> }
    #[derive(serde::Deserialize)]
    struct RecordCid<'a> { cid: Cid<&'a str> }
    fn parse_record_ref<'a>(output: &'a [u8], expected_did: &Did, expected_collection: &Nsid) -> Option<(Rkey<&'a str>, Cid<&'a str>)> {
        let out: RecordRef<'a> = serde_json::from_slice(output).ok()?;
        if !out.cid.is_valid() { return None; }
        match out.uri.authority() { AtIdentifier::Did(did) if did.as_ref() == expected_did.as_ref() => {}, _ => return None }
        if out.uri.collection()?.as_ref() != expected_collection.as_str() { return None; }
        let rkey = out.uri.rkey()?;
        if rkey.as_str().is_empty() { return None; }
        Some((rkey, out.cid))
    }
    fn parse_record_cid<'a>(output: &'a [u8]) -> Option<Cid<&'a str>> {
        let out: RecordCid<'a> = serde_json::from_slice(output).ok()?;
        out.cid.is_valid().then_some(out.cid)
    }
    /// Read-your-writes override
    '''


def self_test() -> None:
    run_checks(positive_fixture())
    cases = {
        "parser": positive_fixture().replace("uri: AtUri<&'a str>", "uri: Value"),
        "get_record": positive_fixture().replace("parse_record_cid(resp.buffer()).map(|cid| cid.as_str().to_owned())", "serde_json::from_slice::<Value>(resp.buffer()).ok()"),
        "forwarding": positive_fixture().replace("Ok(resp) => raw_response(resp.status(), resp.buffer().to_vec()),", "Ok(resp) => Response::new(Body::from(resp.buffer().to_vec())),"),
    }
    for name, fixture in cases.items():
        try:
            run_checks(fixture)
        except CheckError:
            continue
        raise CheckError(f"negative self-test fixture unexpectedly passed: {name}")
    print("check-proxy-response self-test passed")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    try:
        if args.self_test:
            self_test()
        else:
            run_checks(PROXY.read_text())
            print("proxy response checks passed")
    except (CheckError, OSError) as exc:
        print(f"check-proxy-response failed: {exc}", file=sys.stderr)
        return 1
    return 0


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