"""
Tests for Workflow CRUD, Execution, and Chain REST endpoints.

Uses FastAPI TestClient against api_server.app with isolated temp directories
for workflow definitions and runs.
"""

import json
import uuid
from pathlib import Path
from unittest.mock import patch, MagicMock

import pytest
from fastapi.testclient import TestClient


# ── Fixtures ─────────────────────────────────────────────────────────────


@pytest.fixture(autouse=True)
def _isolate_broker_and_workflows(tmp_path, monkeypatch):
    """Redirect all broker + workflow storage to a temp directory and disable rate limiting."""
    # Build the full broker directory tree
    for subdir in (
        "channels", "sessions", "cursors", "tasks", "capabilities",
        "profiles", "jobs", "violations", "memory", "locks", "agents",
        "archive", "api_keys",
        "workflows", "workflows/definitions", "workflows/runs",
    ):
        (tmp_path / subdir).mkdir(parents=True, exist_ok=True)

    import brynq.server as srv
    monkeypatch.setattr(srv, "BROKER_DIR", tmp_path)
    monkeypatch.setattr(srv, "CHANNELS_DIR", tmp_path / "channels")
    monkeypatch.setattr(srv, "SESSIONS_DIR", tmp_path / "sessions")
    monkeypatch.setattr(srv, "CURSORS_DIR", tmp_path / "cursors")

    import brynq.workflow as w
    monkeypatch.setattr(w, "WORKFLOWS_DIR", tmp_path / "workflows")
    monkeypatch.setattr(w, "DEFINITIONS_DIR", tmp_path / "workflows" / "definitions")
    monkeypatch.setattr(w, "RUNS_DIR", tmp_path / "workflows" / "runs")

    # Redirect api_server's KEYS_DIR to temp
    import brynq.api_server as api
    monkeypatch.setattr(api, "KEYS_DIR", tmp_path / "api_keys")

    # Disable the per-IP rate limiter so tests don't hit 429s
    import brynq.api_security as sec
    monkeypatch.setattr(sec._ip_rate_limiter, "is_allowed", lambda ip: (True, {"limit": 9999, "remaining": 9999, "window": 60}))


@pytest.fixture()
def client():
    """TestClient that suppresses the lifespan (no background tasks)."""
    from brynq.api_server import app

    # Disable the lifespan to avoid starting background tasks / orchestrator
    with patch.object(app, "router"):
        pass
    return TestClient(app, raise_server_exceptions=False)


@pytest.fixture()
def api_key(client, tmp_path):
    """Register an agent and return a valid API key."""
    resp = client.post("/auth/register", json={
        "name": "test-agent",
        "platform": "test",
        "capabilities": [],
    })
    assert resp.status_code == 200, resp.text
    return resp.json()["api_key"]


@pytest.fixture()
def auth_headers(api_key):
    return {"X-API-Key": api_key}


# ── Helpers ──────────────────────────────────────────────────────────────


def _sample_workflow_payload() -> dict:
    return {
        "name": "CI Pipeline",
        "description": "Build, test, deploy",
        "steps": [
            {"step_id": "plan", "type": "task", "dependencies": [], "inputs": {}, "config": {}},
            {"step_id": "build", "type": "task", "dependencies": ["plan"], "inputs": {}, "config": {}},
            {"step_id": "test", "type": "task", "dependencies": ["build"], "inputs": {}, "config": {}},
            {"step_id": "deploy", "type": "task", "dependencies": ["test"], "inputs": {}, "config": {}},
        ],
    }


# ── Workflow CRUD Tests ──────────────────────────────────────────────────


class TestWorkflowCRUD:
    def test_create_workflow(self, client, auth_headers):
        resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        assert resp.status_code == 200, resp.text
        data = resp.json()
        assert data["workflow_id"]
        assert data["name"] == "CI Pipeline"
        assert len(data["steps"]) == 4

    def test_create_workflow_no_name(self, client, auth_headers):
        resp = client.post("/workflows", json={"name": "", "steps": []}, headers=auth_headers)
        assert resp.status_code == 422

    def test_create_workflow_invalid_step_type(self, client, auth_headers):
        resp = client.post("/workflows", json={
            "name": "Bad",
            "steps": [{"step_id": "s1", "type": "magic"}],
        }, headers=auth_headers)
        assert resp.status_code == 422

    def test_create_workflow_duplicate_step_ids(self, client, auth_headers):
        resp = client.post("/workflows", json={
            "name": "Dups",
            "steps": [
                {"step_id": "dup", "type": "task"},
                {"step_id": "dup", "type": "task"},
            ],
        }, headers=auth_headers)
        assert resp.status_code == 422

    def test_create_workflow_unknown_dependency(self, client, auth_headers):
        resp = client.post("/workflows", json={
            "name": "Bad Deps",
            "steps": [{"step_id": "s1", "type": "task", "dependencies": ["nonexistent"]}],
        }, headers=auth_headers)
        assert resp.status_code == 422

    def test_list_workflows_empty(self, client, auth_headers):
        resp = client.get("/workflows", headers=auth_headers)
        assert resp.status_code == 200
        assert resp.json()["workflows"] == []

    def test_list_workflows(self, client, auth_headers):
        client.post("/workflows", json={"name": "WF-A", "steps": []}, headers=auth_headers)
        client.post("/workflows", json={"name": "WF-B", "steps": []}, headers=auth_headers)
        resp = client.get("/workflows", headers=auth_headers)
        assert resp.status_code == 200
        assert len(resp.json()["workflows"]) == 2

    def test_get_workflow(self, client, auth_headers):
        create_resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        resp = client.get(f"/workflows/{wf_id}", headers=auth_headers)
        assert resp.status_code == 200
        assert resp.json()["name"] == "CI Pipeline"

    def test_get_workflow_not_found(self, client, auth_headers):
        resp = client.get("/workflows/nonexistent", headers=auth_headers)
        assert resp.status_code == 404

    def test_delete_workflow(self, client, auth_headers):
        create_resp = client.post("/workflows", json={"name": "Delete Me", "steps": []}, headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        resp = client.delete(f"/workflows/{wf_id}", headers=auth_headers)
        assert resp.status_code == 200
        assert resp.json()["ok"] is True
        # Verify it's gone
        resp2 = client.get(f"/workflows/{wf_id}", headers=auth_headers)
        assert resp2.status_code == 404

    def test_delete_workflow_not_found(self, client, auth_headers):
        resp = client.delete("/workflows/nonexistent", headers=auth_headers)
        assert resp.status_code == 404


# ── Workflow Execution Tests ─────────────────────────────────────────────


class TestWorkflowExecution:
    def test_execute_workflow(self, client, auth_headers):
        create_resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        resp = client.post(f"/workflows/{wf_id}/execute", json={"original_input": "test data"}, headers=auth_headers)
        assert resp.status_code == 200, resp.text
        data = resp.json()
        assert data["run_id"]
        assert data["status"] == "running"
        assert data["inputs"]["original_input"] == "test data"
        # First step (no deps) should be ready
        assert data["step_statuses"]["plan"]["status"] == "ready"
        assert data["step_statuses"]["build"]["status"] == "pending"

    def test_execute_workflow_not_found(self, client, auth_headers):
        resp = client.post("/workflows/nonexistent/execute", json={}, headers=auth_headers)
        assert resp.status_code == 404

    def test_get_run(self, client, auth_headers):
        create_resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        exec_resp = client.post(f"/workflows/{wf_id}/execute", json={}, headers=auth_headers)
        run_id = exec_resp.json()["run_id"]
        resp = client.get(f"/workflows/runs/{run_id}", headers=auth_headers)
        assert resp.status_code == 200
        assert resp.json()["run_id"] == run_id
        assert resp.json()["status"] == "running"

    def test_get_run_not_found(self, client, auth_headers):
        resp = client.get("/workflows/runs/nonexistent", headers=auth_headers)
        assert resp.status_code == 404

    def test_complete_step(self, client, auth_headers):
        create_resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        exec_resp = client.post(f"/workflows/{wf_id}/execute", json={}, headers=auth_headers)
        run_id = exec_resp.json()["run_id"]

        # Complete the first step
        resp = client.post(
            f"/workflows/runs/{run_id}/steps/plan/complete",
            json={"result": "Planning done"},
            headers=auth_headers,
        )
        assert resp.status_code == 200, resp.text
        data = resp.json()
        assert data["steps_completed"] == 1
        assert "build" in data["newly_ready"]

    def test_complete_step_run_not_found(self, client, auth_headers):
        resp = client.post(
            "/workflows/runs/nonexistent/steps/plan/complete",
            json={"result": "x"},
            headers=auth_headers,
        )
        assert resp.status_code == 404

    def test_complete_step_step_not_found(self, client, auth_headers):
        create_resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        exec_resp = client.post(f"/workflows/{wf_id}/execute", json={}, headers=auth_headers)
        run_id = exec_resp.json()["run_id"]
        resp = client.post(
            f"/workflows/runs/{run_id}/steps/nonexistent/complete",
            json={"result": "x"},
            headers=auth_headers,
        )
        assert resp.status_code == 404

    def test_full_lifecycle(self, client, auth_headers):
        """Walk a 4-step workflow from start to completion via the API."""
        create_resp = client.post("/workflows", json=_sample_workflow_payload(), headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        exec_resp = client.post(f"/workflows/{wf_id}/execute", json={"original_input": "go"}, headers=auth_headers)
        run_id = exec_resp.json()["run_id"]

        for step_id in ("plan", "build", "test", "deploy"):
            resp = client.post(
                f"/workflows/runs/{run_id}/steps/{step_id}/complete",
                json={"result": f"{step_id} done"},
                headers=auth_headers,
            )
            assert resp.status_code == 200

        # Verify the run is completed
        run_resp = client.get(f"/workflows/runs/{run_id}", headers=auth_headers)
        assert run_resp.json()["status"] == "completed"
        assert run_resp.json()["completed_at"] is not None

    def test_parallel_steps(self, client, auth_headers):
        """Two steps with the same dependency should both become ready."""
        payload = {
            "name": "Parallel",
            "steps": [
                {"step_id": "init", "type": "task", "dependencies": []},
                {"step_id": "branch_a", "type": "task", "dependencies": ["init"]},
                {"step_id": "branch_b", "type": "task", "dependencies": ["init"]},
                {"step_id": "merge", "type": "task", "dependencies": ["branch_a", "branch_b"]},
            ],
        }
        create_resp = client.post("/workflows", json=payload, headers=auth_headers)
        wf_id = create_resp.json()["workflow_id"]
        exec_resp = client.post(f"/workflows/{wf_id}/execute", json={}, headers=auth_headers)
        run_id = exec_resp.json()["run_id"]

        # Complete init -> both branches ready
        resp = client.post(
            f"/workflows/runs/{run_id}/steps/init/complete",
            json={"result": "done"},
            headers=auth_headers,
        )
        newly_ready = resp.json()["newly_ready"]
        assert "branch_a" in newly_ready
        assert "branch_b" in newly_ready

        # Complete one branch -> merge still pending
        client.post(f"/workflows/runs/{run_id}/steps/branch_a/complete", json={"result": "a done"}, headers=auth_headers)
        run_data = client.get(f"/workflows/runs/{run_id}", headers=auth_headers).json()
        assert run_data["step_statuses"]["merge"]["status"] == "pending"

        # Complete other branch -> merge ready
        resp2 = client.post(f"/workflows/runs/{run_id}/steps/branch_b/complete", json={"result": "b done"}, headers=auth_headers)
        assert "merge" in resp2.json()["newly_ready"]


# ── Chain Endpoint Tests ─────────────────────────────────────────────────


class TestChainEndpoint:
    def test_chain_success(self, client, auth_headers, monkeypatch):
        """Chain with 3 steps, mock LLM, verify interpolation and output."""
        call_log = []

        def mock_chat(backend, model, messages, **kwargs):
            prompt = messages[0]["content"]
            call_log.append({"backend": backend, "model": model, "prompt": prompt})
            return f"output_from_{backend}_{model}"

        monkeypatch.setattr("brynq.api_server.chat", mock_chat, raising=False)
        # Also patch at the import location used inside the endpoint
        monkeypatch.setattr("brynq.llm_router.chat", mock_chat)

        resp = client.post("/chain", json={
            "input": "Analyze Q4 revenue",
            "steps": [
                {"backend": "claude", "model": "claude-sonnet-4-20250514", "prompt": "Analyze: {original_input}"},
                {"backend": "gemini", "model": "gemini-2.5-pro", "prompt": "Review: {step_1_output}"},
                {"backend": "ollama", "model": "llama3", "prompt": "Summarize:\n{step_1_output}\n{step_2_output}"},
            ],
        }, headers=auth_headers)

        assert resp.status_code == 200, resp.text
        data = resp.json()
        assert data["status"] == "completed"
        assert data["run_id"]

        # Verify interpolation happened correctly
        assert call_log[0]["prompt"] == "Analyze: Analyze Q4 revenue"
        assert call_log[1]["prompt"] == "Review: output_from_claude_claude-sonnet-4-20250514"
        assert "output_from_claude_claude-sonnet-4-20250514" in call_log[2]["prompt"]
        assert "output_from_gemini_gemini-2.5-pro" in call_log[2]["prompt"]

        # Verify response structure
        assert data["final_output"] == "output_from_ollama_llama3"
        assert "step_1" in data["step_outputs"]
        assert "step_2" in data["step_outputs"]
        assert "step_3" in data["step_outputs"]

    def test_chain_single_step(self, client, auth_headers, monkeypatch):
        """Chain with a single step should work."""
        monkeypatch.setattr("brynq.llm_router.chat", lambda *a, **kw: "single result")

        resp = client.post("/chain", json={
            "input": "Hello",
            "steps": [
                {"backend": "auto", "model": "test", "prompt": "Say hi to {original_input}"},
            ],
        }, headers=auth_headers)

        assert resp.status_code == 200
        data = resp.json()
        assert data["status"] == "completed"
        assert data["final_output"] == "single result"

    def test_chain_llm_failure(self, client, auth_headers, monkeypatch):
        """If an LLM call fails, chain returns 502 with partial results."""
        call_count = [0]

        def mock_chat(backend, model, messages, **kwargs):
            call_count[0] += 1
            if call_count[0] == 2:
                raise RuntimeError("LLM backend unavailable")
            return f"output_{call_count[0]}"

        monkeypatch.setattr("brynq.llm_router.chat", mock_chat)

        resp = client.post("/chain", json={
            "input": "test",
            "steps": [
                {"backend": "auto", "model": "m1", "prompt": "{original_input}"},
                {"backend": "auto", "model": "m2", "prompt": "{step_1_output}"},
                {"backend": "auto", "model": "m3", "prompt": "{step_2_output}"},
            ],
        }, headers=auth_headers)

        assert resp.status_code == 502
        data = resp.json()
        assert data["status"] == "failed"
        assert data["failed_step"] == "step_2"
        assert "LLM backend unavailable" in data["error"]
        # First step output should be present
        assert "step_1" in data["step_outputs"]

    def test_chain_empty_steps(self, client, auth_headers):
        """Chain with no steps should be rejected by Pydantic validation."""
        resp = client.post("/chain", json={
            "input": "test",
            "steps": [],
        }, headers=auth_headers)
        assert resp.status_code == 422

    def test_chain_empty_input(self, client, auth_headers):
        """Chain with empty input should be rejected."""
        resp = client.post("/chain", json={
            "input": "",
            "steps": [{"backend": "auto", "model": "m", "prompt": "go"}],
        }, headers=auth_headers)
        assert resp.status_code == 422

    def test_chain_cleans_up_workflow(self, client, auth_headers, monkeypatch):
        """The ad-hoc workflow definition should be deleted after chain completes."""
        monkeypatch.setattr("brynq.llm_router.chat", lambda *a, **kw: "ok")

        # Count workflows before
        before = client.get("/workflows", headers=auth_headers).json()["workflows"]

        resp = client.post("/chain", json={
            "input": "test",
            "steps": [{"backend": "auto", "model": "m", "prompt": "{original_input}"}],
        }, headers=auth_headers)
        assert resp.status_code == 200

        # No lingering workflow definitions
        after = client.get("/workflows", headers=auth_headers).json()["workflows"]
        assert len(after) == len(before)


# ── Auth Tests (ensure auth is enforced) ─────────────────────────────────


class TestWorkflowAuth:
    def test_create_requires_auth(self, client):
        resp = client.post("/workflows", json={"name": "test", "steps": []})
        assert resp.status_code in (401, 403)

    def test_list_requires_auth(self, client):
        resp = client.get("/workflows")
        assert resp.status_code in (401, 403)

    def test_execute_requires_auth(self, client):
        resp = client.post("/workflows/some-id/execute", json={})
        assert resp.status_code in (401, 403)

    def test_chain_requires_auth(self, client):
        resp = client.post("/chain", json={
            "input": "test",
            "steps": [{"backend": "auto", "model": "m", "prompt": "go"}],
        })
        assert resp.status_code in (401, 403)

    def test_get_run_requires_auth(self, client):
        resp = client.get("/workflows/runs/some-id")
        assert resp.status_code in (401, 403)

    def test_complete_step_requires_auth(self, client):
        resp = client.post("/workflows/runs/r/steps/s/complete", json={"result": "x"})
        assert resp.status_code in (401, 403)

    def test_delete_requires_auth(self, client):
        resp = client.delete("/workflows/some-id")
        assert resp.status_code in (401, 403)
