"""End-to-end tests for screenshot-client-batch against live screenshot service.

Run with:
    AZURE_STORAGE_CONNECTION_STRING="..." \
    RUN_E2E=1 \
    SCREENSHOT_FUNC_BASE_URL="https://your-service.azurewebsites.net" \
        pytest tests/test_e2e.py -v

Stress tests (longer duration, higher load):
    AZURE_STORAGE_CONNECTION_STRING="..." \
    RUN_E2E=1 \
    RUN_STRESS_TEST=1 \
    SCREENSHOT_FUNC_BASE_URL="https://your-service.azurewebsites.net" \
        pytest tests/test_e2e.py -v -k stress

Required environment variables:
    RUN_E2E=1 - Enable E2E tests
    SCREENSHOT_FUNC_BASE_URL - Screenshot service URL
    AZURE_STORAGE_CONNECTION_STRING - Azure Storage connection for blob verification

Optional:
    SCREENSHOT_FUNC_API_KEY - API key for authenticated endpoints
    RUN_STRESS_TEST=1 - Enable stress tests (longer duration, higher load)
"""

from __future__ import annotations

import os
import time
import uuid
from collections.abc import Iterable
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Any

import pytest

from screenshot_batch import ScreenshotServiceClient

# Same test sites as service E2E tests
SITES_UNDER_TEST = [
    "https://landonorris.com/",
    "https://wonjyou.studio/",
    "https://algon.iq/",
]

# Sites for stress testing - diverse set to exercise different rendering scenarios
STRESS_TEST_SITES = [
    "https://example.com/",
    "https://httpbin.org/html",
    "https://www.wikipedia.org/",
    "https://news.ycombinator.com/",
    "https://jsonplaceholder.typicode.com/",
    "https://www.google.com/",
    "https://www.rust-lang.org/",
    "https://go.dev/",
]


def _require_e2e() -> None:
    """Skip test if RUN_E2E environment variable is not set."""
    if os.getenv("RUN_E2E") != "1":
        pytest.skip("Set RUN_E2E=1 to enable E2E tests")


def _require_base_url() -> str:
    """Get base URL from environment."""
    base_url = os.getenv("SCREENSHOT_FUNC_BASE_URL")
    if not base_url:
        pytest.skip("SCREENSHOT_FUNC_BASE_URL not set")
    return base_url.rstrip("/")


def _manifest_files(payload: dict[str, Any]) -> Iterable[dict[str, Any]]:
    """Extract files list from manifest payload."""
    manifest = payload.get("manifest")
    if isinstance(manifest, dict):
        return manifest.get("files", [])
    return payload.get("files", [])


def _verify_blobs_exist(blob_prefix: str, files: list[dict[str, Any]]) -> None:
    """Verify that files in manifest exist in Azure blob storage.

    Args:
        blob_prefix: Base blob path from manifest.
        files: List of file entries from manifest.

    Raises:
        AssertionError: If any files are missing.
    """
    connection_string = os.getenv("AZURE_STORAGE_CONNECTION_STRING")
    if not connection_string:
        pytest.skip("AZURE_STORAGE_CONNECTION_STRING not set - cannot verify blob uploads")

    from azure.storage.blob import BlobServiceClient

    blob_service = BlobServiceClient.from_connection_string(connection_string)
    container_name = os.getenv("AZURE_STORAGE_CONTAINER", "screenshot-output")
    container_client = blob_service.get_container_client(container_name)

    missing = []
    for file_entry in files:
        file_path = file_entry.get("path", "")
        if not file_path:
            continue

        blob_name = f"{blob_prefix}/{file_path}"
        blob_client = container_client.get_blob_client(blob_name)

        try:
            blob_client.get_blob_properties()
        except Exception:
            missing.append(blob_name)

    if missing:
        pytest.fail(
            f"Blob verification failed: {len(missing)} file(s) missing:\n"
            + "\n".join(f"  - {name}" for name in missing[:10])
            + (f"\n  ... and {len(missing) - 10} more" if len(missing) > 10 else "")
        )


class TestScreenshotServiceClientE2E:
    """E2E tests for ScreenshotServiceClient."""

    def test_simple_screenshot_job(self) -> None:
        """Run a simple screenshot job and verify completion."""
        _require_e2e()
        base_url = _require_base_url()

        client = ScreenshotServiceClient(base_url=base_url)
        batch_id = f"e2e-batch-client-{uuid.uuid4().hex[:8]}"

        payload = {
            "jobs": [
                {
                    "job_id": "simple-test",
                    "url": SITES_UNDER_TEST[0],
                    "options": {
                        "capture": {
                            "enabled": True,
                        }
                    },
                }
            ],
            "concurrency": 1,
        }

        # Start batch
        handle = client.start_batch(batch_id, payload)
        assert handle.batch_id == batch_id
        assert handle.job_id is not None

        # Poll until completion
        job_id = str(handle.job_id)
        status = client.poll_run(batch_id, job_id, timeout=180.0, poll_interval=5.0)
        assert status.get("status") in {"completed", "failed"}

        # Get result and verify manifest structure
        result = client.get_result(batch_id, job_id)
        manifest = result.get("manifest", {})
        storage = manifest.get("storage", {})
        blob_prefix = storage.get("prefix")
        assert blob_prefix, "Manifest missing storage.prefix"

        files = list(_manifest_files(result))
        assert isinstance(files, list)

    def test_screenshot_captures_site(self) -> None:
        """End-to-end capture that verifies files are uploaded to blob storage."""
        _require_e2e()
        base_url = _require_base_url()

        client = ScreenshotServiceClient(base_url=base_url)
        batch_id = f"e2e-batch-client-{uuid.uuid4().hex[:8]}"

        payload = {
            "jobs": [
                {
                    "job_id": "capture-test",
                    "url": SITES_UNDER_TEST[0],
                    "options": {
                        "capture": {
                            "enabled": True,
                            "max_pages": 1,
                            "depth": 0,
                            "viewports": ["desktop"],
                            "post_nav_wait_s": 2.0,
                            "timeout_s": 120.0,
                            "max_capture_attempts": 1,
                        }
                    },
                }
            ],
            "concurrency": 1,
        }

        handle = client.start_batch(batch_id, payload)
        job_id = str(handle.job_id)

        status = client.poll_run(batch_id, job_id, timeout=300.0, poll_interval=5.0)
        assert status.get("status") == "completed", f"Job failed: {status}"

        result = client.get_result(batch_id, job_id)
        files = list(_manifest_files(result))
        assert files, "Expected manifest to contain captured files"
        assert any(str(f.get("path", "")).endswith(".png") for f in files), "Expected PNG screenshot"

        # Verify blobs actually exist in Azure storage
        manifest = result.get("manifest", {})
        storage = manifest.get("storage", {})
        blob_prefix = storage.get("prefix")
        assert blob_prefix, "Manifest missing storage.prefix"
        _verify_blobs_exist(blob_prefix, files)

    def test_get_run_status(self) -> None:
        """Verify get_run returns valid status during polling."""
        _require_e2e()
        base_url = _require_base_url()

        client = ScreenshotServiceClient(base_url=base_url)
        batch_id = f"e2e-batch-client-{uuid.uuid4().hex[:8]}"

        payload = {
            "jobs": [
                {
                    "job_id": "status-test",
                    "url": SITES_UNDER_TEST[1],
                    "options": {"capture": {"enabled": True}},
                }
            ],
            "concurrency": 1,
        }

        handle = client.start_batch(batch_id, payload)
        job_id = str(handle.job_id)

        # Check immediate status
        status = client.get_run(batch_id, job_id)
        assert "status" in status
        assert status["status"] in {"queued", "running", "completed", "failed", "cancelled", "cancel_requested"}

        # Poll to completion
        final_status = client.poll_run(batch_id, job_id, timeout=180.0, poll_interval=5.0)
        assert final_status.get("status") in {"completed", "failed"}

    def test_list_runs(self) -> None:
        """Verify list_runs returns batch runs."""
        _require_e2e()
        base_url = _require_base_url()

        client = ScreenshotServiceClient(base_url=base_url)
        batch_id = f"e2e-batch-client-{uuid.uuid4().hex[:8]}"

        payload = {
            "jobs": [
                {
                    "job_id": "list-test",
                    "url": SITES_UNDER_TEST[2],
                    "options": {"capture": {"enabled": True}},
                }
            ],
            "concurrency": 1,
        }

        handle = client.start_batch(batch_id, payload)
        job_id = str(handle.job_id)

        # List runs for this batch
        runs = client.list_runs(batch_id)
        assert isinstance(runs, list)
        assert len(runs) >= 1

        # Find our run
        our_run = next((r for r in runs if r.get("job_id") == job_id), None)
        assert our_run is not None, f"Expected to find job_id {job_id} in runs list"

        # Clean up - poll to completion
        client.poll_run(batch_id, job_id, timeout=180.0, poll_interval=5.0)

    def test_from_env_factory(self) -> None:
        """Verify from_env() factory creates working client."""
        _require_e2e()
        _require_base_url()  # Ensures SCREENSHOT_FUNC_BASE_URL is set

        client = ScreenshotServiceClient.from_env()
        batch_id = f"e2e-from-env-{uuid.uuid4().hex[:8]}"

        payload = {
            "jobs": [
                {
                    "job_id": "env-test",
                    "url": SITES_UNDER_TEST[0],
                    "options": {"capture": {"enabled": True}},
                }
            ],
            "concurrency": 1,
        }

        handle = client.start_batch(batch_id, payload)
        assert handle.batch_id == batch_id

        job_id = str(handle.job_id)
        status = client.poll_run(batch_id, job_id, timeout=180.0, poll_interval=5.0)
        assert status.get("status") in {"completed", "failed"}

    def test_screenshot_multiple_sites(self) -> None:
        """Run batch with multiple sites and verify all complete."""
        _require_e2e()
        base_url = _require_base_url()

        client = ScreenshotServiceClient(base_url=base_url)
        batch_id = f"e2e-multi-site-{uuid.uuid4().hex[:8]}"

        # Create jobs for all test sites
        jobs = [
            {
                "job_id": f"site-{i}",
                "url": url,
                "options": {
                    "capture": {
                        "enabled": True,
                        "max_pages": 1,
                        "depth": 0,
                        "viewports": ["desktop"],
                        "post_nav_wait_s": 2.0,
                        "timeout_s": 120.0,
                    }
                },
            }
            for i, url in enumerate(SITES_UNDER_TEST)
        ]

        payload = {
            "jobs": jobs,
            "concurrency": 2,
        }

        handle = client.start_batch(batch_id, payload)
        job_id = str(handle.job_id)

        # Poll with longer timeout for multiple sites
        status = client.poll_run(batch_id, job_id, timeout=600.0, poll_interval=10.0)
        assert status.get("status") == "completed", f"Batch failed: {status}"

        # Verify result manifest
        result = client.get_result(batch_id, job_id)
        files = list(_manifest_files(result))
        assert files, "Expected manifest to contain captured files"

        # Should have screenshots for each site
        png_files = [f for f in files if str(f.get("path", "")).endswith(".png")]
        assert len(png_files) >= len(SITES_UNDER_TEST), (
            f"Expected at least {len(SITES_UNDER_TEST)} PNGs, got {len(png_files)}"
        )

        # Verify blobs exist
        manifest = result.get("manifest", {})
        storage = manifest.get("storage", {})
        blob_prefix = storage.get("prefix")
        assert blob_prefix, "Manifest missing storage.prefix"
        _verify_blobs_exist(blob_prefix, files)


# =============================================================================
# Stress Tests
# =============================================================================


def _require_stress_test() -> None:
    """Skip unless RUN_STRESS_TEST=1 is set.

    Stress tests are separate from regular e2e tests because they:
    - Take longer to complete (10+ minutes)
    - Generate higher load on the service
    - May incur additional Azure costs
    """
    if os.getenv("RUN_STRESS_TEST") != "1":
        pytest.skip("Set RUN_STRESS_TEST=1 to enable stress tests")


def _slug_from_url(url: str) -> str:
    """Extract a slug from URL for job naming."""
    from urllib.parse import urlparse

    parsed = urlparse(url)
    return parsed.netloc.replace(".", "-").replace("www-", "")


@dataclass
class BatchRunResult:
    """Result of a single batch run for aggregation."""

    batch_id: str
    job_id: str
    status: str
    job_total: int
    job_completed: int
    job_failed: int
    elapsed: float
    error: str | None = None


def _run_single_batch(
    client: ScreenshotServiceClient,
    batch_id: str,
    jobs_payload: list[dict[str, Any]],
    concurrency: int,
) -> BatchRunResult:
    """Submit and poll a single batch, returning results."""
    start_time = time.time()
    try:
        handle = client.start_batch(
            batch_id,
            {
                "jobs": jobs_payload,
                "concurrency": concurrency,
            },
        )

        status = client.poll_run(
            batch_id,
            str(handle.job_id),
            timeout=300.0,
            poll_interval=10.0,
        )
        elapsed = time.time() - start_time

        return BatchRunResult(
            batch_id=batch_id,
            job_id=str(handle.job_id),
            status=status.get("status", "unknown"),
            job_total=status.get("job_total", len(jobs_payload)),
            job_completed=status.get("job_completed", 0),
            job_failed=status.get("job_failed", 0),
            elapsed=elapsed,
        )
    except Exception as e:
        return BatchRunResult(
            batch_id=batch_id,
            job_id="",
            status="error",
            job_total=len(jobs_payload),
            job_completed=0,
            job_failed=len(jobs_payload),
            elapsed=time.time() - start_time,
            error=str(e),
        )


@pytest.mark.parametrize(
    "concurrency,job_count,description",
    [
        pytest.param(4, 8, "moderate load - 8 jobs @ concurrency 4", id="moderate-4x8"),
        pytest.param(4, 16, "high load - 16 jobs @ concurrency 4", id="high-4x16"),
        pytest.param(4, 24, "burst load - 24 jobs @ concurrency 4", id="burst-4x24"),
    ],
)
def test_stress_single_batch(
    concurrency: int,
    job_count: int,
    description: str,
) -> None:
    """Stress test: single batch with many jobs.

    Tests the service's ability to process many screenshot jobs within
    a single batch via the batch-client. Uses fixed concurrency=4.
    """
    _require_e2e()
    _require_stress_test()
    base_url = _require_base_url()

    # Build job list by cycling through test sites
    jobs_payload = []
    for i in range(job_count):
        url = STRESS_TEST_SITES[i % len(STRESS_TEST_SITES)]
        slug = f"{_slug_from_url(url)}-{i:03d}"
        jobs_payload.append(
            {
                "job_id": slug,
                "url": url,
                "options": {
                    "capture": {
                        "enabled": True,
                        "max_pages": 1,
                        "depth": 0,
                        "viewports": ["desktop"],
                        "pre_capture_wait_s": 0.5,
                        "post_nav_wait_s": 0.1,
                        "timeout_s": 60.0,
                        "max_capture_attempts": 2,
                    }
                },
            }
        )

    client = ScreenshotServiceClient(base_url=base_url)
    batch_id = f"e2e-stress-{concurrency}x{job_count}-{uuid.uuid4().hex[:6]}"

    result = _run_single_batch(client, batch_id, jobs_payload, concurrency)

    # Log results
    print(f"\nStress test results ({description}):")
    print(f"  Batch ID: {result.batch_id}")
    print(f"  Concurrency: {concurrency}")
    print(f"  Jobs: {result.job_completed}/{result.job_total} completed, {result.job_failed} failed")
    print(f"  Duration: {result.elapsed:.1f}s")
    if result.job_completed > 0:
        print(f"  Throughput: {result.job_completed / result.elapsed:.2f} jobs/sec")
    if result.error:
        print(f"  Error: {result.error}")

    # Assertions
    assert result.status == "completed", (
        f"Stress test failed: status={result.status}, "
        f"completed={result.job_completed}/{result.job_total}, failed={result.job_failed}"
    )

    max_allowed_failures = max(1, job_count // 10)
    assert result.job_failed <= max_allowed_failures, (
        f"Too many job failures: {result.job_failed}/{result.job_total} (max allowed: {max_allowed_failures})"
    )


@pytest.mark.parametrize(
    "parallel_batches,jobs_per_batch,concurrency,description",
    [
        pytest.param(2, 4, 4, "2 batches x 4 jobs - baseline", id="2x4"),
        pytest.param(4, 8, 4, "4 batches x 8 jobs - medium load", id="4x8"),
        pytest.param(8, 4, 4, "8 batches x 4 jobs - more parallelism", id="8x4"),
        pytest.param(32, 4, 4, "32 batches x 4 jobs - scale-out test", id="32x4"),
        # pytest.param(96, 4, 4, "96 batches x 4 jobs - heavy scale-out test", id="96x4"),
    ],
)
def test_stress_parallel_batches(
    parallel_batches: int,
    jobs_per_batch: int,
    concurrency: int,
    description: str,
) -> None:
    """Stress test: multiple batches submitted in parallel.

    Tests the service's ability to handle concurrent API requests
    triggering multiple orchestrations simultaneously via the batch-client.
    """
    _require_e2e()
    _require_stress_test()
    base_url = _require_base_url()

    client = ScreenshotServiceClient(base_url=base_url)
    batch_prefix = f"e2e-parallel-{uuid.uuid4().hex[:6]}"

    # Prepare job payloads for each batch
    all_batch_jobs: list[tuple[str, list[dict[str, Any]]]] = []
    for batch_idx in range(parallel_batches):
        batch_id = f"{batch_prefix}-b{batch_idx:02d}"
        jobs_payload = []
        for job_idx in range(jobs_per_batch):
            url = STRESS_TEST_SITES[(batch_idx * jobs_per_batch + job_idx) % len(STRESS_TEST_SITES)]
            slug = f"{_slug_from_url(url)}-b{batch_idx:02d}-j{job_idx:02d}"
            jobs_payload.append(
                {
                    "job_id": slug,
                    "url": url,
                    "options": {
                        "capture": {
                            "enabled": True,
                            "max_pages": 1,
                            "depth": 0,
                            "viewports": ["desktop"],
                            "pre_capture_wait_s": 0.5,
                            "post_nav_wait_s": 0.1,
                            "timeout_s": 60.0,
                            "max_capture_attempts": 2,
                        }
                    },
                }
            )
        all_batch_jobs.append((batch_id, jobs_payload))

    # Submit all batches in parallel using ThreadPoolExecutor
    start_time = time.time()
    results: list[BatchRunResult] = []

    with ThreadPoolExecutor(max_workers=parallel_batches) as executor:
        futures = {
            executor.submit(_run_single_batch, client, batch_id, jobs, concurrency): batch_id
            for batch_id, jobs in all_batch_jobs
        }

        for future in as_completed(futures):
            batch_id = futures[future]
            try:
                result = future.result()
                results.append(result)
            except Exception as e:
                results.append(
                    BatchRunResult(
                        batch_id=batch_id,
                        job_id="",
                        status="error",
                        job_total=jobs_per_batch,
                        job_completed=0,
                        job_failed=jobs_per_batch,
                        elapsed=0,
                        error=str(e),
                    )
                )

    total_elapsed = time.time() - start_time

    # Aggregate results
    total_jobs = sum(r.job_total for r in results)
    total_completed = sum(r.job_completed for r in results)
    total_failed = sum(r.job_failed for r in results)
    batches_completed = sum(1 for r in results if r.status == "completed")
    batches_failed = sum(1 for r in results if r.status != "completed")

    # Log results
    print(f"\nParallel stress test results ({description}):")
    print(f"  Batches: {batches_completed}/{parallel_batches} completed, {batches_failed} failed")
    print(f"  Jobs: {total_completed}/{total_jobs} completed, {total_failed} failed")
    print(f"  Total duration: {total_elapsed:.1f}s")
    if total_completed > 0:
        print(f"  Throughput: {total_completed / total_elapsed:.2f} jobs/sec")

    for r in results:
        status_str = f"{r.status}"
        if r.error:
            status_str += f" ({r.error[:50]}...)" if len(r.error) > 50 else f" ({r.error})"
        print(f"    {r.batch_id}: {r.job_completed}/{r.job_total} - {status_str}")

    # Assertions
    max_batch_failures = max(1, parallel_batches // 4)
    assert batches_failed <= max_batch_failures, (
        f"Too many batch failures: {batches_failed}/{parallel_batches} (max allowed: {max_batch_failures})"
    )

    max_job_failures = max(2, total_jobs // 10)
    assert total_failed <= max_job_failures, (
        f"Too many job failures: {total_failed}/{total_jobs} (max allowed: {max_job_failures})"
    )
