"""Step definitions for cost budget tracking and enforcement."""

from __future__ import annotations

from behave import given, then, when

from cleveragents.application.services.cost_tracking import CostTrackingService
from cleveragents.core.exceptions import BudgetExceededError
from cleveragents.domain.models.core.cost_budget import (
    BudgetLevel,
    OrgCostAccumulator,
    SessionCostBudget,
    ThreadSafeOrgCostAccumulator,
)
from cleveragents.domain.models.core.cost_metadata import CostMetadata


@given("a cost tracking service is initialized")
def step_init_cost_tracking_service(context: object) -> None:
    """Initialize cost tracking service."""
    context.cost_service = CostTrackingService()


@given("a session cost budget of ${amount} is configured")
def step_configure_session_budget(context: object, amount: str) -> None:
    """Configure session cost budget."""
    budget_amount = float(amount)
    context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)


@given("a cost metadata tracker is created")
def step_create_cost_metadata(context: object) -> None:
    """Create cost metadata tracker."""
    context.cost_metadata = CostMetadata()


@when(
    "an LLM call costs ${cost} with {input_tokens} input tokens and {output_tokens} output tokens"
)
def step_record_llm_call(
    context: object, cost: str, input_tokens: str, output_tokens: str
) -> None:
    """Record an LLM call cost."""
    cost_amount = float(cost)
    input_tok = int(input_tokens)
    output_tok = int(output_tokens)

    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        None,
        input_tokens=input_tok,
        output_tokens=output_tok,
        cost=cost_amount,
        provider="openai",
    )


@then("the cost metadata should show:")
def step_verify_cost_metadata(context: object) -> None:
    """Verify cost metadata values."""
    for row in context.table:
        field = row["field"]
        expected = float(row["value"]) if "." in row["value"] else int(row["value"])
        actual = getattr(context.cost_metadata, field)
        assert actual == expected, f"Expected {field}={expected}, got {actual}"


@then("the session budget should show ${remaining} remaining")
def step_verify_remaining_budget(context: object, remaining: str) -> None:
    """Verify remaining budget."""
    expected = float(remaining)
    actual = context.session_budget.remaining()
    assert actual == expected, f"Expected remaining=${expected}, got ${actual}"


@given("the session budget is ${amount}")
def step_set_session_budget(context: object, amount: str) -> None:
    """Set session budget amount."""
    budget_amount = float(amount)
    context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)
    context.cost_metadata = CostMetadata()


@when("an LLM call costs ${cost:f}")
def step_record_single_cost(context: object, cost: float) -> None:
    """Record a single LLM call cost."""
    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        None,
        input_tokens=100,
        output_tokens=50,
        cost=cost,
        provider="openai",
    )


@when("another LLM call costs ${cost}")
def step_record_another_cost(context: object, cost: str) -> None:
    """Record another LLM call cost."""
    cost_amount = float(cost)
    try:
        context.cost_service.record_cost(
            context.cost_metadata,
            context.session_budget,
            None,
            input_tokens=100,
            output_tokens=50,
            cost=cost_amount,
            provider="openai",
        )
        context.cost_service.enforce_budget(
            context.session_budget, None, context.cost_metadata, provider="openai"
        )
        context.budget_exceeded = False
    except BudgetExceededError as e:
        context.budget_exceeded = True
        context.budget_error = e


@then("the second call should raise BudgetExceededError")
def step_verify_budget_exceeded(context: object) -> None:
    """Verify budget exceeded error was raised."""
    assert context.budget_exceeded, "Expected BudgetExceededError to be raised"


@then("the error should indicate plan budget exceeded")
def step_verify_error_type(context: object) -> None:
    """Verify error indicates plan budget exceeded."""
    assert context.budget_error.budget_type == BudgetLevel.SESSION.value


@when("an LLM call from {provider} costs ${cost}")
def step_record_provider_cost(context: object, provider: str, cost: str) -> None:
    """Record cost from specific provider."""
    cost_amount = float(cost)
    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        None,
        input_tokens=100,
        output_tokens=50,
        cost=cost_amount,
        provider=provider.lower(),
    )


@then("the cost metadata should show provider breakdown:")
def step_verify_provider_breakdown(context: object) -> None:
    """Verify provider cost breakdown."""
    for row in context.table:
        provider = row["provider"]
        expected = float(row["cost"])
        actual = context.cost_metadata.provider_costs.get(provider, 0.0)
        assert actual == expected, f"Expected {provider}=${expected}, got ${actual}"


@when("${amount} has been spent")
def step_spend_amount(context: object, amount: str) -> None:
    """Spend a specific amount."""
    cost_amount = float(amount)
    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        None,
        input_tokens=100,
        output_tokens=50,
        cost=cost_amount,
        provider="openai",
    )


@then("budget status should show:")
def step_verify_budget_status(context: object) -> None:
    """Verify budget status."""
    status = context.cost_service.get_budget_status(context.session_budget, None)
    session_status = status.get("session", {})

    for row in context.table:
        field = row["field"]
        expected = float(row["value"])
        actual = session_status.get(field)
        assert actual == expected, f"Expected {field}={expected}, got {actual}"


@when("a check for ${cost} cost is performed")
def step_check_budget(context: object, cost: str) -> None:
    """Check if cost fits in budget."""
    cost_amount = float(cost)
    org_accumulator = getattr(context, "org_accumulator", None)
    context.budget_check = context.cost_service.check_budget(
        context.session_budget, org_accumulator, cost_amount
    )


@then("the budget check should return warning=true")
def step_verify_warning(context: object) -> None:
    """Verify warning flag is set."""
    assert context.budget_check.warning is True


@then("the check should still be allowed")
def step_verify_allowed(context: object) -> None:
    """Verify check is still allowed."""
    assert context.budget_check.allowed is True


@when("an LLM call costs ${cost} from {provider} using {model}")
def step_record_cost_with_model(
    context: object, cost: str, provider: str, model: str
) -> None:
    """Record cost with provider and model."""
    cost_amount = float(cost)
    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        None,
        input_tokens=100,
        output_tokens=50,
        cost=cost_amount,
        provider=provider.lower(),
        model=model,
    )
    try:
        context.cost_service.enforce_budget(
            context.session_budget,
            None,
            context.cost_metadata,
            provider=provider.lower(),
            model=model,
        )
        context.budget_exceeded = False
    except BudgetExceededError as e:
        context.budget_exceeded = True
        context.budget_error = e


@then("a BudgetExceededError should be raised")
def step_verify_error_raised(context: object) -> None:
    """Verify error was raised."""
    assert context.budget_exceeded, "Expected BudgetExceededError"


@then("the cost metadata should record an exhaustion event with:")
def step_verify_exhaustion_event(context: object) -> None:
    """Verify exhaustion event was recorded."""
    assert len(context.cost_metadata.budget_exhaustion_events) > 0
    event = context.cost_metadata.budget_exhaustion_events[0]

    for row in context.table:
        field = row["field"]
        expected = row["value"]
        if field in ("limit", "used"):
            expected = float(expected)
            actual = getattr(event, field)
        else:
            actual = getattr(event, field)
        assert actual == expected, f"Expected {field}={expected}, got {actual}"


@given("no session budget is configured")
def step_no_session_budget(context: object) -> None:
    """Configure no session budget."""
    context.session_budget = None
    context.cost_metadata = CostMetadata()


@when("multiple LLM calls totaling ${total} are made")
def step_multiple_calls(context: object, total: str) -> None:
    """Make multiple LLM calls."""
    total_cost = float(total)
    # Simulate 5 calls
    per_call = total_cost / 5
    for _ in range(5):
        context.cost_service.record_cost(
            context.cost_metadata,
            context.session_budget,
            None,
            input_tokens=100,
            output_tokens=50,
            cost=per_call,
            provider="openai",
        )


@then("all calls should succeed")
def step_verify_all_succeeded(context: object) -> None:
    """Verify all calls succeeded."""
    # If we got here, no exception was raised
    assert True


@then("the cost metadata should track the total cost")
def step_verify_total_tracked(context: object) -> None:
    """Verify total cost is tracked."""
    assert context.cost_metadata.total_cost > 0


@when("a budget check is performed with a negative cost")
def step_check_negative_cost(context: object) -> None:
    """Check with negative cost."""
    try:
        context.cost_service.check_budget(context.session_budget, None, -1.0)
        context.error_raised = False
    except ValueError as e:
        context.error_raised = True
        context.error = e


@then("a ValueError should be raised for cost validation")
def step_verify_value_error(context: object) -> None:
    """Verify ValueError was raised."""
    assert context.error_raised, "Expected ValueError"


@when("recording cost with negative input tokens")
def step_record_negative_tokens(context: object) -> None:
    """Record cost with negative tokens."""
    try:
        context.cost_service.record_cost(
            context.cost_metadata,
            context.session_budget,
            None,
            input_tokens=-100,
            output_tokens=50,
            cost=1.0,
            provider="openai",
        )
        context.error_raised = False
    except ValueError as e:
        context.error_raised = True
        context.error = e


@given("an organization budget of ${amount}")
def step_configure_org_budget(context: object, amount: str) -> None:
    """Configure organization budget."""
    org = OrgCostAccumulator(org_id="test-org", max_cost_usd=float(amount))
    context.org_accumulator = ThreadSafeOrgCostAccumulator(org)


@given("a session budget of ${amount}")
def step_configure_session_budget_2(context: object, amount: str) -> None:
    """Configure session budget (alternative step for org scenarios)."""
    budget_amount = float(amount)
    context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)


@when("${amount} has been spent at session level")
def step_spend_at_session_level(context: object, amount: str) -> None:
    """Record cost at session and org levels."""
    cost_amount = float(amount)
    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        context.org_accumulator,
        input_tokens=100,
        output_tokens=50,
        cost=cost_amount,
        provider="openai",
    )


@when("${amount} has been spent at organization level")
def step_spend_at_org_level(context: object, amount: str) -> None:
    """Record cost at organization level."""
    cost_amount = float(amount)
    context.cost_service.record_cost(
        context.cost_metadata,
        context.session_budget,
        context.org_accumulator,
        input_tokens=100,
        output_tokens=50,
        cost=cost_amount,
        provider="openai",
    )


@then("the check should fail due to organization budget")
def step_verify_org_budget_denied(context: object) -> None:
    """Verify check was denied due to org budget."""
    assert context.budget_check.allowed is False


@then('the exceeded_level should be "org"')
def step_verify_exceeded_level_org(context: object) -> None:
    """Verify exceeded level is ORG."""
    assert context.budget_check.exceeded_level == BudgetLevel.ORG
