from __future__ import annotations

from pathlib import Path
import hashlib
import json
import sqlite3
import sys

import pytest

ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/"active-runtime/app"))

from aios_agent_task_initial_execution_projection import (
    InitialExecutionProjectionError,
    InitialExecutionProjectionRequest,
    InitialExecutionProjectionStore,
    SourceCaptureProjectionRecoveryRequest,
)
from aios_online_agent_execution import (
    GovernedExecutionBinding,
    project_execution_envelope,
)

AUTHORIZATION="owner-initial-execution-projection-recovery-20260820"
TASK_ID="task-initial-projection-a"
JOB_ID="job-initial-projection-a"
BASE="a"*40
BRANCH="agent/task-initial-projection-a"
IDEMPOTENCY="task-initial-projection-binding-v1"
CAPABILITY="agent.runtime.openclaw"


@pytest.fixture(autouse=True)
def _auth(monkeypatch):
    monkeypatch.setenv(
        "AIOS_AGENT_TASK_INITIAL_EXECUTION_PROJECTION_AUTHORIZATION_ID",
        AUTHORIZATION,
    )


def canon(value):
    return hashlib.sha256(
        json.dumps(value,sort_keys=True,separators=(",",":")).encode()
    ).hexdigest()


def binding():
    return GovernedExecutionBinding.from_mapping({
        "task_id":TASK_ID,
        "job_id":JOB_ID,
        "repository":"aios/aios",
        "base_oid":BASE,
        "expected_branch_oid":BASE,
        "target_branch":BRANCH,
        "branch_mode":"CREATE_NEW",
        "generation":1,
        "idempotency_key":IDEMPOTENCY,
    })


def runtime_input():
    return {
        "goal":"Supervise release preparation",
        "instructions":"Observe and resume only the same lineage.",
        "success_criteria":[{"id":"heartbeat","description":"Heartbeat observed"}],
        "constraints":{"production_write":False},
        "scope":{"scope_type":"legacy","tenant_id":None,"project_id":None},
        "trace_id":TASK_ID,
    }


def envelope(*, job_id=JOB_ID, action="agent_runtime", capability=CAPABILITY):
    bound=binding()
    if job_id != JOB_ID:
        bound=GovernedExecutionBinding.from_mapping({**bound.public(),"job_id":job_id})
    return project_execution_envelope(
        bound,
        {
            "action":action,
            "adapter":"openclaw" if action=="agent_runtime" else "source_capture",
            "capability_scope":[capability],
            "runtime_input":runtime_input() if action=="agent_runtime" else None,
        },
    )


def make_db(tmp_path, *, claimed_by=None, lease=False, metadata=None, caps=None, event="TASK_CREATED"):
    path=tmp_path/"state.sqlite3"
    con=sqlite3.connect(path)
    con.executescript("""
    CREATE TABLE tasks(
      task_id TEXT PRIMARY KEY,
      goal TEXT NOT NULL,
      instructions TEXT NOT NULL,
      success_criteria_json TEXT NOT NULL,
      required_capabilities_json TEXT NOT NULL,
      constraints_json TEXT NOT NULL,
      metadata_json TEXT NOT NULL,
      deadline TEXT,
      priority INTEGER NOT NULL,
      max_attempts INTEGER NOT NULL,
      attempt INTEGER NOT NULL DEFAULT 0,
      status TEXT NOT NULL,
      claimed_by TEXT,
      lease_token_hash TEXT,
      lease_expires_at TEXT,
      result_json TEXT,
      result_digest TEXT,
      result_idempotency_key TEXT,
      verification_json TEXT,
      revision INTEGER NOT NULL DEFAULT 1,
      created_at TEXT NOT NULL,
      updated_at TEXT NOT NULL
    );
    CREATE TABLE events(
      event_id INTEGER PRIMARY KEY AUTOINCREMENT,
      task_id TEXT,
      agent_id TEXT,
      event_type TEXT NOT NULL,
      payload_json TEXT NOT NULL,
      created_at TEXT NOT NULL
    );
    """)
    metadata={} if metadata is None else metadata
    caps=[] if caps is None else caps
    con.execute(
        "INSERT INTO tasks VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        (
            TASK_ID,"Supervise release preparation",
            "Observe and resume only the same lineage.",
            json.dumps([{"id":"heartbeat","description":"Heartbeat observed"}]),
            json.dumps(caps),json.dumps({"production_write":False}),json.dumps(metadata),
            None,100,20,0,"QUEUED",claimed_by,
            "lease-hash" if lease else None,
            "2099-01-01T00:00:00Z" if lease else None,
            None,None,None,None,1,"2026-08-20T00:00:00Z","2026-08-20T00:00:00Z",
        ),
    )
    con.execute(
        "INSERT INTO events(task_id,agent_id,event_type,payload_json,created_at) VALUES (?,?,?,?,?)",
        (TASK_ID,None,event,"{}","2026-08-20T00:00:00Z"),
    )
    con.commit();con.close()
    return path


def request(*, digest=None, key="initial-projection-a-v1"):
    return InitialExecutionProjectionRequest.from_mapping({
        "expected_status":"QUEUED",
        "expected_revision":1,
        "expected_attempt":0,
        "expected_metadata_sha256":canon({}),
        "expected_required_capabilities_sha256":canon([]),
        "expected_job_id":JOB_ID,
        "expected_execution_digest":digest or envelope()["execution_digest"],
        "required_capability":CAPABILITY,
        "adapter":"openclaw",
        "authorization_id":AUTHORIZATION,
        "idempotency_key":key,
        "reason":"initial governed supervisor runtime projection",
    })


def resolver(_task_id,_job_id,_capability,_adapter,received_runtime_input):
    assert received_runtime_input==runtime_input()
    return envelope()


def source_capture_result():
    old=envelope(action="source_capture",capability="git.read")
    return {
        "task_id":TASK_ID,
        "job_id":JOB_ID,
        "repository":"aios/aios",
        "base_oid":BASE,
        "expected_branch_oid":BASE,
        "target_branch":BRANCH,
        "generation":1,
        "status":"SUCCEEDED",
        "classification":"NEEDS_REWORK_SOURCE_CAPTURE",
        "execution_mode":"source_capture",
        "execution_digest":old["execution_digest"],
        "changed_files":[],
        "writer_lease":"NONE",
        "verification_state":"NEEDS_REWORK",
        "direct_main_write":False,
        "production_write":False,
        "force_push":False,
        "protected_ref_bypass":False,
        "secret_material_exposed":False,
        "source_snapshot":{"schema":"test.snapshot.v1","payload_sha256":"d"*64},
    }


def make_source_capture_rework_db(tmp_path, *, mutate=None, events=None):
    path=make_db(tmp_path)
    result=source_capture_result()
    verification={
        "decision":"NEEDS_REWORK",
        "reason":"SOURCE_CAPTURE_COMPLETE_RUNTIME_PROJECTION_REQUIRED",
    }
    if mutate:
        mutate(result,verification)
    event_types=events or [
        "TASK_CLAIMED","TASK_HEARTBEAT","RESULT_SUBMITTED","TASK_VERIFIED"
    ]
    con=sqlite3.connect(path)
    con.execute(
        "UPDATE tasks SET attempt=1,status='QUEUED',claimed_by=NULL,revision=5,"
        "result_json=?,result_digest=?,result_idempotency_key=?,verification_json=?",
        (
            json.dumps(result,sort_keys=True,separators=(",",":")),
            canon(result),"source-capture-result-v1",
            json.dumps(verification,sort_keys=True,separators=(",",":")),
        ),
    )
    for event_type in event_types:
        payload={}
        if event_type=="TASK_VERIFIED":
            payload={"decision":"NEEDS_REWORK","next_state":"QUEUED"}
        con.execute(
            "INSERT INTO events(task_id,agent_id,event_type,payload_json,created_at) "
            "VALUES (?,?,?,?,?)",
            (TASK_ID,"generic-worker",event_type,json.dumps(payload),"2026-08-20T00:01:00Z"),
        )
    con.commit();con.close()
    return path,result,verification


def recovery_request(result,verification,*,digest=None,key="source-capture-recovery-a-v1"):
    return SourceCaptureProjectionRecoveryRequest.from_mapping({
        "expected_status":"QUEUED",
        "expected_revision":5,
        "expected_attempt":1,
        "expected_metadata_sha256":canon({}),
        "expected_required_capabilities_sha256":canon([]),
        "previous_verification_sha256":canon(verification),
        "expected_source_capture_result_sha256":canon(result),
        "expected_source_capture_execution_digest":result["execution_digest"],
        "expected_job_id":JOB_ID,
        "expected_execution_digest":digest or envelope()["execution_digest"],
        "required_capability":CAPABILITY,
        "adapter":"openclaw",
        "authorization_id":AUTHORIZATION,
        "idempotency_key":key,
        "reason":"first durable runtime projection after governed source capture",
    })


def test_initial_projection_cas_persists_envelope_capability_and_event(tmp_path):
    path=make_db(tmp_path)
    result=InitialExecutionProjectionStore(path,resolver).project(TASK_ID,request())
    assert result["revision"]==2
    assert result["attempt"]==0
    assert result["idempotent"] is False
    con=sqlite3.connect(path)
    row=con.execute(
        "SELECT status,attempt,revision,claimed_by,metadata_json,required_capabilities_json FROM tasks"
    ).fetchone()
    assert row[:4]==("QUEUED",0,2,None)
    metadata=json.loads(row[4])
    assert metadata["execution"]["action"]=="agent_runtime"
    assert metadata["execution"]["adapter"]=="openclaw"
    assert metadata["execution"]["governance"]["writer_lease"] is False
    assert json.loads(row[5])==[CAPABILITY]
    events=con.execute("SELECT event_type,payload_json FROM events ORDER BY event_id").fetchall()
    assert [event[0] for event in events]==[
        "TASK_CREATED","TASK_INITIAL_EXECUTION_PROJECTION_AUTHORIZED"
    ]
    payload=json.loads(events[-1][1])
    assert payload["job_id"]==JOB_ID
    assert payload["main_write"] is False
    assert payload["production_write"] is False


def test_initial_projection_idempotent_exact_replay(tmp_path):
    path=make_db(tmp_path)
    store=InitialExecutionProjectionStore(path,resolver)
    first=store.project(TASK_ID,request())
    second=store.project(TASK_ID,request())
    assert first["idempotent"] is False
    assert second["idempotent"] is True
    assert second["execution_digest"]==first["execution_digest"]
    con=sqlite3.connect(path)
    assert con.execute("SELECT COUNT(*) FROM task_initial_execution_projections").fetchone()[0]==1
    assert con.execute("SELECT COUNT(*) FROM events").fetchone()[0]==2


def test_source_capture_recovery_projects_first_runtime_envelope(tmp_path):
    path,result,verification=make_source_capture_rework_db(tmp_path)
    response=InitialExecutionProjectionStore(path,resolver).project_from_source_capture(
        TASK_ID,recovery_request(result,verification)
    )
    assert response["revision"]==6
    assert response["attempt"]==1
    assert response["recovery_from_source_capture"] is True
    assert response["idempotent"] is False
    con=sqlite3.connect(path)
    row=con.execute(
        "SELECT status,attempt,revision,claimed_by,metadata_json,"
        "required_capabilities_json,result_json,verification_json FROM tasks"
    ).fetchone()
    assert row[:4]==("QUEUED",1,6,None)
    metadata=json.loads(row[4])
    assert metadata["execution"]["action"]=="agent_runtime"
    assert metadata["execution"]["adapter"]=="openclaw"
    assert metadata["execution"]["governance"]["writer_lease"] is False
    assert json.loads(row[5])==[CAPABILITY]
    assert json.loads(row[6])==result
    assert json.loads(row[7])==verification
    events=con.execute("SELECT event_type,payload_json FROM events ORDER BY event_id").fetchall()
    assert [event[0] for event in events]==[
        "TASK_CREATED","TASK_CLAIMED","TASK_HEARTBEAT","RESULT_SUBMITTED",
        "TASK_VERIFIED","TASK_INITIAL_EXECUTION_PROJECTION_AUTHORIZED",
    ]
    payload=json.loads(events[-1][1])
    assert payload["recovery_from_source_capture"] is True
    assert payload["source_capture_execution_digest"]==result["execution_digest"]
    assert payload["main_write"] is False
    assert payload["production_write"] is False


def test_source_capture_recovery_idempotent_exact_replay(tmp_path):
    path,result,verification=make_source_capture_rework_db(tmp_path)
    store=InitialExecutionProjectionStore(path,resolver)
    first=store.project_from_source_capture(TASK_ID,recovery_request(result,verification))
    second=store.project_from_source_capture(TASK_ID,recovery_request(result,verification))
    assert first["idempotent"] is False
    assert second["idempotent"] is True
    assert second["execution_digest"]==first["execution_digest"]
    con=sqlite3.connect(path)
    assert con.execute(
        "SELECT COUNT(*) FROM task_source_capture_projection_recoveries"
    ).fetchone()[0]==1
    assert con.execute("SELECT COUNT(*) FROM events").fetchone()[0]==6


@pytest.mark.parametrize(
    "mutate,error",
    [
        (lambda result,verification: result.__setitem__("changed_files",["x"]),
         "SOURCE_CAPTURE_RESULT_CONTRACT_INVALID"),
        (lambda result,verification: result.__setitem__("job_id","job-other"),
         "SOURCE_CAPTURE_RESULT_CONTRACT_INVALID"),
        (lambda result,verification: verification.__setitem__("decision","SUCCEEDED"),
         "SOURCE_CAPTURE_NEEDS_REWORK_VERIFICATION_REQUIRED"),
    ],
)
def test_source_capture_recovery_rejects_unsafe_result_or_verification(
    tmp_path,mutate,error
):
    path,result,verification=make_source_capture_rework_db(tmp_path,mutate=mutate)
    with pytest.raises(InitialExecutionProjectionError,match=error):
        InitialExecutionProjectionStore(path,resolver).project_from_source_capture(
            TASK_ID,recovery_request(result,verification)
        )


@pytest.mark.parametrize(
    "events",
    [
        ["TASK_CLAIMED","RESULT_SUBMITTED","TASK_VERIFIED"],
        ["TASK_CLAIMED","TASK_HEARTBEAT","RESULT_SUBMITTED","TASK_VERIFIED","TASK_CLAIMED"],
        ["TASK_CLAIMED","TASK_HEARTBEAT","TASK_VERIFIED","RESULT_SUBMITTED"],
    ],
)
def test_source_capture_recovery_rejects_invalid_event_chain(tmp_path,events):
    path,result,verification=make_source_capture_rework_db(tmp_path,events=events)
    with pytest.raises(
        InitialExecutionProjectionError,match="SOURCE_CAPTURE_RECOVERY_EVENT_CHAIN_INVALID"
    ):
        InitialExecutionProjectionStore(path,resolver).project_from_source_capture(
            TASK_ID,recovery_request(result,verification)
        )


def test_source_capture_recovery_rejects_execution_digest_drift(tmp_path):
    path,result,verification=make_source_capture_rework_db(tmp_path)
    with pytest.raises(InitialExecutionProjectionError,match="EXECUTION_DIGEST_DRIFT"):
        InitialExecutionProjectionStore(path,resolver).project_from_source_capture(
            TASK_ID,recovery_request(result,verification,digest="f"*64)
        )


def test_source_capture_recovery_request_is_strict(tmp_path):
    _path,result,verification=make_source_capture_rework_db(tmp_path)
    data=recovery_request(result,verification).public()
    data["metadata_patch"]={"execution":"forbidden"}
    with pytest.raises(InitialExecutionProjectionError,match="FIELDS_INVALID"):
        SourceCaptureProjectionRecoveryRequest.from_mapping(data)


@pytest.mark.parametrize(
    "kwargs,error",
    [
        ({"claimed_by":"worker-a"},"TASK_ALREADY_CLAIMED"),
        ({"lease":True},"ACTIVE_LEASE_DENY"),
        ({"metadata":{"execution":{}}},"TASK_EXECUTION_ALREADY_PROJECTED"),
        ({"caps":["git.read"]},"TASK_INITIAL_CAPABILITIES_MUST_BE_EMPTY"),
        ({"event":"TASK_CLAIMED"},"TASK_INITIAL_EVENT_CHAIN_REQUIRED"),
    ],
)
def test_initial_projection_rejects_noninitial_state(tmp_path,kwargs,error):
    path=make_db(tmp_path,**kwargs)
    with pytest.raises(InitialExecutionProjectionError,match=error):
        InitialExecutionProjectionStore(path,resolver).project(TASK_ID,request())


def test_initial_projection_rejects_binding_drift(tmp_path):
    path=make_db(tmp_path)
    def wrong(*_args):
        return envelope(job_id="job-other")
    with pytest.raises(InitialExecutionProjectionError,match="GOVERNED_BINDING_CONFLICT"):
        InitialExecutionProjectionStore(path,wrong).project(TASK_ID,request())


def test_initial_projection_rejects_nonruntime_envelope(tmp_path):
    path=make_db(tmp_path)
    def wrong(*_args):
        return envelope(action="source_capture",capability="git.read")
    with pytest.raises(InitialExecutionProjectionError):
        InitialExecutionProjectionStore(path,wrong).project(TASK_ID,request())


def test_initial_projection_rejects_execution_digest_drift(tmp_path):
    path=make_db(tmp_path)
    with pytest.raises(InitialExecutionProjectionError,match="EXECUTION_DIGEST_DRIFT"):
        InitialExecutionProjectionStore(path,resolver).project(
            TASK_ID,request(digest="f"*64)
        )


def test_request_is_narrow_and_initial_state_bound():
    data=request().public()
    data["arbitrary_metadata"]={"x":1}
    with pytest.raises(InitialExecutionProjectionError,match="FIELDS_INVALID"):
        InitialExecutionProjectionRequest.from_mapping(data)
    data=request().public();data["expected_attempt"]=1
    with pytest.raises(InitialExecutionProjectionError,match="INITIAL_STATE_REQUIRED"):
        InitialExecutionProjectionRequest.from_mapping(data)


def test_initial_projection_unit_is_unprivileged_and_owner_credential_bound():
    unit=(ROOT/"deploy/systemd/aios-agent-task-initial-execution-projection.service").read_text()
    assert "User=aios-agent-task" in unit
    assert "Group=aios-agent-task" in unit
    assert "LoadCredential=agent_task_control_token:" in unit
    assert "Environment=AIOS_GOVERNED_WRITE_API_URL=http://127.0.0.1:18892" in unit
    assert "RestrictAddressFamilies=AF_UNIX AF_INET" in unit
    assert "ReadWritePaths=/var/lib/aios-agent-task" in unit
    assert "User=root" not in unit