from collections.abc import Iterator
from dataclasses import dataclass

import pytest
from sqlalchemy import Engine
from sqlmodel import Session

from domain.tasks.subtasks import SubTask, SubTaskStatus
from domain.tasks.tasks import Source, TaskStatus
from orchestration.task_orchestrator import TaskOrchestrator
from repositories.tasks.subtask_postgres_repository import SubTaskPostgresRepository
from repositories.tasks.task_postgres_repository import TaskPostgresRepository


@dataclass
class Harness:
    orchestrator: TaskOrchestrator
    tasks: TaskPostgresRepository
    subtasks: SubTaskPostgresRepository


@pytest.fixture
def harness(db_engine: Engine) -> Iterator[Harness]:
    with Session(db_engine) as session:
        tasks = TaskPostgresRepository(session=session)
        subtasks = SubTaskPostgresRepository(session=session)
        yield Harness(
            orchestrator=TaskOrchestrator(task_repo=tasks, subtask_repo=subtasks),
            tasks=tasks,
            subtasks=subtasks,
        )


def test_create_task_with_subtask_creates_both_in_waiting(
    harness: Harness,
) -> None:
    # act
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test",
        inputs={"foo": "bar"},
        source=Source.PORTFOLIO,
        source_id="abc",
    )

    # assert
    assert task.status is TaskStatus.WAITING
    assert subtask.status is SubTaskStatus.WAITING
    assert subtask.task_id == task.id
    assert subtask.inputs == {"foo": "bar"}


def test_start_subtask_cascades_to_in_progress(harness: Harness) -> None:
    # arrange
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )

    # act
    started = harness.orchestrator.start_subtask(
        subtask.id, cloud_logs_url="https://example/log"
    )

    # assert
    assert started.status is SubTaskStatus.IN_PROGRESS
    assert started.cloud_logs_url == "https://example/log"
    assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS


def test_complete_subtask_cascades_to_complete(harness: Harness) -> None:
    # arrange
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    harness.orchestrator.start_subtask(subtask.id)

    # act
    harness.orchestrator.complete_subtask(subtask.id, {"value": 42})

    # assert
    done_subtask = harness.subtasks.get(subtask.id)
    done_task = harness.tasks.get(task.id)
    assert done_subtask.outputs == {"result": {"value": 42}}
    assert done_task.status is TaskStatus.COMPLETE
    assert done_task.job_completed is not None


def test_fail_subtask_cascades_to_failed(harness: Harness) -> None:
    # arrange
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )

    # act
    harness.orchestrator.fail_subtask(subtask.id, RuntimeError("boom"))

    # assert
    failed_subtask = harness.subtasks.get(subtask.id)
    failed_task = harness.tasks.get(task.id)
    assert failed_subtask.outputs == {"error": "boom"}
    assert failed_task.status is TaskStatus.FAILED


def test_failed_subtask_locks_task_failed_even_with_others_complete(
    harness: Harness,
) -> None:
    # arrange
    task, first = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    second = SubTask.create(task_id=task.id)
    harness.subtasks.create(second)

    # act
    harness.orchestrator.complete_subtask(first.id)
    harness.orchestrator.fail_subtask(second.id, RuntimeError("nope"))

    # assert
    assert harness.tasks.get(task.id).status is TaskStatus.FAILED


def test_mixed_complete_and_in_progress_keeps_task_in_progress(
    harness: Harness,
) -> None:
    # arrange
    task, first = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    second = SubTask.create(task_id=task.id)
    harness.subtasks.create(second)

    # act
    harness.orchestrator.complete_subtask(first.id)
    harness.orchestrator.start_subtask(second.id)

    # assert
    assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS


def test_run_subtask_happy_path_returns_result_and_cascades_complete(
    harness: Harness,
) -> None:
    # arrange
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )

    # act
    result = harness.orchestrator.run_subtask(subtask.id, work=lambda: {"answer": 42})

    # assert
    assert result == {"answer": 42}
    assert harness.subtasks.get(subtask.id).status is SubTaskStatus.COMPLETE
    assert harness.tasks.get(task.id).status is TaskStatus.COMPLETE


def test_create_child_subtask_adds_waiting_child_without_changing_parent_status(
    harness: Harness,
) -> None:
    # arrange
    task, first = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    harness.orchestrator.start_subtask(first.id)
    assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS

    # act
    child = harness.orchestrator.create_child_subtask(
        task.id, inputs={"split": "a"}
    )

    # assert
    persisted_child = harness.subtasks.get(child.id)
    assert persisted_child.task_id == task.id
    assert persisted_child.status is SubTaskStatus.WAITING
    assert persisted_child.inputs == {"split": "a"}
    assert persisted_child.id != first.id
    # Cascade is a no-op: parent stays IN_PROGRESS.
    assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS


def test_run_subtask_failing_work_marks_failed_and_reraises(
    harness: Harness,
) -> None:
    # arrange
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )

    def boom() -> None:
        raise RuntimeError("boom")

    # act / assert
    with pytest.raises(RuntimeError, match="boom"):
        harness.orchestrator.run_subtask(subtask.id, work=boom)

    assert harness.subtasks.get(subtask.id).status is SubTaskStatus.FAILED
    assert harness.tasks.get(task.id).status is TaskStatus.FAILED


def test_run_subtasks_creates_runs_and_completes_a_whole_batch(
    harness: Harness,
) -> None:
    """run_subtasks fans the parent into one child per item, runs each, and
    leaves every child COMPLETE with the parent cascaded to COMPLETE."""
    # arrange — a parent task whose coordinator subtask is done
    task, coordinator = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    harness.orchestrator.complete_subtask(coordinator.id)

    # act
    results = harness.orchestrator.run_subtasks(
        task.id,
        [{"property_id": 1}, {"property_id": 2}, {"property_id": 3}],
        work=lambda st: (st.inputs or {})["property_id"] * 10,
    )

    # assert — results in order, all children COMPLETE, parent COMPLETE
    assert results == [10, 20, 30]
    children = [s for s in harness.subtasks.list_by_task(task.id) if s.id != coordinator.id]
    assert len(children) == 3
    assert all(c.status is SubTaskStatus.COMPLETE for c in children)
    assert {(c.inputs or {})["property_id"] for c in children} == {1, 2, 3}
    assert harness.tasks.get(task.id).status is TaskStatus.COMPLETE


def test_run_subtasks_isolates_a_failing_item_and_continues(
    harness: Harness,
) -> None:
    """A raising item is marked FAILED with its error recorded, its siblings still
    complete (result None for the failure), and the parent cascades to FAILED."""
    # arrange
    task, coordinator = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    harness.orchestrator.complete_subtask(coordinator.id)

    def _work(st: SubTask) -> int:
        pid = (st.inputs or {})["property_id"]
        if pid == 2:
            raise RuntimeError("property 2 exploded")
        return pid

    # act — does NOT raise; failure is isolated
    results = harness.orchestrator.run_subtasks(
        task.id,
        [{"property_id": 1}, {"property_id": 2}, {"property_id": 3}],
        work=_work,
    )

    # assert
    assert results == [1, None, 3]
    children = {
        (c.inputs or {})["property_id"]: c
        for c in harness.subtasks.list_by_task(task.id)
        if c.id != coordinator.id
    }
    assert children[1].status is SubTaskStatus.COMPLETE
    assert children[3].status is SubTaskStatus.COMPLETE
    assert children[2].status is SubTaskStatus.FAILED
    assert children[2].outputs == {"error": "property 2 exploded"}
    # any FAILED child → parent FAILED
    assert harness.tasks.get(task.id).status is TaskStatus.FAILED


def test_cascade_short_circuits_once_task_already_failed(harness: Harness) -> None:
    """Once the Task is FAILED, completing another SubTask leaves it FAILED — the
    terminal state is not recomputed away."""
    # arrange — two children; fail the first so the task is FAILED
    task, coordinator = harness.orchestrator.create_task_with_subtask(
        task_source="manual:test"
    )
    second = harness.orchestrator.create_child_subtask(task.id)
    harness.orchestrator.fail_subtask(coordinator.id, RuntimeError("boom"))
    assert harness.tasks.get(task.id).status is TaskStatus.FAILED

    # act — complete the other child
    harness.orchestrator.complete_subtask(second.id)

    # assert — task stays FAILED (cascade short-circuited on the terminal state)
    assert harness.tasks.get(task.id).status is TaskStatus.FAILED