"""Unit tests for execution core batch runtime."""

import asyncio
import json
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock

import pytest

from dotpromptz.adapters._base import Adapter, GeneratedImage, GenerateResponse
from dotpromptz.executor import (
    BatchResult,
    _compute_delay,
    _execute_with_retry,
    _extract_fixed_output_content,
    execute_batch,
    execute_plan,
    extract_output_content,
    part_to_output_dict,
    resolve_batch_adapter,
    serialize_output,
    write_output_file,
)
from dotpromptz.planner import ExecutionPlan, ExecutionPlanItem
from dotpromptz.typing import (
    AdapterConfig,
    MediaContent,
    MediaPart,
    PromptOutputConfig,
    RenderedPrompt,
    RetryConfig,
    Role,
    RuntimeConfig,
    TextPart,
)


def _make_adapter(*, text: str = 'response', side_effect: Exception | None = None) -> Adapter:
    """Create a mock adapter with configurable async generate behavior."""
    adapter = MagicMock(spec=Adapter)
    if side_effect:
        adapter.generate = AsyncMock(side_effect=side_effect)
    else:
        adapter.generate = AsyncMock(return_value=GenerateResponse(text=text))
    return adapter


def _make_rendered_prompt(
    *,
    adapter: str | AdapterConfig = 'openai',
    model: str = 'gpt-4o',
    output: PromptOutputConfig | None = None,
) -> RenderedPrompt:
    """Build a minimal rendered prompt for batch execution tests."""
    payload: dict[str, Any] = {
        'version': 1,
        'adapter': adapter,
        'config': {'model': model},
        'messages': [
            {
                'role': Role.USER,
                'content': [
                    TextPart(text='hello'),
                ],
            }
        ],
    }
    if output is not None:
        payload['output'] = output
    return RenderedPrompt.model_validate(payload)


def _make_execution_plan(
    *,
    tmp_path: Path,
    output_config: PromptOutputConfig,
    rendered: RenderedPrompt | None = None,
    stem: str = 'result',
) -> ExecutionPlan:
    """Build a minimal execution plan for execute_plan tests."""
    rendered_prompt = rendered or _make_rendered_prompt(output=output_config)
    ext = 'png' if output_config.format == 'image' else output_config.format
    output_path = tmp_path / f'{stem}.{ext}'
    return ExecutionPlan(
        compiled=MagicMock(),
        output_config=output_config,
        runtime=None,
        output_dir=tmp_path,
        fmt=output_config.format,
        batch_time='20260305_000000',
        force=False,
        items=[
            ExecutionPlanItem(
                index=0,
                variables={},
                rendered=rendered_prompt,
                output_stem=stem,
                output_path=output_path,
            )
        ],
    )


def _find_log_record(caplog: pytest.LogCaptureFixture, message: str) -> Any:
    """Return the first captured record matching the message text."""
    for record in caplog.records:
        if record.getMessage() == message:
            return record
    raise AssertionError(f'Log record not found: {message}')


class TestExecuteBatch:
    """Tests for batch execution behavior."""

    @pytest.mark.asyncio
    async def test_basic_batch(self) -> None:
        adapter = _make_adapter(text='result')
        rendered = [_make_rendered_prompt(), _make_rendered_prompt()]
        inputs = [{'topic': 'AI'}, {'topic': 'ML'}]

        results = await execute_batch(rendered, inputs, adapter)

        assert len(results) == 2
        assert all(item.status == 'success' for item in results)
        assert [item.index for item in results] == [0, 1]
        assert [item.trace for item in results] == [(0,), (1,)]

    @pytest.mark.asyncio
    async def test_length_mismatch_raises(self) -> None:
        adapter = _make_adapter(text='result')
        rendered = [_make_rendered_prompt()]
        inputs = [{'topic': 'AI'}, {'topic': 'ML'}]

        with pytest.raises(ValueError, match='same length'):
            await execute_batch(rendered, inputs, adapter)

    @pytest.mark.asyncio
    async def test_max_workers_zero_raises(self) -> None:
        adapter = _make_adapter(text='unused')
        rendered = [_make_rendered_prompt()]

        with pytest.raises(ValueError, match='max_workers'):
            await execute_batch(rendered, [{}], adapter, max_workers=0)

    @pytest.mark.asyncio
    async def test_retry_transient_error(self) -> None:
        call_count = 0

        async def _generate(_: Any) -> GenerateResponse:
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise RuntimeError('transient')
            return GenerateResponse(text='ok')

        adapter = MagicMock(spec=Adapter)
        adapter.generate = AsyncMock(side_effect=_generate)

        runtime = RuntimeConfig(retry=RetryConfig(max_retries=2, retry_delay=0.001))
        results = await execute_batch([_make_rendered_prompt()], [{}], adapter, runtime=runtime)

        assert len(results) == 1
        assert results[0].status == 'success'
        assert results[0].attempts == 2

    @pytest.mark.asyncio
    async def test_schema_validation_failure_returns_error(self) -> None:
        adapter = _make_adapter(text='{"age": 30}')
        output = PromptOutputConfig(
            format='json',
            file_name='result',
            schema_value={
                'type': 'object',
                'properties': {'name': {'type': 'string'}},
                'required': ['name'],
                'additionalProperties': False,
            },
        )
        rendered = [_make_rendered_prompt(output=output)]

        results = await execute_batch(rendered, [{}], adapter)

        assert len(results) == 1
        assert results[0].status == 'error'
        assert results[0].error is not None
        assert 'Schema validation failed' in results[0].error

    @pytest.mark.asyncio
    async def test_on_item_complete_callback(self) -> None:
        adapter = _make_adapter(text='done')
        completed: list[BatchResult] = []

        def _on_complete(result: BatchResult) -> None:
            completed.append(result)

        rendered = [_make_rendered_prompt(), _make_rendered_prompt()]
        inputs = [{'a': 1}, {'a': 2}]
        results = await execute_batch(rendered, inputs, adapter, on_item_complete=_on_complete)

        assert len(completed) == 2
        assert len(results) == 2

    @pytest.mark.asyncio
    async def test_batch_progress_metadata_tracks_input_and_completion_order(self) -> None:
        slow_rendered = _make_rendered_prompt().model_copy(
            update={'messages': [{'role': Role.USER, 'content': [TextPart(text='slow')]}]}
        )
        fast_rendered = _make_rendered_prompt().model_copy(
            update={'messages': [{'role': Role.USER, 'content': [TextPart(text='fast')]}]}
        )

        async def _generate(rendered: RenderedPrompt) -> GenerateResponse:
            if rendered == slow_rendered:
                text = 'slow'
            else:
                text = 'fast'
            if text == 'slow':
                await asyncio.sleep(0.02)
            else:
                await asyncio.sleep(0.001)
            return GenerateResponse(text=text)

        adapter = MagicMock(spec=Adapter)
        adapter.generate = AsyncMock(side_effect=_generate)
        rendered = [slow_rendered, fast_rendered]

        results = await execute_batch(rendered, [{}, {}], adapter, max_workers=2)

        assert [item.item_no for item in results] == [1, 2]
        assert [item.total_items for item in results] == [2, 2]
        assert [item.completed_items for item in results] == [2, 1]

    @pytest.mark.asyncio
    async def test_retry_log_includes_progress_context(self, caplog: pytest.LogCaptureFixture) -> None:
        call_count = 0

        async def _generate(_: Any) -> GenerateResponse:
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise RuntimeError('transient')
            return GenerateResponse(text='ok')

        adapter = MagicMock(spec=Adapter)
        adapter.generate = AsyncMock(side_effect=_generate)
        runtime = RuntimeConfig(retry=RetryConfig(max_retries=1, retry_delay=0.001))
        caplog.set_level('WARNING')

        await execute_batch([_make_rendered_prompt()], [{}], adapter, runtime=runtime)

        assert 'Retrying after error.' in caplog.text
        record = _find_log_record(caplog, 'Retrying after error.')
        assert record._dotpromptz_fields['item_no'] == 1
        assert record._dotpromptz_fields['total_items'] == 1

    @pytest.mark.asyncio
    async def test_batch_returns_success_and_error_results(
        self,
    ) -> None:
        adapter = MagicMock(spec=Adapter)
        adapter.generate = AsyncMock(side_effect=[GenerateResponse(text='ok'), RuntimeError('boom')])
        rendered = [_make_rendered_prompt(), _make_rendered_prompt()]

        results = await execute_batch(rendered, [{}, {}], adapter, max_workers=2)

        assert [item.status for item in results] == ['success', 'error']


class TestResolveBatchAdapter:
    """Tests for adapter resolution from rendered prompt metadata."""

    def test_adapter_config_object(self) -> None:
        from unittest.mock import patch

        from dotpromptz.credentials import Credential, CredentialPool

        cred = Credential(name='test', adapter='openai', api_key='sk-test')
        pool = CredentialPool([cred])
        rendered = _make_rendered_prompt(adapter=AdapterConfig(name='openai'))

        with patch('dotpromptz.executor.get_adapter') as mock_get_adapter:
            resolve_batch_adapter(rendered, pool=pool)
            mock_get_adapter.assert_called_once_with('openai', credential=cred)

    def test_string_adapter(self) -> None:
        from unittest.mock import patch

        from dotpromptz.credentials import Credential, CredentialPool

        cred = Credential(name='main', adapter='anthropic', api_key='sk-test')
        pool = CredentialPool([cred])
        rendered = _make_rendered_prompt(adapter='anthropic')

        with patch('dotpromptz.executor.get_adapter') as mock_get_adapter:
            resolve_batch_adapter(rendered, pool=pool)
            mock_get_adapter.assert_called_once_with('anthropic', credential=cred)

    def test_none_adapter_raises_value_error(self) -> None:
        rendered = _make_rendered_prompt(adapter='openai').model_copy(update={'adapter': None})

        with pytest.raises(ValueError, match='No adapter specified'):
            resolve_batch_adapter(rendered)


class TestSerializeOutput:
    """Tests for output serialization helpers."""

    def test_json_dict(self) -> None:
        result = serialize_output({'key': 'value'}, 'json')
        assert json.loads(result) == {'key': 'value'}

    def test_yaml_format(self) -> None:
        result = serialize_output({'name': 'Alice', 'age': 30}, 'yaml')
        assert 'name: Alice' in result
        assert 'age: 30' in result

    def test_txt_format(self) -> None:
        result = serialize_output({'key': 'value'}, 'txt')
        assert result == str({'key': 'value'})


class TestPartToOutputDict:
    """Tests for part serialization."""

    def test_text_part(self) -> None:
        result = part_to_output_dict(TextPart(text='hello'))
        assert result == {'text': 'hello'}

    def test_media_part(self) -> None:
        part = MediaPart(media=MediaContent(url='https://example.com/img.png', contentType='image/png'))
        result = part_to_output_dict(part)
        assert result['media']['url'] == 'https://example.com/img.png'


class TestWriteOutputFile:
    """Tests for output file writing helper."""

    def test_writes_text_file(self, tmp_path: Path) -> None:
        cfg = PromptOutputConfig(format='json', file_name='output')
        result = write_output_file('{"a": 1}', cfg, tmp_path, 'result')
        assert result == tmp_path / 'result.json'
        assert result.read_text(encoding='utf-8') == '{"a": 1}'

    def test_existing_file_without_force_raises(self, tmp_path: Path) -> None:
        cfg = PromptOutputConfig(format='txt', file_name='output')
        filepath = tmp_path / 'out.txt'
        filepath.write_text('old', encoding='utf-8')
        with pytest.raises(FileExistsError, match='--force'):
            write_output_file('new', cfg, tmp_path, 'out')

    def test_writes_image_output_as_png(self, tmp_path: Path) -> None:
        cfg = PromptOutputConfig(format='image', file_name='output')
        result = write_output_file(b'\x89PNG', cfg, tmp_path, 'result')
        assert result == tmp_path / 'result.png'
        assert result.read_bytes() == b'\x89PNG'


class TestExtractOutputContent:
    """Tests for extracting serialized text from responses."""

    def test_text_json(self) -> None:
        response = GenerateResponse(text='{"key": "value"}')
        result = extract_output_content(response, 'json')
        assert json.loads(result) == {'key': 'value'}

    def test_empty_response(self) -> None:
        response = GenerateResponse()
        result = extract_output_content(response, 'json')
        assert result == ''

    def test_image_response(self) -> None:
        response = GenerateResponse(images=[GeneratedImage(data=b'\x89PNG', mime_type='image/png')])
        result = extract_output_content(response, 'image')
        assert result == b'\x89PNG'

    def test_image_response_requires_png(self) -> None:
        response = GenerateResponse(images=[GeneratedImage(data=b'\xff\xd8', mime_type='image/jpeg')])
        with pytest.raises(ValueError, match='image/png'):
            extract_output_content(response, 'image')

    def test_image_response_with_anthropic_reports_unsupported(self) -> None:
        response = GenerateResponse(text='no image')
        with pytest.raises(ValueError, match='Anthropic adapter does not support output.format=image'):
            extract_output_content(response, 'image', adapter_name='anthropic')


class TestExtractFixedOutputContent:
    """Tests for extracting configured fixed output content."""

    def test_json_fixed_value(self) -> None:
        output = PromptOutputConfig(
            format='json',
            file_name='result',
            schema_value={
                'type': 'object',
                'properties': {'name': {'type': 'string'}},
                'required': ['name'],
            },
            fixed={'name': 'Alice'},
        )

        result = _extract_fixed_output_content(output, 'json')

        assert result is not None
        assert json.loads(result) == {'name': 'Alice'}

    def test_invalid_fixed_value_raises(self) -> None:
        output = PromptOutputConfig(
            format='json',
            file_name='result',
            schema_value={
                'type': 'object',
                'properties': {'name': {'type': 'string'}},
                'required': ['name'],
            },
            fixed={'age': 30},
        )

        with pytest.raises(ValueError, match='Schema validation failed'):
            _extract_fixed_output_content(output, 'json')


class TestComputeDelay:
    """Tests for retry delay calculation."""

    def test_linear_delay(self) -> None:
        cfg = RetryConfig(max_retries=3, retry_delay=2.0)
        assert _compute_delay(0, cfg) == 2.0
        assert _compute_delay(1, cfg) == 4.0


class TestExecuteWithRetry:
    """Tests for low-level retry executor."""

    @pytest.mark.asyncio
    async def test_success_no_retry(self) -> None:
        async def _ok(attempt: int) -> str:
            assert attempt == 0
            return 'ok'

        result, attempts = await _execute_with_retry(_ok, retry_cfg=None, timeout=None)
        assert result == 'ok'
        assert attempts == 1

    @pytest.mark.asyncio
    async def test_timeout_all_attempts_raises(self) -> None:
        async def _slow(attempt: int) -> str:
            await asyncio.sleep(10)
            return 'never'

        cfg = RetryConfig(max_retries=1, retry_delay=0.001)
        with pytest.raises(asyncio.TimeoutError):
            await _execute_with_retry(_slow, retry_cfg=cfg, timeout=0.01)


class TestExecutePlan:
    """Tests for plan-level execution and output handling."""

    @pytest.mark.asyncio
    async def test_item_failure_logs_without_writing_error_output(
        self,
        tmp_path: Path,
        monkeypatch: pytest.MonkeyPatch,
        caplog: pytest.LogCaptureFixture,
    ) -> None:
        output_config = PromptOutputConfig(format='json', file_name='result')
        plan = _make_execution_plan(tmp_path=tmp_path, output_config=output_config)

        monkeypatch.setattr('dotpromptz.executor.resolve_batch_adapter', lambda rendered: _make_adapter(text='ok'))

        async def _fake_execute_batch(
            rendered_prompts: list[RenderedPrompt],
            input_list: list[dict[str, Any]],
            adapter: Adapter,
            *,
            prompt_file: Path | None = None,
            trace_prefix: tuple[int, ...] = (),
            max_workers: int = 5,
            on_item_complete: Any = None,
            runtime: RuntimeConfig | None = None,
            abort_event: asyncio.Event | None = None,
        ) -> list[BatchResult]:
            result = BatchResult(
                index=0,
                input={},
                prompt_file=prompt_file,
                trace=trace_prefix + (0,),
                status='error',
                error='boom',
                attempts=2,
                elapsed=0.01,
            )
            if on_item_complete is not None:
                on_item_complete(result)
            return [result]

        monkeypatch.setattr('dotpromptz.executor.execute_batch', _fake_execute_batch)
        caplog.set_level('ERROR')

        report = await execute_plan(plan)

        assert report.batch_results[0].status == 'error'
        assert report.written_files == []
        assert not (tmp_path / 'result.json').exists()
        assert 'Batch item execution failed.' in caplog.text
        record = _find_log_record(caplog, 'Batch item execution failed.')
        assert record._dotpromptz_fields['item_no'] == 1
        assert record._dotpromptz_fields['total_items'] == 1
        assert record._dotpromptz_fields['completed_items'] == 1

    @pytest.mark.asyncio
    async def test_success_logs_input_and_completion_progress(
        self,
        tmp_path: Path,
        monkeypatch: pytest.MonkeyPatch,
        caplog: pytest.LogCaptureFixture,
    ) -> None:
        output_config = PromptOutputConfig(format='json', file_name='result')
        plan = _make_execution_plan(tmp_path=tmp_path, output_config=output_config)

        monkeypatch.setattr('dotpromptz.executor.resolve_batch_adapter', lambda rendered: _make_adapter(text='ok'))
        caplog.set_level('INFO')

        report = await execute_plan(plan)

        assert report.batch_results[0].status == 'success'
        assert (tmp_path / 'result.json').exists()
        assert 'Batch item completed.' in caplog.text
        completed_record = _find_log_record(caplog, 'Batch item completed.')
        assert completed_record._dotpromptz_fields['item_no'] == 1
        assert completed_record._dotpromptz_fields['total_items'] == 1
        assert completed_record._dotpromptz_fields['completed_items'] == 1

    @pytest.mark.asyncio
    async def test_image_extract_failure_logs_without_writing_error_file(
        self,
        tmp_path: Path,
        monkeypatch: pytest.MonkeyPatch,
        caplog: pytest.LogCaptureFixture,
    ) -> None:
        output_config = PromptOutputConfig(format='image', file_name='result')
        plan = _make_execution_plan(tmp_path=tmp_path, output_config=output_config)

        monkeypatch.setattr('dotpromptz.executor.resolve_batch_adapter', lambda rendered: _make_adapter(text='ok'))

        async def _fake_execute_batch(
            rendered_prompts: list[RenderedPrompt],
            input_list: list[dict[str, Any]],
            adapter: Adapter,
            *,
            prompt_file: Path | None = None,
            trace_prefix: tuple[int, ...] = (),
            max_workers: int = 5,
            on_item_complete: Any = None,
            runtime: RuntimeConfig | None = None,
            abort_event: asyncio.Event | None = None,
        ) -> list[BatchResult]:
            result = BatchResult(
                index=0,
                input={},
                prompt_file=prompt_file,
                trace=trace_prefix + (0,),
                status='success',
                response=GenerateResponse(text='no image'),
                error=None,
            )
            if on_item_complete is not None:
                on_item_complete(result)
            return [result]

        monkeypatch.setattr('dotpromptz.executor.execute_batch', _fake_execute_batch)
        caplog.set_level('WARNING')

        report = await execute_plan(plan)

        assert report.batch_results[0].status == 'success'
        assert report.written_files == []
        assert not (tmp_path / 'result.png').exists()
        assert not (tmp_path / 'result.error.txt').exists()
        assert 'Failed to extract image response content.' in caplog.text

    @pytest.mark.asyncio
    async def test_fixed_output_writes_without_calling_adapter(
        self,
        tmp_path: Path,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        output_config = PromptOutputConfig(
            format='json',
            file_name='result',
            schema_value={
                'type': 'object',
                'properties': {'name': {'type': 'string'}},
                'required': ['name'],
            },
            fixed={'name': 'Alice'},
        )
        rendered = _make_rendered_prompt(output=output_config).model_copy(update={'adapter': None})
        plan = _make_execution_plan(tmp_path=tmp_path, output_config=output_config, rendered=rendered)

        def _unexpected_adapter(_: RenderedPrompt) -> Adapter:
            raise AssertionError('resolve_batch_adapter should not be called for fixed output')

        monkeypatch.setattr('dotpromptz.executor.resolve_batch_adapter', _unexpected_adapter)

        report = await execute_plan(plan)

        assert report.batch_results[0].status == 'success'
        assert report.batch_results[0].attempts == 0
        assert report.written_files == [tmp_path / 'result.json']
        assert json.loads((tmp_path / 'result.json').read_text(encoding='utf-8')) == {'name': 'Alice'}

    @pytest.mark.asyncio
    async def test_fixed_output_schema_validation_failure_returns_error(
        self,
        tmp_path: Path,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        output_config = PromptOutputConfig(
            format='json',
            file_name='result',
            schema_value={
                'type': 'object',
                'properties': {'name': {'type': 'string'}},
                'required': ['name'],
            },
            fixed={'age': 30},
        )
        rendered = _make_rendered_prompt(output=output_config).model_copy(update={'adapter': None})
        plan = _make_execution_plan(tmp_path=tmp_path, output_config=output_config, rendered=rendered)

        def _unexpected_adapter(_: RenderedPrompt) -> Adapter:
            raise AssertionError('resolve_batch_adapter should not be called for fixed output')

        monkeypatch.setattr('dotpromptz.executor.resolve_batch_adapter', _unexpected_adapter)

        report = await execute_plan(plan)

        assert report.batch_results[0].status == 'error'
        assert report.batch_results[0].error is not None
        assert 'Schema validation failed' in report.batch_results[0].error
        assert report.written_files == []
        assert not (tmp_path / 'result.json').exists()
