"""Tests for TES models."""

from __future__ import annotations

from ai4gh.specs.tes.models import (
    TESCreateTaskResponse,
    TESExecutor,
    TESInput,
    TESListTasksResponse,
    TESOutput,
    TESResources,
    TESState,
    TESTask,
    TESTaskLog,
    TESView,
)


class TestTESState:
    def test_enum_values(self) -> None:
        assert TESState.UNKNOWN == "UNKNOWN"
        assert TESState.QUEUED == "QUEUED"
        assert TESState.RUNNING == "RUNNING"
        assert TESState.COMPLETE == "COMPLETE"
        assert TESState.EXECUTOR_ERROR == "EXECUTOR_ERROR"
        assert TESState.CANCELED == "CANCELED"
        assert TESState.PREEMPTED == "PREEMPTED"

    def test_from_string(self) -> None:
        assert TESState("RUNNING") is TESState.RUNNING


class TestTESView:
    def test_enum_values(self) -> None:
        assert TESView.MINIMAL == "MINIMAL"
        assert TESView.BASIC == "BASIC"
        assert TESView.FULL == "FULL"


class TestTESInput:
    def test_required_fields(self) -> None:
        inp = TESInput(path="/data/input.bam")
        assert inp.path == "/data/input.bam"
        assert inp.url is None
        assert inp.content is None

    def test_full(self) -> None:
        inp = TESInput(
            name="reads",
            description="Input BAM file",
            url="s3://bucket/input.bam",
            path="/data/input.bam",
            type="FILE",
            streamable=True,
        )
        assert inp.name == "reads"
        assert inp.url == "s3://bucket/input.bam"
        assert inp.streamable is True

    def test_inline_content(self) -> None:
        inp = TESInput(
            path="/data/script.sh",
            content="#!/bin/bash\necho hello",
        )
        assert inp.content.startswith("#!/bin/bash")


class TestTESOutput:
    def test_required_fields(self) -> None:
        out = TESOutput(url="s3://bucket/output.bam", path="/data/output.bam")
        assert out.url == "s3://bucket/output.bam"
        assert out.path == "/data/output.bam"

    def test_with_name(self) -> None:
        out = TESOutput(
            name="aligned",
            url="s3://bucket/aligned.bam",
            path="/data/aligned.bam",
            type="FILE",
        )
        assert out.name == "aligned"


class TestTESExecutor:
    def test_required_fields(self) -> None:
        ex = TESExecutor(image="ubuntu:20.04", command=["echo", "hello"])
        assert ex.image == "ubuntu:20.04"
        assert ex.command == ["echo", "hello"]
        assert ex.workdir is None

    def test_full(self) -> None:
        ex = TESExecutor(
            image="biocontainers/bwa:0.7.17",
            command=["bwa", "mem", "ref.fa", "reads.fq"],
            workdir="/work",
            stdin="/data/input.txt",
            stdout="/logs/stdout.log",
            stderr="/logs/stderr.log",
            env={"THREADS": "4"},
        )
        assert ex.workdir == "/work"
        assert ex.env == {"THREADS": "4"}


class TestTESResources:
    def test_defaults(self) -> None:
        res = TESResources()
        assert res.cpu_cores is None
        assert res.ram_gb is None
        assert res.preemptible is None

    def test_full(self) -> None:
        res = TESResources(
            cpu_cores=4,
            ram_gb=8.0,
            disk_gb=100.0,
            preemptible=True,
            zones=["us-east-1a"],
        )
        assert res.cpu_cores == 4
        assert res.ram_gb == 8.0
        assert res.preemptible is True


class TestTESTaskLog:
    def test_minimal(self) -> None:
        log = TESTaskLog()
        assert log.logs is None
        assert log.start_time is None

    def test_with_data(self) -> None:
        log = TESTaskLog(
            start_time="2024-01-01T00:00:00Z",
            end_time="2024-01-01T01:00:00Z",
            metadata={"host": "node-1"},
            system_logs=["Starting task"],
        )
        assert log.metadata == {"host": "node-1"}
        assert len(log.system_logs) == 1


class TestTESTask:
    def test_minimal(self) -> None:
        task = TESTask()
        assert task.id is None
        assert task.state is None
        assert task.executors is None

    def test_full_task(self) -> None:
        task = TESTask(
            id="task-1",
            state=TESState.COMPLETE,
            name="alignment",
            description="Align reads to reference",
            executors=[TESExecutor(image="bwa:latest", command=["bwa", "mem"])],
            inputs=[TESInput(path="/data/reads.fq", url="s3://bucket/reads.fq")],
            outputs=[TESOutput(url="s3://bucket/aligned.bam", path="/data/aligned.bam")],
            resources=TESResources(cpu_cores=4, ram_gb=8.0),
            tags={"project": "genomics"},
        )
        assert task.id == "task-1"
        assert task.state is TESState.COMPLETE
        assert len(task.executors) == 1
        assert task.tags == {"project": "genomics"}

    def test_from_dict(self) -> None:
        task = TESTask.model_validate(
            {
                "id": "task-2",
                "state": "RUNNING",
                "name": "variant_calling",
            }
        )
        assert task.state is TESState.RUNNING


class TestTESCreateTaskResponse:
    def test_creation(self) -> None:
        resp = TESCreateTaskResponse(id="task-new")
        assert resp.id == "task-new"


class TestTESListTasksResponse:
    def test_empty(self) -> None:
        resp = TESListTasksResponse()
        assert resp.tasks is None
        assert resp.next_page_token is None

    def test_with_tasks(self) -> None:
        resp = TESListTasksResponse(
            tasks=[TESTask(id="t1"), TESTask(id="t2")],
            next_page_token="page-2",
        )
        assert len(resp.tasks) == 2
        assert resp.next_page_token == "page-2"


class TestExtraFields:
    def test_allows_extra(self) -> None:
        task = TESTask.model_validate({"id": "task-1", "vendor_specific": "value"})
        assert task.vendor_specific == "value"
