from __future__ import annotations

import uuid

import pytest

from codeflash_api.jit._router import (
    SYSTEM_PROMPT,
    USER_PROMPT,
    _validate_request,
)
from codeflash_api.jit.schemas import JitRewriteRequest
from codeflash_api.optimize._context import BaseOptimizerContext
from codeflash_api.optimize.schemas import (
    OptimizeErrorResponse,
    OptimizeResponse,
    OptimizeResponseItem,
)


class TestJitRewriteRequest:
    """Tests for JitRewriteRequest schema."""

    def test_defaults(self) -> None:
        """
        Default fields are set correctly.
        """
        req = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
        )
        assert 1 == req.n_candidates
        assert "python" == req.language
        assert req.dependency_code is None
        assert req.python_version is None

    def test_version_resolution(self) -> None:
        """
        language_version backfills python_version.
        """
        req = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
            language_version="3.12.9",
        )
        assert "3.12.9" == req.python_version

    def test_version_resolution_skipped_for_non_python(
        self,
    ) -> None:
        """
        language_version does not backfill for non-python.
        """
        req = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
            language_version="3.12.9",
            language="java",
        )
        assert req.python_version is None

    def test_explicit_python_version_not_overridden(
        self,
    ) -> None:
        """
        Explicit python_version is not overridden by language_version.
        """
        req = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
            python_version="3.11.0",
            language_version="3.12.9",
        )
        assert "3.11.0" == req.python_version

    def test_no_optimize_specific_fields(self) -> None:
        """
        JitRewriteRequest lacks optimize-specific fields.
        """
        req = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
        )
        assert not hasattr(req, "baseline_runtime_ns")
        assert not hasattr(req, "loop_count")
        assert not hasattr(req, "line_profiler_results")
        assert not hasattr(req, "test_input_examples")
        assert not hasattr(req, "is_async")
        assert not hasattr(req, "is_numerical_code")


class TestPromptTemplates:
    """Tests for JIT prompt template loading."""

    def test_system_prompt_loaded(self) -> None:
        """
        System prompt is loaded and non-empty.
        """
        assert len(SYSTEM_PROMPT) > 0
        assert "JIT" in SYSTEM_PROMPT

    def test_user_prompt_loaded(self) -> None:
        """
        User prompt is loaded and contains placeholder.
        """
        assert len(USER_PROMPT) > 0
        assert "{source_code}" in USER_PROMPT

    def test_system_prompt_contains_numba(self) -> None:
        """
        System prompt covers Numba.
        """
        assert "Numba" in SYSTEM_PROMPT
        assert "@njit" in SYSTEM_PROMPT

    def test_system_prompt_contains_torch(self) -> None:
        """
        System prompt covers PyTorch.
        """
        assert "torch.compile" in SYSTEM_PROMPT

    def test_system_prompt_contains_tensorflow(self) -> None:
        """
        System prompt covers TensorFlow.
        """
        assert "tf.function" in SYSTEM_PROMPT

    def test_system_prompt_contains_jax(self) -> None:
        """
        System prompt covers JAX.
        """
        assert "jax.jit" in SYSTEM_PROMPT


class TestValidateRequest:
    """Tests for _validate_request."""

    def test_empty_source_code(self) -> None:
        """
        Empty source code raises 400.
        """
        data = JitRewriteRequest(
            source_code="",
            trace_id=str(uuid.uuid4()),
            python_version="3.12.0",
        )
        ctx = BaseOptimizerContext("s", "u", "")
        with pytest.raises(Exception, match="empty"):
            _validate_request(data, ctx)

    def test_invalid_trace_id(self) -> None:
        """
        Invalid trace ID raises 400.
        """
        data = JitRewriteRequest(
            source_code="x = 1",
            trace_id="not-a-uuid",
            python_version="3.12.0",
        )
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        with pytest.raises(Exception, match="trace"):
            _validate_request(data, ctx)

    def test_missing_python_version(self) -> None:
        """
        Missing python_version raises 400.
        """
        data = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
        )
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        with pytest.raises(Exception, match="version"):
            _validate_request(data, ctx)

    def test_invalid_python_version(self) -> None:
        """
        Invalid python version format raises 400.
        """
        data = JitRewriteRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
            python_version="2.7.18",
        )
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        with pytest.raises(Exception, match="version"):
            _validate_request(data, ctx)

    def test_syntax_error_source(self) -> None:
        """
        Invalid Python syntax raises 400.
        """
        data = JitRewriteRequest(
            source_code="def f(\n",
            trace_id=str(uuid.uuid4()),
            python_version="3.12.0",
        )
        ctx = BaseOptimizerContext("s", "u", "def f(\n")
        with pytest.raises(Exception, match="syntax|source"):
            _validate_request(data, ctx)

    def test_valid_request(self) -> None:
        """
        Valid request returns python version tuple.
        """
        data = JitRewriteRequest(
            source_code="x = 1\n",
            trace_id=str(uuid.uuid4()),
            python_version="3.12.9",
        )
        ctx = BaseOptimizerContext.get_dynamic_context("s", "u", "x = 1\n")
        result = _validate_request(data, ctx)
        assert (3, 12, 9) == result


class TestJitReusesOptimizeSchemas:
    """Tests that JIT endpoint reuses optimize response types."""

    def test_response_model(self) -> None:
        """
        OptimizeResponse wraps items for JIT results too.
        """
        resp = OptimizeResponse(
            optimizations=[
                OptimizeResponseItem(
                    source_code="import numba\n@numba.njit\ndef f(): pass",
                    explanation="Added JIT",
                    optimization_id="jit-abc",
                )
            ]
        )
        assert 1 == len(resp.optimizations)
        assert "numba" in resp.optimizations[0].source_code

    def test_error_response(self) -> None:
        """
        OptimizeErrorResponse works for JIT errors.
        """
        err = OptimizeErrorResponse(
            error="Could not generate any optimizations."
        )
        assert "optimizations" in err.error
