import json
from pathlib import Path
from unittest.mock import patch

import pytest
from screenshot import ScreenshotOptions

from screenshot_batch import (
    ScreenshotBatchCoordinator,
    ScreenshotBatchResult,
    ScreenshotJobSpec,
    ScreenshotMultiRunCoordinator,
)
from screenshot_batch.cli import (
    _build_spec,
    _iter_jsonl,
    _make_progress_callback,
)
from screenshot_batch.cli import (
    main as cli_main,
)
from screenshot_batch.coordinator import _unwrap_manifest


def test_job_spec_payload_merges_metadata():
    spec = ScreenshotJobSpec.from_url("https://example.com", metadata={"priority": "high"})
    payload = spec.to_payload(default_metadata={"source": "test"})

    assert payload["job_id"]
    assert payload["metadata"]["source"] == "test"
    assert payload["metadata"]["priority"] == "high"
    assert payload["options"]["capture"]["enabled"] is True


def test_unwrap_manifest_prefers_nested_payload():
    manifest = {"manifest": {"storage_path": "/tmp/example", "files": []}}
    inner = _unwrap_manifest(manifest)
    assert inner is manifest["manifest"]


@patch("screenshot_batch.coordinator.ScreenshotBatchCoordinator.run_batch_sync")
def test_multi_run_coordinator_chunks(mock_run_batch_sync, tmp_path):
    def fake_run_batch_sync(self, specs):
        return ScreenshotBatchResult(
            batch_id="batch",
            job_id=f"run-{len(calls)}",
            status="completed",
            job_total=len(specs),
            job_completed=len(specs),
            job_failed=0,
            status_payload={},
            manifest_payload={},
            local_root=str(tmp_path / f"run-{len(calls)}"),
        )

    calls = []

    def record_call(*args, **kwargs):
        specs_arg = args[-1] if args else kwargs.get("specs")
        calls.append(tuple(spec.job_id for spec in specs_arg))
        return fake_run_batch_sync(None, specs_arg)

    mock_run_batch_sync.side_effect = record_call

    runner = ScreenshotMultiRunCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
        download_results=False,
    )

    specs = [
        ScreenshotJobSpec(
            job_id=f"job-{index}",
            url=f"https://example.com/{index}",
            options=ScreenshotOptions(),
        )
        for index in range(5)
    ]

    summaries = runner.run_jobs(specs, chunk_size=2, max_parallel_runs=1)

    assert len(calls) == 3
    assert len(summaries) == 3
    assert [summary.result.job_total for summary in summaries] == [2, 2, 1]


@patch("screenshot_batch.coordinator.ScreenshotBatchCoordinator.run_batch_sync")
def test_progress_callback_invoked(mock_run_batch_sync, tmp_path):
    def fake_run(specs):
        return ScreenshotBatchResult(
            batch_id="batch",
            job_id="run",
            status="completed",
            job_total=len(specs),
            job_completed=len(specs),
            job_failed=0,
            status_payload={},
            manifest_payload={},
            local_root=None,
        )

    mock_run_batch_sync.side_effect = fake_run

    runner = ScreenshotMultiRunCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=None,
        download_results=False,
    )

    specs = [
        ScreenshotJobSpec(job_id=f"job-{i}", url=f"https://example.com/{i}", options=ScreenshotOptions())
        for i in range(3)
    ]

    events = []

    def callback(completed, total, summary):
        events.append((completed, total, summary.chunk_index))

    summaries = runner.run_jobs(specs, chunk_size=1, max_parallel_runs=1, progress_callback=callback)

    assert events
    assert events[-1][0] == events[-1][1]
    assert all(summary.job_results for summary in summaries)


def test_iter_jsonl_and_build_spec(tmp_path):
    sample = tmp_path / "jobs.jsonl"
    records = [
        {
            "url": "https://example.com",
            "job_id": "job-1",
            "metadata": {"priority": "high"},
            "options": {"enabled": True, "max_pages": 2},
        },
        {
            "website": "https://example.org",
        },
    ]
    sample.write_text("\n".join(json.dumps(item) for item in records), encoding="utf-8")

    rows = list(_iter_jsonl(sample))
    assert len(rows) == 2

    spec = _build_spec(rows[0], {"source": "test"})
    assert spec.job_id == "job-1"
    assert spec.metadata["source"] == "test"
    assert spec.metadata["priority"] == "high"


def test_cli(tmp_path, monkeypatch):
    input_path = tmp_path / "jobs.jsonl"
    input_path.write_text(
        "\n".join(
            [
                json.dumps({"url": "https://example.com"}),
                json.dumps({"company": "invalid"}),
            ]
        )
        + "\n",
        encoding="utf-8",
    )

    monkeypatch.setenv("SCREENSHOT_FUNC_BASE_URL", "https://example.com")

    output_path = tmp_path / "results.jsonl"

    with patch("screenshot_batch.coordinator.ScreenshotMultiRunCoordinator.run_jobs", return_value=[]):
        exit_code = cli_main(
            [
                str(input_path),
                "batch-1",
                "--chunk-size",
                "4",
                "--output",
                str(output_path),
            ]
        )
    assert exit_code == 0

    lines = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines() if line.strip()]
    assert len(lines) == 2
    assert lines[0]["status"] == "unknown"
    assert lines[0]["blob_url"] is None
    assert lines[1]["status"] == "skipped_missing_url"
    assert lines[1]["blob_url"] is None


def test_collect_job_results_single_job_defaults_to_success(tmp_path):
    runner = ScreenshotMultiRunCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
        download_results=True,
    )

    run_dir = tmp_path / "chunk-00001"
    run_dir.mkdir(parents=True, exist_ok=True)

    result = ScreenshotBatchResult(
        batch_id="batch",
        job_id="run",
        status="completed",
        job_total=1,
        job_completed=1,
        job_failed=0,
        status_payload={},
        manifest_payload={},
        local_root=str(run_dir),
    )

    spec = ScreenshotJobSpec(
        job_id="job-1",
        url="https://example.com",
        options=ScreenshotOptions(),
    )

    outcomes = runner._collect_job_results((spec,), result, run_dir)
    summary = outcomes["job-1"]
    assert summary.status == "succeeded"
    assert summary.local_path == run_dir


@patch("screenshot_batch.coordinator.ScreenshotBatchCoordinator.run_batch_sync")
@patch.object(ScreenshotMultiRunCoordinator, "_upload_run_artifacts")
def test_run_jobs_assigns_blob_urls(mock_upload, mock_run, tmp_path):
    run_dir = tmp_path / "chunk-00001"
    job_dir = run_dir / "job-1"
    (job_dir / "metadata").mkdir(parents=True, exist_ok=True)
    metadata_path = job_dir / "metadata" / "external_screenshots.json"
    metadata_path.write_text(json.dumps({"entries": [{"status": "captured"}], "errors": []}), encoding="utf-8")

    def fake_run_batch_sync(specs):
        run_dir.mkdir(parents=True, exist_ok=True)
        return ScreenshotBatchResult(
            batch_id="batch",
            job_id="run",
            status="completed",
            job_total=len(specs),
            job_completed=len(specs),
            job_failed=0,
            status_payload={},
            manifest_payload={},
            local_root=str(run_dir),
        )

    mock_run.side_effect = fake_run_batch_sync
    mock_upload.return_value = {"job-1": "https://blob.example/yc-companies/job-1"}

    runner = ScreenshotMultiRunCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
        download_results=True,
        upload_results=True,
    )
    runner._uploader = lambda path, blob_path=None: None  # type: ignore[attr-defined]

    specs = [ScreenshotJobSpec(job_id="job-1", url="https://example.com", options=ScreenshotOptions())]

    summaries = runner.run_jobs(specs, chunk_size=1, max_parallel_runs=1)
    summary = summaries[0].job_results["job-1"]
    assert summary.status == "succeeded"
    assert summary.blob_url == "https://blob.example/yc-companies/job-1"
    mock_upload.assert_called_once()


def test_build_blob_name_with_manifest_prefix(tmp_path):
    runner = ScreenshotBatchCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
    )

    result = runner._build_blob_name("prefix", "/screens/result.png", tmp_path / "result.png")
    assert result == "prefix/screens/result.png"


def test_build_blob_name_without_manifest_prefix(tmp_path):
    runner = ScreenshotBatchCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
    )

    result = runner._build_blob_name(None, "/screens/result.png", tmp_path / "artifact.txt")
    assert result == "screens/result.png"


def test_collect_job_results_without_local_root(tmp_path):
    runner = ScreenshotMultiRunCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
        download_results=False,
    )

    specs = (
        ScreenshotJobSpec(job_id="job-1", url="https://example.com/1", options=ScreenshotOptions()),
        ScreenshotJobSpec(job_id="job-2", url="https://example.com/2", options=ScreenshotOptions()),
    )

    result = ScreenshotBatchResult(
        batch_id="batch",
        job_id="run",
        status="processing",
        job_total=2,
        job_completed=2,
        job_failed=0,
        status_payload={},
        manifest_payload={},
        local_root=None,
    )

    outcomes = runner._collect_job_results(specs, result, None)
    assert len(outcomes) == 2
    assert all(summary.status == "processing" for summary in outcomes.values())
    assert all(summary.local_path is None for summary in outcomes.values())


def test_cli_requires_base_url(monkeypatch, tmp_path):
    monkeypatch.delenv("SCREENSHOT_FUNC_BASE_URL", raising=False)
    input_path = tmp_path / "jobs.jsonl"
    input_path.write_text(json.dumps({"url": "https://example.com"}), encoding="utf-8")

    exit_code = cli_main([str(input_path), "batch-1"])
    assert exit_code == 1


def test_make_progress_callback_none():
    assert _make_progress_callback(0) is None


def test_ensure_source_file_downloads(monkeypatch, tmp_path):
    runner = ScreenshotBatchCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
    )

    monkeypatch.setattr(
        "screenshot_batch.coordinator.ScreenshotBatchCoordinator._build_blob_name",
        lambda self, manifest_prefix, rel_path, target: "prefix/job.json",
    )

    def fake_download(self, blob_name: str, target: Path) -> bool:  # type: ignore[unused-variable]
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(b"payload")
        return True

    monkeypatch.setattr(
        "screenshot_batch.coordinator.ScreenshotBatchCoordinator._download_from_screenshot_storage",
        fake_download,
    )

    target = tmp_path / "downloads" / "job.json"
    assert runner._ensure_source_file(target, "prefix", "job.json")
    assert target.read_bytes() == b"payload"


def test_download_blob_prefix(tmp_path):
    class DummyStorage:
        def __init__(self) -> None:
            self.invocations = 0

        def download_prefix(self, prefix: str, dest_dir: Path) -> bool:
            dest_dir.mkdir(parents=True, exist_ok=True)
            (dest_dir / "artifact.txt").write_text("data", encoding="utf-8")
            self.invocations += 1
            return True

        def download_file(self, blob_name: str, dest: Path) -> bool:  # pragma: no cover - helper completeness
            return False

    storage = DummyStorage()
    runner = ScreenshotBatchCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
        storage=storage,
    )

    run_root = tmp_path / "run"
    run_root.mkdir()
    assert runner._download_blob_prefix("prefix/", run_root)
    assert storage.invocations == 1
    assert (run_root / "artifact.txt").exists()


@pytest.mark.asyncio
async def test_materialise_manifest_copies_files(tmp_path):
    runner = ScreenshotBatchCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
    )

    storage_root = tmp_path / "storage"
    file_path = storage_root / "screens" / "base.json"
    file_path.parent.mkdir(parents=True, exist_ok=True)
    file_path.write_text("body", encoding="utf-8")

    manifest = {
        "storage_path": str(storage_root),
        "files": [{"path": "screens/base.json"}],
    }

    result = await runner._materialise_manifest("run-1", manifest)

    assert result is not None
    assert result == runner.store_dir / "run-1"
    assert (result / "screens" / "base.json").read_text(encoding="utf-8") == "body"


@pytest.mark.asyncio
async def test_materialise_manifest_uses_blob_prefix(monkeypatch, tmp_path):
    runner = ScreenshotBatchCoordinator(
        base_url="https://example.com",
        batch_id="batch",
        store_dir=tmp_path,
    )

    monkeypatch.setattr(
        "screenshot_batch.coordinator.ScreenshotBatchCoordinator._download_blob_prefix",
        lambda self, blob_prefix, run_root: True,
    )

    manifest = {
        "storage_path": "",
        "files": [],
        "blob_prefix": "prefix",
    }

    result = await runner._materialise_manifest("run-2", manifest)

    assert result == runner.store_dir / "run-2"
