# grantflow/tests/test_integration_hitl_flows_exports.py

from fastapi.testclient import TestClient


from grantflow.api.app import app

from grantflow.tests.integration_helpers import drain_hitl_to_done, wait_for_terminal_status

client = TestClient(app)


def test_hitl_pause_resume_flow_supports_export_payload_and_export():
    response = client.post(
        "/generate",
        json={
            "donor_id": "usaid",
            "input_context": {"project": "Education", "country": "Kenya"},
            "llm_mode": False,
            "hitl_enabled": True,
        },
    )
    assert response.status_code == 200
    job_id = response.json()["job_id"]

    status = wait_for_terminal_status(client, job_id)
    assert status["status"] in {"pending_hitl", "done"}
    if status["status"] != "done":
        status = drain_hitl_to_done(client, job_id, initial_status=status, max_cycles=12)
    assert status["status"] == "done"

    export_payload_response = client.get(f"/status/{job_id}/export-payload")
    assert export_payload_response.status_code == 200
    export_payload_body = export_payload_response.json()
    payload = export_payload_body.get("payload")
    assert isinstance(payload, dict)
    payload_state = payload.get("state")
    assert isinstance(payload_state, dict)
    assert payload_state.get("hitl_pending") is False
    assert "strategy" not in payload_state
    assert "donor_strategy" not in payload_state

    critic_findings = payload.get("critic_findings")
    assert isinstance(critic_findings, list)
    for finding in critic_findings:
        assert isinstance(finding, dict)
        finding_id = str(finding.get("finding_id") or "").strip()
        assert finding_id
        assert str(finding.get("id") or "").strip() == finding_id

    export = client.post("/export", json={"payload": payload, "format": "both"})
    assert export.status_code == 200
    assert export.headers["content-type"].startswith("application/zip")
    assert export.content[:2] == b"PK"


def test_hitl_reject_then_resume_flow_supports_export_payload_and_export():
    response = client.post(
        "/generate",
        json={
            "donor_id": "usaid",
            "input_context": {"project": "Education", "country": "Kenya"},
            "llm_mode": False,
            "hitl_enabled": True,
            "hitl_checkpoints": ["toc"],
        },
    )
    assert response.status_code == 200
    job_id = response.json()["job_id"]

    status = wait_for_terminal_status(client, job_id)
    assert status["status"] == "pending_hitl"
    assert status["checkpoint_stage"] == "toc"
    first_checkpoint_id = str(status["checkpoint_id"] or "")
    assert first_checkpoint_id

    reject = client.post(
        "/hitl/approve",
        json={"checkpoint_id": first_checkpoint_id, "approved": False, "feedback": "Need rewrite before MEL"},
    )
    assert reject.status_code == 200
    reject_body = reject.json()
    assert reject_body["status"] == "rejected"

    resume = client.post(f"/resume/{job_id}", json={})
    assert resume.status_code == 200
    assert resume.json()["resuming_from"] == "architect"

    status = wait_for_terminal_status(client, job_id)
    assert status["status"] == "pending_hitl"
    assert status["checkpoint_stage"] == "toc"
    second_checkpoint_id = str(status["checkpoint_id"] or "")
    assert second_checkpoint_id
    assert second_checkpoint_id != first_checkpoint_id

    approve = client.post(
        "/hitl/approve",
        json={"checkpoint_id": second_checkpoint_id, "approved": True, "feedback": "Rewrite accepted"},
    )
    assert approve.status_code == 200

    resume = client.post(f"/resume/{job_id}", json={})
    assert resume.status_code == 200
    assert resume.json()["resuming_from"] == "mel"

    status = wait_for_terminal_status(client, job_id)
    if status["status"] != "done":
        status = drain_hitl_to_done(client, job_id, initial_status=status, max_cycles=12)
    assert status["status"] == "done"

    export_payload_response = client.get(f"/status/{job_id}/export-payload")
    assert export_payload_response.status_code == 200
    payload = export_payload_response.json().get("payload")
    assert isinstance(payload, dict)
    payload_state = payload.get("state")
    assert isinstance(payload_state, dict)
    assert payload_state.get("hitl_pending") is False
    assert "strategy" not in payload_state
    assert "donor_strategy" not in payload_state

    export = client.post("/export", json={"payload": payload, "format": "both"})
    assert export.status_code == 200
    assert export.headers["content-type"].startswith("application/zip")
    assert export.content[:2] == b"PK"


def test_hitl_logframe_reject_then_resume_flow_supports_export_payload_and_export():
    response = client.post(
        "/generate",
        json={
            "donor_id": "usaid",
            "input_context": {"project": "Education", "country": "Kenya"},
            "llm_mode": False,
            "hitl_enabled": True,
            "hitl_checkpoints": ["logframe"],
        },
    )
    assert response.status_code == 200
    job_id = response.json()["job_id"]

    status = wait_for_terminal_status(client, job_id)
    assert status["status"] == "pending_hitl"
    assert status["checkpoint_stage"] == "logframe"
    first_checkpoint_id = str(status["checkpoint_id"] or "")
    assert first_checkpoint_id

    reject = client.post(
        "/hitl/approve",
        json={"checkpoint_id": first_checkpoint_id, "approved": False, "feedback": "Need MEL adjustments"},
    )
    assert reject.status_code == 200
    reject_body = reject.json()
    assert reject_body["status"] == "rejected"

    resume = client.post(f"/resume/{job_id}", json={})
    assert resume.status_code == 200
    assert resume.json()["resuming_from"] == "mel"

    status = wait_for_terminal_status(client, job_id)
    assert status["status"] == "pending_hitl"
    assert status["checkpoint_stage"] == "logframe"
    second_checkpoint_id = str(status["checkpoint_id"] or "")
    assert second_checkpoint_id
    assert second_checkpoint_id != first_checkpoint_id

    approve = client.post(
        "/hitl/approve",
        json={"checkpoint_id": second_checkpoint_id, "approved": True, "feedback": "MEL accepted"},
    )
    assert approve.status_code == 200

    resume = client.post(f"/resume/{job_id}", json={})
    assert resume.status_code == 200
    assert resume.json()["resuming_from"] == "critic"

    status = wait_for_terminal_status(client, job_id)
    if status["status"] != "done":
        status = drain_hitl_to_done(client, job_id, initial_status=status, max_cycles=12)
    assert status["status"] == "done"

    export_payload_response = client.get(f"/status/{job_id}/export-payload")
    assert export_payload_response.status_code == 200
    payload = export_payload_response.json().get("payload")
    assert isinstance(payload, dict)
    payload_state = payload.get("state")
    assert isinstance(payload_state, dict)
    assert payload_state.get("hitl_pending") is False
    assert "strategy" not in payload_state
    assert "donor_strategy" not in payload_state

    export = client.post("/export", json={"payload": payload, "format": "both"})
    assert export.status_code == 200
    assert export.headers["content-type"].startswith("application/zip")
    assert export.content[:2] == b"PK"
