from __future__ import annotations

import importlib
import os
from pathlib import Path
from types import SimpleNamespace
from typing import Final, cast

import pytest

import litellm

from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface
from ...shared.reporting.strategy import ModuleCaseSpec
from ...shared.tracing.profiler import FunctionTraceEvent
from ...shared.tracing.steps import PipelineStep
from .models import RouteFixture, RouteSpec, TraceScenario, TraceSuite
from .reporting import TraceArtifact
from .runner import run_trace_cases, run_trace_scenario, scenario_nodeids, validate_trace_suite
from .sdk.execution import SdkCall, collect_trace, execute_trace


def _fixture(_base_url: str) -> RouteFixture:
    return RouteFixture(kwargs={}, provider_responses=())


def _case(*, surface: Surface = "sdk", function: SdkFunction = "ocr") -> HarnessCase:
    return HarnessCase(
        strategy_id="trace_parity",
        strategy_label="Trace parity",
        sdk_function=function,
        spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"),
        surface=surface,
    )


def test_scenario_filtering_and_occurrence_node_ids() -> None:
    suite: Final = TraceSuite(
        route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
        scenarios=(
            TraceScenario("sync-one", _fixture, asynchronous=False),
            TraceScenario("async-one", _fixture, asynchronous=True),
            TraceScenario("async-two", _fixture, asynchronous=True),
        ),
    )
    case: Final = _case()

    nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"}))

    assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",)


def test_runner_arguments_select_scenarios(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
    case: Final = _case()
    selected: list[frozenset[str]] = []

    def capture_case(
        _run: HarnessRun,
        _case: HarnessCase,
        scenarios: frozenset[str],
        _on_update: object,
    ) -> None:
        selected.append(scenarios)

    monkeypatch.setattr(runner, "_run_case", capture_case)

    exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral",))

    assert exit_code == 0
    assert selected == [frozenset({"mistral"})]


def test_python_trace_preserves_native_ocr_dispatch_setting(monkeypatch: pytest.MonkeyPatch) -> None:
    execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution")
    route: Final = RouteSpec("ocr", ("ocr", "aocr"), _fixture)
    observed: list[str | None] = []

    def collect(
        _function: SdkCall,
        _fixture: RouteFixture,
        *,
        asynchronous: bool,
    ) -> SimpleNamespace:
        observed.append(os.environ.get("LITELLM_RUST"))
        return SimpleNamespace(
            events=(FunctionTraceEvent(0, None, "aocr" if asynchronous else "ocr"),),
            error=None,
        )

    monkeypatch.setattr(execution, "_collect", collect)

    monkeypatch.setenv("LITELLM_RUST", "0")
    collect_trace(route, asynchronous=False)
    monkeypatch.setenv("LITELLM_RUST", "1")
    collect_trace(route, asynchronous=True)

    assert observed == ["0", "1"]
    assert os.environ["LITELLM_RUST"] == "1"


def test_expected_provider_failure_omits_feedback_banner(
    capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
) -> None:
    loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case")
    suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
    scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error")
    monkeypatch.setattr(litellm, "suppress_debug_info", False)

    result: Final = execute_trace(suite.route, scenario, "sdk")

    assert result.python_error is None
    assert "Give Feedback / Get Help" not in capsys.readouterr().out
    assert litellm.suppress_debug_info is False


@pytest.mark.parametrize("asynchronous", (False, True))
def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None:
    loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
    suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
    name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek"
    scenario: Final = next(item for item in suite.scenarios if item.name == name)

    trace: Final = execute_trace(suite.route, scenario, "sdk")

    assert trace.python_error is None
    url: Final = next(
        event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url")
    )
    project: Final = next(
        event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project")
    )
    location: Final = next(
        event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location")
    )
    assert project.parent_id == location.parent_id == url.id
    assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python)


@pytest.mark.parametrize("asynchronous", (False, True))
def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None:
    loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
    suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
    name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials"
    scenario: Final = next(item for item in suite.scenarios if item.name == name)
    monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials")
    monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key")

    trace: Final = execute_trace(suite.route, scenario, "sdk")

    assert trace.python_error is None
    validate: Final = next(
        event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment")
    )
    helpers: Final = (
        "VertexBase.safe_get_vertex_ai_project",
        "VertexBase.safe_get_vertex_ai_credentials",
        "VertexBase.get_access_token",
    )
    assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers
    token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token"))
    load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth"))
    refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth"))
    assert load.parent_id == token.id
    assert refresh.parent_id == load.id
    assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials"
    assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key"


def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None:
    route: Final = RouteSpec("ocr", ("ocr", "aocr"), _fixture)
    duplicate: Final = TraceSuite(
        route=route,
        scenarios=(
            TraceScenario("sync-same", _fixture, asynchronous=False),
            TraceScenario("sync-same", _fixture, asynchronous=False),
        ),
    )
    unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, asynchronous=False),))
    case: Final = _case()

    assert validate_trace_suite(duplicate, case) is not None
    assert validate_trace_suite(unsafe, case) is not None


def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None:
    invalid_name: Final = TraceSuite(
        route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
        scenarios=(TraceScenario("bedrock", _fixture, asynchronous=True),),
    )
    wrong_function: Final = TraceSuite(
        route=RouteSpec("messages", ("create", "acreate"), _fixture),
        scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
    )
    wrong_surface: Final = TraceSuite(
        route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
        scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
    )
    case: Final = _case()

    assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "")
    assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "")
    assert "requires the sdk surface" in (validate_trace_suite(wrong_surface, _case(surface="gateway")) or "")


def test_invalid_route_dispatch_records_harness_error() -> None:
    case: Final = _case()
    run: Final = HarnessRun.from_cases((case,))
    result: Final = run.results[case.key]
    suite: Final = TraceSuite(
        route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
        scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
    )
    nodeid: Final = "trace:gateway:ocr:sync-one"

    run_trace_scenario(run, result, suite, suite.scenarios[0], "gateway", nodeid, lambda _: None)

    assert result.outcomes[nodeid] is RunStatus.ERROR
    assert run.failures == [(nodeid, "trace scenarios only run on the sdk surface")]


def test_python_trace_without_errors_passes(monkeypatch: pytest.MonkeyPatch) -> None:
    runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
    case: Final = _case()
    run: Final = HarnessRun.from_cases((case,))
    result: Final = run.results[case.key]
    suite: Final = TraceSuite(
        route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
        scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
    )
    trace: Final = TraceArtifact.from_traces(
        surface="sdk",
        sdk_function="ocr",
        scenario="sync-one",
        python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),),
    )
    monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace)

    run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None)

    assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED
    assert run.failures == []