"""Unit tests for the RunPod SDK compatibility layer (gmi_ieops.runpod)."""

import asyncio
import os
import time
import threading
import pytest
import requests
from unittest.mock import patch, MagicMock, AsyncMock


class TestImportCompat:
    """Verify that the import surface matches RunPod SDK."""

    def test_import_runpod_from_gmi_ieops(self):
        from gmi_ieops import runpod
        assert hasattr(runpod, "serverless")

    def test_serverless_start_exists(self):
        from gmi_ieops import runpod
        assert callable(runpod.serverless.start)

    def test_serverless_progress_update_exists(self):
        from gmi_ieops import runpod
        assert callable(runpod.serverless.progress_update)

    def test_rp_upload_import(self):
        from gmi_ieops.runpod.serverless.utils import rp_upload
        assert callable(rp_upload.upload_image)
        assert callable(rp_upload.files)

    def test_rp_logger_import(self):
        from gmi_ieops.runpod.serverless.modules.rp_logger import RunPodLogger
        log = RunPodLogger()
        assert callable(log.info)
        assert callable(log.error)

    def test_worker_state_import(self):
        from gmi_ieops.runpod.serverless.modules.worker_state import WORKER_ID, JobsProgress
        assert isinstance(WORKER_ID, str)
        assert len(WORKER_ID) > 0

    def test_rp_handler_is_generator(self):
        from gmi_ieops.runpod.serverless.modules.rp_handler import is_generator

        def sync_fn():
            return 1

        def gen_fn():
            yield 1

        async def async_fn():
            return 1

        async def async_gen_fn():
            yield 1

        assert is_generator(sync_fn) is False
        assert is_generator(gen_fn) is True
        assert is_generator(async_fn) is False
        assert is_generator(async_gen_fn) is True


class TestHealthServer:
    """Verify the health server starts and responds correctly."""

    def test_health_server_responds_200(self):
        from gmi_ieops.runpod.serverless.worker import _start_health_server

        port = 19876  # unlikely to conflict
        _start_health_server(port=port)
        time.sleep(0.5)

        resp = requests.get(f"http://localhost:{port}/health", timeout=2)
        assert resp.status_code == 200
        assert resp.json() == {"status": "ok"}

    def test_health_server_404_on_other_paths(self):
        # Reuse port from previous test (server still running as daemon thread)
        port = 19876
        resp = requests.get(f"http://localhost:{port}/other", timeout=2)
        assert resp.status_code == 404

    def test_health_server_port_conflict_does_not_crash(self):
        """Starting health server on an occupied port should log warning, not crash."""
        from gmi_ieops.runpod.serverless.worker import _start_health_server
        # Port 19876 is already in use from above — should handle gracefully
        _start_health_server(port=19876)  # should not raise


class TestWorkerLocalMode:
    """Verify that without RUNPOD_WEBHOOK_GET_JOB, worker detects local mode."""

    def test_is_local_without_env(self):
        from gmi_ieops.runpod.serverless.worker import _is_local

        # No RUNPOD_WEBHOOK_GET_JOB → local mode
        env_backup = os.environ.pop("RUNPOD_WEBHOOK_GET_JOB", None)
        try:
            config = {"rp_args": {}}
            assert _is_local(config) is True
        finally:
            if env_backup:
                os.environ["RUNPOD_WEBHOOK_GET_JOB"] = env_backup

    def test_is_local_with_test_input(self):
        from gmi_ieops.runpod.serverless.worker import _is_local

        config = {"rp_args": {"test_input": '{"input": {}}'}}
        assert _is_local(config) is True

    def test_is_not_local_with_webhook(self):
        from gmi_ieops.runpod.serverless.worker import _is_local

        os.environ["RUNPOD_WEBHOOK_GET_JOB"] = "http://fake/v2/test/job-take/$ID"
        try:
            config = {"rp_args": {}}
            assert _is_local(config) is False
        finally:
            del os.environ["RUNPOD_WEBHOOK_GET_JOB"]


class TestJobsProgress:
    """Verify JobsProgress singleton tracks jobs correctly."""

    def test_add_and_get_job_list(self):
        from gmi_ieops.runpod.serverless.modules.worker_state import JobsProgress, Job

        progress = JobsProgress()
        progress.clear()

        progress.add(Job("job-1"))
        progress.add(Job("job-2"))

        job_list = progress.get_job_list()
        assert "job-1" in job_list
        assert "job-2" in job_list

        progress.remove(Job("job-1"))
        job_list = progress.get_job_list()
        assert "job-1" not in job_list
        assert "job-2" in job_list

        progress.clear()
        assert progress.get_job_list() is None

    def test_job_count(self):
        from gmi_ieops.runpod.serverless.modules.worker_state import JobsProgress, Job

        progress = JobsProgress()
        progress.clear()

        assert progress.get_job_count() == 0
        progress.add(Job("j1"))
        assert progress.get_job_count() == 1
        progress.add(Job("j2"))
        assert progress.get_job_count() == 2
        progress.clear()


class TestHandlerSignatureCompat:
    """Verify the handler function signature matches RunPod SDK expectations."""

    def test_sync_handler_returns_dict(self):
        """ComfyUI-style handler: receives job dict, returns result dict."""
        from gmi_ieops.runpod.serverless.modules.rp_job import run_job
        import asyncio

        def handler(job):
            return {"text": "hello", "worker_id": "test"}

        job = {"id": "test-job-1", "input": {"prompt": "hi"}}
        result = asyncio.run(run_job(handler, job))

        assert result["output"]["text"] == "hello"
        assert result["output"]["worker_id"] == "test"

    def test_sync_handler_returns_error(self):
        from gmi_ieops.runpod.serverless.modules.rp_job import run_job
        import asyncio

        def handler(job):
            return {"error": "something went wrong"}

        job = {"id": "test-job-2", "input": {}}
        result = asyncio.run(run_job(handler, job))

        assert "error" in result

    def test_sync_handler_exception(self):
        from gmi_ieops.runpod.serverless.modules.rp_job import run_job

        def handler(job):
            raise ValueError("test error")

        job = {"id": "test-job-3", "input": {}}
        result = asyncio.run(run_job(handler, job))

        assert "error" in result

    def test_generator_handler(self):
        """RunPod SDK returns generator output as-is (consumed by rp_http.stream_result)."""
        from gmi_ieops.runpod.serverless.modules.rp_job import run_job
        import inspect

        def handler(job):
            yield {"token": "hello"}
            yield {"token": " world"}

        job = {"id": "test-job-gen", "input": {}}
        result = asyncio.run(run_job(handler, job))

        # run_job wraps generator in output — streaming is handled by rp_http
        assert inspect.isgenerator(result["output"])

    def test_async_handler(self):
        from gmi_ieops.runpod.serverless.modules.rp_job import run_job

        async def handler(job):
            return {"result": "async-ok"}

        job = {"id": "test-job-async", "input": {}}
        result = asyncio.run(run_job(handler, job))

        assert result["output"]["result"] == "async-ok"


class TestUploadCompat:
    """Verify rp_upload API surface matches RunPod SDK."""

    def test_upload_image_without_bucket_returns_local(self):
        """Without BUCKET_ENDPOINT_URL, upload_image should not crash."""
        from gmi_ieops.runpod.serverless.utils import rp_upload
        import tempfile

        # Create a temp file to upload
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
            f.write(b"fake image data")
            temp_path = f.name

        try:
            # Without BUCKET_ENDPOINT_URL, should return local path
            env_backup = os.environ.pop("BUCKET_ENDPOINT_URL", None)
            try:
                result = rp_upload.upload_image("test-job", temp_path)
                # Should return a string (local path or presigned URL)
                assert isinstance(result, str)
            finally:
                if env_backup:
                    os.environ["BUCKET_ENDPOINT_URL"] = env_backup
        finally:
            os.remove(temp_path)

    def test_get_boto_client_without_creds(self):
        """get_boto_client should handle missing credentials gracefully."""
        from gmi_ieops.runpod.serverless.utils.rp_upload import get_boto_client

        env_backup = {}
        for key in ["BUCKET_ENDPOINT_URL", "BUCKET_ACCESS_KEY_ID", "BUCKET_SECRET_ACCESS_KEY"]:
            env_backup[key] = os.environ.pop(key, None)

        try:
            result = get_boto_client()
            # Should return (None, None) or handle gracefully
            assert result is not None  # returns a tuple
        finally:
            for key, val in env_backup.items():
                if val is not None:
                    os.environ[key] = val


class TestHeartbeatCompat:
    """Verify heartbeat (rp_ping) API surface."""

    def test_heartbeat_class_exists(self):
        from gmi_ieops.runpod.serverless.modules.rp_ping import Heartbeat
        hb = Heartbeat()
        assert hasattr(hb, "start_ping")
        assert hasattr(hb, "PING_URL")
        assert hasattr(hb, "PING_INTERVAL")

    def test_heartbeat_reads_env(self):
        """Heartbeat should read RUNPOD_WEBHOOK_PING and RUNPOD_PING_INTERVAL."""
        os.environ["RUNPOD_WEBHOOK_PING"] = "http://test-platform/v2/ep/ping/$RUNPOD_POD_ID"
        os.environ["RUNPOD_PING_INTERVAL"] = "5000"
        os.environ["RUNPOD_POD_ID"] = "test-pod-123"

        try:
            from importlib import reload
            import gmi_ieops.runpod.serverless.modules.rp_ping as ping_mod
            reload(ping_mod)

            hb = ping_mod.Heartbeat()
            assert "test-platform" in hb.PING_URL
            assert hb.PING_INTERVAL == 5
        finally:
            os.environ.pop("RUNPOD_WEBHOOK_PING", None)
            os.environ.pop("RUNPOD_PING_INTERVAL", None)
            os.environ.pop("RUNPOD_POD_ID", None)


class TestRefreshWorker:
    """Verify REFRESH_WORKER environment variable support."""

    def test_refresh_worker_env_default_false(self):
        """REFRESH_WORKER defaults to false."""
        from gmi_ieops.runpod.serverless import modules as _m
        # The serverless __init__.py reads REFRESH_WORKER
        from gmi_ieops.runpod.serverless import worker
        # Default should not crash; worker module should load fine
        assert worker is not None

    def test_refresh_worker_env_true(self):
        os.environ["REFRESH_WORKER"] = "true"
        try:
            from importlib import reload
            import gmi_ieops.runpod.serverless as sl
            reload(sl)
            # Should not crash
            assert sl is not None
        finally:
            os.environ.pop("REFRESH_WORKER", None)


class TestOutputUpload:
    """Verify SDK output upload interception (_upload_output)."""

    def test_upload_skipped_when_not_configured(self):
        """Without IEOPS_WEBDAV_URL, _upload_output passes through original data."""
        from gmi_ieops.runpod.serverless.modules import rp_http

        original_url = rp_http._WEBDAV_URL
        rp_http._WEBDAV_URL = ""
        try:
            job_data = {"output": {"text": "hello"}}
            result = asyncio.run(rp_http._upload_output("job-1", job_data))
            assert result is job_data  # unchanged
        finally:
            rp_http._WEBDAV_URL = original_url

    def test_upload_skipped_when_no_owner_oid(self):
        """Without IEOPS_OWNER_OID, _upload_output passes through."""
        from gmi_ieops.runpod.serverless.modules import rp_http

        original_url, original_oid = rp_http._WEBDAV_URL, rp_http._OWNER_OID
        rp_http._WEBDAV_URL = "http://fake:9007"
        rp_http._OWNER_OID = ""
        try:
            job_data = {"output": {"text": "hello"}}
            result = asyncio.run(rp_http._upload_output("job-2", job_data))
            assert result is job_data
        finally:
            rp_http._WEBDAV_URL = original_url
            rp_http._OWNER_OID = original_oid

    def test_upload_skipped_for_in_progress(self):
        """IN_PROGRESS status should not be uploaded."""
        from gmi_ieops.runpod.serverless.modules import rp_http

        original_url, original_oid = rp_http._WEBDAV_URL, rp_http._OWNER_OID
        rp_http._WEBDAV_URL = "http://fake:9007"
        rp_http._OWNER_OID = "org-1"
        try:
            job_data = {"status": "IN_PROGRESS", "output": {"partial": True}}
            result = asyncio.run(rp_http._upload_output("job-3", job_data))
            assert result is job_data
        finally:
            rp_http._WEBDAV_URL = original_url
            rp_http._OWNER_OID = original_oid

    def test_upload_skipped_when_no_output(self):
        """Job data without 'output' key should not be uploaded."""
        from gmi_ieops.runpod.serverless.modules import rp_http

        original_url, original_oid = rp_http._WEBDAV_URL, rp_http._OWNER_OID
        rp_http._WEBDAV_URL = "http://fake:9007"
        rp_http._OWNER_OID = "org-1"
        try:
            job_data = {"error": "something failed"}
            result = asyncio.run(rp_http._upload_output("job-4", job_data))
            assert result is job_data
        finally:
            rp_http._WEBDAV_URL = original_url
            rp_http._OWNER_OID = original_oid

    def test_upload_graceful_degradation_on_connection_error(self):
        """When WebDAV is unreachable, returns original result (no crash)."""
        from gmi_ieops.runpod.serverless.modules import rp_http

        original_url, original_oid, original_uid = rp_http._WEBDAV_URL, rp_http._OWNER_OID, rp_http._OWNER_UID
        rp_http._WEBDAV_URL = "http://127.0.0.1:19999"  # nothing listening
        rp_http._OWNER_OID = "org-1"
        rp_http._OWNER_UID = "user-1"
        try:
            job_data = {"output": {"text": "hello"}}
            result = asyncio.run(rp_http._upload_output("job-5", job_data))
            assert result is job_data  # degraded to original
        finally:
            rp_http._WEBDAV_URL = original_url
            rp_http._OWNER_OID = original_oid
            rp_http._OWNER_UID = original_uid

    def test_upload_returns_file_url_on_success(self):
        """Mock WebDAV to verify successful upload returns _file_url."""
        from gmi_ieops.runpod.serverless.modules import rp_http
        from unittest.mock import AsyncMock, patch, MagicMock
        from contextlib import asynccontextmanager

        original_url, original_oid, original_uid = rp_http._WEBDAV_URL, rp_http._OWNER_OID, rp_http._OWNER_UID
        rp_http._WEBDAV_URL = "http://fake-webdav:9007"
        rp_http._OWNER_OID = "test-org"
        rp_http._OWNER_UID = "test-user"

        try:
            # Mock aiohttp.ClientSession to return 201 for MKCOL and PUT
            mock_resp = MagicMock()
            mock_resp.status = 201

            @asynccontextmanager
            async def mock_request(*args, **kwargs):
                yield mock_resp

            @asynccontextmanager
            async def mock_put(*args, **kwargs):
                yield mock_resp

            @asynccontextmanager
            async def mock_session_ctx(*args, **kwargs):
                session = MagicMock()
                session.request = mock_request
                session.put = mock_put
                yield session

            with patch("gmi_ieops.runpod.serverless.modules.rp_http.aiohttp.ClientSession", return_value=MagicMock(__aenter__=mock_session_ctx().__aenter__, __aexit__=mock_session_ctx().__aexit__)):
                # Simpler approach: patch at module level
                pass

            # Use a simpler mock approach
            import aiohttp as _aiohttp

            class FakeResp:
                status = 201
                async def __aenter__(self): return self
                async def __aexit__(self, *a): pass

            class FakeSession:
                async def __aenter__(self): return self
                async def __aexit__(self, *a): pass
                def request(self, *a, **kw): return FakeResp()
                def put(self, *a, **kw): return FakeResp()

            with patch.object(rp_http.aiohttp, "ClientSession", return_value=FakeSession()):
                job_data = {"output": {"text": "hello", "score": 0.99}}
                result = asyncio.run(rp_http._upload_output("job-mock", job_data))

            assert result["output"]["_file_uploaded"] is True
            assert result["output"]["_file_url"] == "/api/v1/files/test-org/test-user/outputs/output-job-mock.json"
            assert result["output"]["_original_size"] > 0
        finally:
            rp_http._WEBDAV_URL = original_url
            rp_http._OWNER_OID = original_oid
            rp_http._OWNER_UID = original_uid


class TestSubmoduleImports:
    """Verify all submodules that ComfyUI worker might transitively import."""

    def test_rp_download(self):
        from gmi_ieops.runpod.serverless.utils import rp_download
        assert callable(rp_download.download_files_from_urls)

    def test_rp_cuda(self):
        from gmi_ieops.runpod.serverless.utils import rp_cuda
        assert hasattr(rp_cuda, "is_available")

    def test_rp_validator(self):
        from gmi_ieops.runpod.serverless.utils import rp_validator
        assert hasattr(rp_validator, "validate")

    def test_rp_debugger(self):
        from gmi_ieops.runpod.serverless.utils import rp_debugger
        assert hasattr(rp_debugger, "Checkpoints")

    def test_rp_cleanup(self):
        from gmi_ieops.runpod.serverless.utils import rp_cleanup
        assert callable(rp_cleanup.clean)

    def test_http_client(self):
        from gmi_ieops.runpod.http_client import AsyncClientSession, SyncClientSession, TooManyRequests
        assert callable(AsyncClientSession)
        assert SyncClientSession is not None
        assert TooManyRequests is not None

    def test_error_module(self):
        from gmi_ieops.runpod.error import RunPodError
        assert issubclass(RunPodError, Exception)
