import pytest

from dobby.models import (
    Artifact,
    InfraConfig,
    InvalidTransitionError,
    Priority,
    Project,
    ProjectPolicies,
    Result,
    Schedule,
    ScheduleType,
    Task,
    TaskStatus,
)
from dobby.store.sqlite import SCHEMA_VERSION


async def test_create_and_get_task(store):
    task = Task(project="spare", description="Test task")
    task_id = await store.create_task(task)
    retrieved = await store.get_task(task_id)
    assert retrieved.project == "spare"
    assert retrieved.description == "Test task"
    assert retrieved.status == TaskStatus.PENDING


async def test_get_nonexistent_task_raises(store):
    with pytest.raises(KeyError):
        await store.get_task("t-nonexistent")


async def test_update_task(store):
    task = Task(project="spare", description="Test")
    task_id = await store.create_task(task)
    await store.update_task(task_id, status=TaskStatus.RUNNING.value)
    updated = await store.get_task(task_id)
    assert updated.status == TaskStatus.RUNNING


async def test_list_tasks_filter_by_project(store):
    await store.create_task(Task(project="spare", description="A"))
    await store.create_task(Task(project="memagent", description="B"))
    await store.create_task(Task(project="spare", description="C"))
    tasks = await store.list_tasks(project="spare")
    assert len(tasks) == 2
    assert all(t.project == "spare" for t in tasks)


async def test_list_tasks_filter_by_status(store):
    t1 = Task(project="spare", description="A")
    t2 = Task(project="spare", description="B")
    await store.create_task(t1)
    tid2 = await store.create_task(t2)
    await store.update_task(tid2, status=TaskStatus.RUNNING.value)
    tasks = await store.list_tasks(status="pending")
    assert len(tasks) == 1


async def test_get_pending_tasks_priority_order(store):
    await store.create_task(Task(project="a", description="low", priority=Priority.LOW))
    await store.create_task(Task(project="b", description="urgent", priority=Priority.URGENT))
    await store.create_task(Task(project="c", description="normal", priority=Priority.NORMAL))
    pending = await store.get_pending_tasks(limit=10)
    assert pending[0].priority == Priority.URGENT
    assert pending[1].priority == Priority.NORMAL
    assert pending[2].priority == Priority.LOW


async def test_get_scheduled_tasks(store):
    await store.create_task(Task(
        project="spare", description="Nightly",
        schedule=Schedule(type=ScheduleType.CRON, expr="0 3 * * *"),
    ))
    await store.create_task(Task(project="spare", description="One-off"))
    scheduled = await store.get_scheduled_tasks()
    assert len(scheduled) == 1
    assert scheduled[0].schedule.type == ScheduleType.CRON


async def test_store_and_get_result(store):
    result = Result(
        task_id="t-20260321-001",
        status="completed",
        summary="Done",
        artifacts=[Artifact(type="pr", url="https://example.com/pull/1")],
        metrics={"loss": 0.5},
    )
    await store.store_result(result)
    retrieved = await store.get_result("t-20260321-001")
    assert retrieved.summary == "Done"
    assert retrieved.metrics == {"loss": 0.5}
    assert len(retrieved.artifacts) == 1
    assert retrieved.artifacts[0].type == "pr"


async def test_get_result_nonexistent(store):
    result = await store.get_result("t-nonexistent")
    assert result is None


async def test_list_results_by_project(store):
    # Need tasks first for project association
    await store.create_task(Task(id="t-1", project="spare", description="A"))
    await store.create_task(Task(id="t-2", project="memagent", description="B"))
    await store.store_result(Result(task_id="t-1", status="completed", summary="A done"))
    await store.store_result(Result(task_id="t-2", status="completed", summary="B done"))
    results = await store.list_results(project="spare")
    assert len(results) == 1
    assert results[0].task_id == "t-1"


async def test_list_results_last_n(store):
    for i in range(5):
        tid = f"t-{i}"
        await store.create_task(Task(id=tid, project="spare", description=f"Task {i}"))
        await store.store_result(Result(task_id=tid, status="completed", summary=f"Done {i}"))
    results = await store.list_results(last_n=3)
    assert len(results) == 3


async def test_register_and_get_project(store):
    project = Project(
        name="spare",
        repo="/path/to/spare",
        git_remote="git@github.com:user/spare.git",
        context="RL framework",
        infra=[InfraConfig(name="gpu-1", type="ssh", host="user@gpu")],
    )
    await store.register_project(project)
    retrieved = await store.get_project("spare")
    assert retrieved.repo == "/path/to/spare"
    assert retrieved.git_remote == "git@github.com:user/spare.git"
    assert len(retrieved.infra) == 1


async def test_get_project_nonexistent(store):
    result = await store.get_project("nope")
    assert result is None


async def test_list_projects(store):
    await store.register_project(Project(name="a", repo="/a"))
    await store.register_project(Project(name="b", repo="/b"))
    projects = await store.list_projects()
    assert len(projects) == 2


async def test_register_project_upsert(store):
    await store.register_project(Project(name="spare", repo="/old"))
    await store.register_project(Project(name="spare", repo="/new"))
    p = await store.get_project("spare")
    assert p.repo == "/new"


# ---------------------------------------------------------------------------
# DB migration / versioning
# ---------------------------------------------------------------------------

async def test_schema_version_set(store):
    version = await store.get_schema_version()
    assert version == SCHEMA_VERSION


async def test_schema_version_fresh_db(tmp_path):
    """A fresh DB should be created at the latest schema version."""
    from dobby.store.sqlite import SqliteStore
    s = SqliteStore(str(tmp_path / "fresh.db"))
    await s.initialize()
    version = await s.get_schema_version()
    assert version == SCHEMA_VERSION
    await s.close()


async def test_migration_v1_to_v2(tmp_path):
    """Simulate a v1 DB and verify migration adds events/attempt_number."""
    import json

    import aiosqlite
    db_path = str(tmp_path / "v1.db")
    # Create v1 schema manually
    async with aiosqlite.connect(db_path) as db:
        await db.execute("PRAGMA journal_mode=WAL")
        await db.execute("CREATE TABLE tasks (id TEXT PRIMARY KEY, data TEXT NOT NULL)")
        await db.execute(
            "CREATE TABLE results"
            " (id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL, data TEXT NOT NULL)"
        )
        await db.execute("CREATE TABLE projects (name TEXT PRIMARY KEY, data TEXT NOT NULL)")
        await db.execute(
            "CREATE TABLE schema_version"
            " (id INTEGER PRIMARY KEY CHECK (id = 1), version INTEGER NOT NULL)"
        )
        await db.execute("INSERT INTO schema_version (id, version) VALUES (1, 1)")
        # Insert a v1 task (no events/attempt_number)
        task_data = json.dumps({
            "id": "t-test-v1", "project": "p", "description": "old task",
            "mode": "quick", "priority": "normal", "exit_criteria": None,
            "schedule": {
                "type": "once", "expr": None, "every": None, "until": None, "depends_on": None,
            },
            "infra": None, "timeout": None, "checkpoint": False, "plan": [],
            "session_id": None, "backend": "claude",
            "created_at": "2026-01-01T00:00:00+00:00",
            "started_at": None, "last_run_at": None, "status": "pending",
        })
        await db.execute("INSERT INTO tasks (id, data) VALUES (?, ?)", ("t-test-v1", task_data))
        # Insert a v1 result (no attempt_number)
        result_data = json.dumps({
            "task_id": "t-test-v1", "status": "completed", "summary": "old result",
            "artifacts": [], "metrics": {}, "gpu_info": None,
            "cost": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "cost_usd": 0.0},
            "started_at": None, "completed_at": "2026-01-01T01:00:00+00:00",
        })
        await db.execute(
            "INSERT INTO results (task_id, data) VALUES (?, ?)", ("t-test-v1", result_data)
        )
        await db.commit()

    # Now open with SqliteStore — should run migration
    from dobby.store.sqlite import SqliteStore
    s = SqliteStore(db_path)
    await s.initialize()
    version = await s.get_schema_version()
    assert version == SCHEMA_VERSION

    task = await s.get_task("t-test-v1")
    assert task.attempt_number == 0
    assert len(task.events) >= 1
    assert task.events[0].detail == "created"

    results = await s.list_results(task_id="t-test-v1")
    assert results[0].attempt_number == 0
    await s.close()


# ---------------------------------------------------------------------------
# State transition guardrails in store
# ---------------------------------------------------------------------------

async def test_update_task_invalid_transition_raises(store):
    task = Task(project="spare", description="Test")
    await store.create_task(task)
    with pytest.raises(InvalidTransitionError):
        # PENDING -> COMPLETED is not allowed (must go through RUNNING first)
        await store.update_task(task.id, status=TaskStatus.COMPLETED.value)


async def test_update_task_records_event(store):
    task = Task(project="spare", description="Test")
    await store.create_task(task)
    await store.update_task(task.id, status=TaskStatus.RUNNING.value)
    updated = await store.get_task(task.id)
    assert len(updated.events) == 2  # created + running
    assert updated.events[1].status == "running"


# ---------------------------------------------------------------------------
# Search and sort in list_tasks
# ---------------------------------------------------------------------------

async def test_list_tasks_search(store):
    await store.create_task(Task(project="spare", description="Add warmup scheduler"))
    await store.create_task(Task(project="spare", description="Fix broken tests"))
    await store.create_task(Task(project="spare", description="Refactor scheduler code"))
    tasks = await store.list_tasks(search="scheduler")
    assert len(tasks) == 2


async def test_list_tasks_sort_priority(store):
    await store.create_task(Task(project="a", description="low", priority=Priority.LOW))
    await store.create_task(Task(project="b", description="urgent", priority=Priority.URGENT))
    await store.create_task(Task(project="c", description="normal", priority=Priority.NORMAL))
    tasks = await store.list_tasks(sort="-priority")
    assert tasks[0].priority == Priority.URGENT
    assert tasks[2].priority == Priority.LOW


# ---------------------------------------------------------------------------
# Schema migration v2 -> v3
# ---------------------------------------------------------------------------

async def test_migration_v2_to_v3(tmp_path):
    """Simulate a v2 DB and verify migration adds next_run_at/max_failures/consecutive_failures."""
    import json

    import aiosqlite
    db_path = str(tmp_path / "v2.db")
    async with aiosqlite.connect(db_path) as db:
        await db.execute("PRAGMA journal_mode=WAL")
        await db.execute("CREATE TABLE tasks (id TEXT PRIMARY KEY, data TEXT NOT NULL)")
        await db.execute(
            "CREATE TABLE results"
            " (id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL, data TEXT NOT NULL)"
        )
        await db.execute("CREATE TABLE projects (name TEXT PRIMARY KEY, data TEXT NOT NULL)")
        await db.execute(
            "CREATE TABLE schema_version"
            " (id INTEGER PRIMARY KEY CHECK (id = 1), version INTEGER NOT NULL)"
        )
        await db.execute("INSERT INTO schema_version (id, version) VALUES (1, 2)")
        task_data = json.dumps({
            "id": "t-test-v2", "project": "p", "description": "v2 task",
            "mode": "quick", "priority": "normal", "exit_criteria": None,
            "schedule": {
                "type": "once", "expr": None, "every": None, "until": None, "depends_on": None,
            },
            "infra": None, "timeout": None, "checkpoint": False, "plan": [],
            "session_id": None, "backend": "claude",
            "created_at": "2026-01-01T00:00:00+00:00",
            "started_at": None, "last_run_at": None, "attempt_number": 0,
            "events": [
                {"status": "pending", "timestamp": "2026-01-01T00:00:00+00:00",
                 "detail": "created"},
            ],
            "status": "pending",
        })
        await db.execute("INSERT INTO tasks (id, data) VALUES (?, ?)", ("t-test-v2", task_data))
        await db.commit()

    from dobby.store.sqlite import SqliteStore
    s = SqliteStore(db_path)
    await s.initialize()
    version = await s.get_schema_version()
    assert version == SCHEMA_VERSION

    task = await s.get_task("t-test-v2")
    assert task.max_failures == 3
    assert task.consecutive_failures == 0
    assert task.next_run_at is None
    await s.close()


# ---------------------------------------------------------------------------
# Failure tracking persistence
# ---------------------------------------------------------------------------

async def test_update_task_consecutive_failures(store):
    task = Task(
        project="spare", description="Recurring fail",
        schedule=Schedule(type=ScheduleType.INTERVAL, every="30m"),
    )
    await store.create_task(task)
    await store.update_task(task.id, status=TaskStatus.RUNNING.value)
    await store.update_task(
        task.id,
        status=TaskStatus.ACTIVE.value,
        consecutive_failures=1,
    )
    updated = await store.get_task(task.id)
    assert updated.consecutive_failures == 1
    assert updated.max_failures == 3


# ---------------------------------------------------------------------------
# get_dependent_tasks
# ---------------------------------------------------------------------------

async def test_get_dependent_tasks(store):
    parent = Task(project="spare", description="Build")
    pid = await store.create_task(parent)
    child1 = Task(
        project="spare", description="Test after build",
        schedule=Schedule(type=ScheduleType.AFTER, depends_on=pid),
    )
    child2 = Task(
        project="other", description="Deploy after build",
        schedule=Schedule(type=ScheduleType.AFTER, depends_on=pid),
    )
    unrelated = Task(project="spare", description="Unrelated")
    await store.create_task(child1)
    await store.create_task(child2)
    await store.create_task(unrelated)

    deps = await store.get_dependent_tasks(pid)
    assert len(deps) == 2
    dep_ids = {d.id for d in deps}
    assert child1.id in dep_ids
    assert child2.id in dep_ids


async def test_get_dependent_tasks_empty(store):
    task = Task(project="spare", description="Lonely")
    tid = await store.create_task(task)
    deps = await store.get_dependent_tasks(tid)
    assert deps == []


# ---------------------------------------------------------------------------
# Project policies persistence
# ---------------------------------------------------------------------------

async def test_project_policies_persist(store):
    project = Project(
        name="policed",
        repo="/path/policed",
        policies=ProjectPolicies(
            default_backend="codex",
            default_timeout=3600,
            max_concurrent=2,
            default_priority="high",
            auto_chain=True,
        ),
    )
    await store.register_project(project)
    retrieved = await store.get_project("policed")
    assert retrieved.policies.default_backend == "codex"
    assert retrieved.policies.default_timeout == 3600
    assert retrieved.policies.max_concurrent == 2
    assert retrieved.policies.default_priority == "high"
    assert retrieved.policies.auto_chain is True


# ---------------------------------------------------------------------------
# Schema migration v4 -> v5
# ---------------------------------------------------------------------------

async def test_migration_v4_to_v5(tmp_path):
    """Simulate a v4 DB and verify migration adds policies to projects."""
    import json

    import aiosqlite
    db_path = str(tmp_path / "v4.db")
    async with aiosqlite.connect(db_path) as db:
        await db.execute("PRAGMA journal_mode=WAL")
        await db.execute("CREATE TABLE tasks (id TEXT PRIMARY KEY, data TEXT NOT NULL)")
        await db.execute(
            "CREATE TABLE results"
            " (id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL, data TEXT NOT NULL)"
        )
        await db.execute("CREATE TABLE projects (name TEXT PRIMARY KEY, data TEXT NOT NULL)")
        await db.execute(
            "CREATE TABLE schema_version"
            " (id INTEGER PRIMARY KEY CHECK (id = 1), version INTEGER NOT NULL)"
        )
        await db.execute("INSERT INTO schema_version (id, version) VALUES (1, 4)")
        # Insert a v4 project (no policies field)
        project_data = json.dumps({
            "name": "oldproj", "repo": "/old", "git_remote": None,
            "context": "test", "claude_md": True, "infra": [],
        })
        await db.execute(
            "INSERT INTO projects (name, data) VALUES (?, ?)", ("oldproj", project_data)
        )
        await db.commit()

    from dobby.store.sqlite import SqliteStore
    s = SqliteStore(db_path)
    await s.initialize()
    version = await s.get_schema_version()
    assert version == SCHEMA_VERSION

    project = await s.get_project("oldproj")
    assert project.policies.auto_chain is False
    assert project.policies.max_concurrent is None
    assert project.policies.default_backend is None
    await s.close()
