"""Actual provider counters survive invalid actions and concurrent task writes."""
from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
import json
from pathlib import Path
import subprocess
import sys
from types import SimpleNamespace

import pytest
from pydantic import BaseModel

from backend.browser import llm_factory
from backend.prospecting import analysis, dialogue, usage
from backend.prospecting.contracts import ProspectingBrief


@pytest.fixture
def ledger(monkeypatch, tmp_path):
    monkeypatch.setattr(usage, "task_doc_dir", lambda tid: tmp_path / tid)
    usage._storage_errors.clear()
    return tmp_path


def one_call(task_id, *, phase="browser", status="succeeded", raw=None):
    with usage.usage_scope(task_id, phase):
        call = usage.start_call("deepseek-flash")
        usage.record_response(call, raw)
        usage.finish_call(call, status)
    return call


def test_new_task_unknown_and_zero_usage_are_distinct(ledger):
    usage.initialize_task("new-task")
    assert usage.snapshot("new-task")["input_tokens"] is None
    one_call("new-task", raw={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
    one_call("new-task")
    state = usage.snapshot("new-task")
    assert state["requests"] == 2 and state["input_tokens"] == state["output_tokens"] == 0
    assert state["unknown_usage_requests"] == state["incomplete_usage_requests"] == 1
    assert state["legacy_incomplete"] is False
    assert state["cache_hit_tokens"] is None


def test_old_task_cannot_be_relabelled_as_fully_instrumented(ledger):
    one_call("old-task", raw={"prompt_tokens": 8, "completion_tokens": 3, "total_tokens": 11})
    usage.initialize_task("old-task")
    state = usage.snapshot("old-task")
    assert state["legacy_incomplete"] is True and state["requests"] == 1
    assert state["coverage"] == "recorded_calls_only"
    assert usage.snapshot("no-records")["total_tokens"] is None


def test_response_is_idempotent_and_private_fields_never_persist(ledger):
    usage.initialize_task("task")
    with usage.usage_scope("task", "browser"):
        call = usage.start_call("deepseek-flash")
        raw = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13,
               "prompt_cache_hit_tokens": 8, "prompt_cache_miss_tokens": 2,
               "api_key": "NEVER-PERSIST", "prompt": "PRIVATE-PROMPT", "screenshot": "PRIVATE-IMAGE"}
        usage.record_response(call, raw, model="deepseek-flash-actual")
        usage.record_response(call, {"prompt_tokens": 999})
        usage.finish_call(call, "succeeded")
        usage.finish_call(call, "schema_error")
    state = usage.snapshot("task")
    assert state["requests"] == 1 and state["input_tokens"] == 10 and state["output_tokens"] == 3
    assert state["cache_hit_tokens"] == 8 and state["cache_miss_tokens"] == 2
    assert state["failed_requests"] == 0 and state["models"] == ["deepseek-flash-actual"]
    text = (ledger / "task" / "model-usage.jsonl").read_text()
    assert all(value not in text for value in ("NEVER-PERSIST", "PRIVATE-PROMPT", "PRIVATE-IMAGE"))


def test_nested_phase_inherits_task_and_scope_is_restored(ledger):
    with usage.usage_scope("task", "browser"):
        first = usage.start_call("deepseek-flash")
        with usage.usage_scope(phase="compression"):
            compressed = usage.start_call("deepseek-flash")
            usage.record_response(compressed, {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12})
            usage.finish_call(compressed, "succeeded")
        assert usage.current_scope() == ("task", "browser")
        usage.finish_call(first, "connection_error")
    assert usage.current_scope() == (None, "unknown") and usage.start_call("deepseek-flash") is None
    state = usage.snapshot("task")
    assert state["requests"] == 2 and {row["phase"] for row in state["phases"]} == {"browser", "compression"}
    assert state["unknown_usage_requests"] == 1


def test_deepseek_cache_top_level_precedes_nested_and_partial_stays_unknown():
    state = usage.normalize_usage(SimpleNamespace(prompt_tokens=10, completion_tokens=3, total_tokens=13,
                                                 prompt_cache_hit_tokens=8, prompt_cache_miss_tokens=2,
                                                 prompt_tokens_details=SimpleNamespace(cached_tokens=0)))
    assert state["cache_hit_tokens"] == 8 and state["cache_miss_tokens"] == 2
    state = usage.normalize_usage({"prompt_tokens": 10, "prompt_tokens_details": {"cached_tokens": 7}})
    assert state["cache_hit_tokens"] == 7 and state["cache_miss_tokens"] == 3
    assert state["output_tokens"] is state["total_tokens"] is None
    assert usage.normalize_usage({"prompt_tokens": True, "completion_tokens": -1})["input_tokens"] is None


def test_threaded_appends_do_not_lose_calls(ledger):
    usage.initialize_task("shared")
    with ThreadPoolExecutor(max_workers=8) as pool:
        list(pool.map(lambda i: one_call("shared", raw={"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}), range(32)))
    state = usage.snapshot("shared")
    assert state["requests"] == 32 and state["total_tokens"] == 192
    assert state["pending_requests"] == 0 and state["ledger_error"] is False


def test_process_appends_are_locked(ledger):
    usage.initialize_task("shared-process")
    script = """from pathlib import Path
import sys
from backend.prospecting import usage
usage.task_doc_dir=lambda tid:Path(sys.argv[1])/tid
for _ in range(4):
 with usage.usage_scope('shared-process','analysis'):
  c=usage.start_call('deepseek-flash')
  usage.record_response(c,{'prompt_tokens':2,'completion_tokens':1,'total_tokens':3})
  usage.finish_call(c,'succeeded')
"""
    processes = [subprocess.Popen([sys.executable, "-c", script, str(ledger)], stdout=subprocess.PIPE, stderr=subprocess.PIPE) for _ in range(3)]
    for process in processes:
        stdout, stderr = process.communicate(timeout=15)
        assert process.returncode == 0, (stdout.decode(), stderr.decode())
    state = usage.snapshot("shared-process")
    assert state["requests"] == 12 and state["total_tokens"] == 36 and not state["ledger_error"]


def test_truncated_ledger_preserves_known_data_and_marks_gap(ledger):
    usage.initialize_task("damaged")
    path = ledger / "damaged" / "model-usage.jsonl"
    with path.open("ab") as stream:
        stream.write(b'{"interrupted":')
    one_call("damaged", raw={"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3})
    state = usage.snapshot("damaged")
    assert state["requests"] == 1 and state["total_tokens"] == 3 and state["ledger_error"]


def test_storage_failure_cannot_turn_a_model_response_into_paid_retry(ledger):
    (ledger / "bad").write_text("not a directory")
    one_call("bad", raw={"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3})
    state = usage.snapshot("bad")
    assert state["ledger_error"] and state["total_tokens"] is None


@pytest.mark.parametrize("field,bad_value", [
    ("at", {"secret": "PRIVATE-NESTED-METADATA"}),
    ("at", "PRIVATE-ARBITRARY-TEXT"),
    ("at", "2026-10-29T10:00:00"),
    ("kind", {"secret": "PRIVATE-KIND"}),
    ("kind", "not-an-event"),
    ("call_id", "wrong-call"),
    ("event_id", "f" * 32 + ":response"),
    ("phase", {"secret": "PRIVATE-PHASE"}),
    ("model", {"secret": "PRIVATE-MODEL"}),
    ("prompt", "PRIVATE-UNEXPECTED-CONTENT"),
    ("finish_reason", "PRIVATE-UNKNOWN-FINISH"),
    ("usage", {"input_tokens": "PRIVATE-NON-NUMERIC"}),
])
def test_valid_json_bad_metadata_is_skipped_without_leak(ledger, field, bad_value):
    usage.initialize_task("bad-metadata")
    one_call("bad-metadata", raw={"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3})
    path = ledger / "bad-metadata" / "model-usage.jsonl"
    rows = [json.loads(line) for line in path.read_text().splitlines()]
    bad = next(row for row in rows if row.get("kind") == "response")
    bad[field] = bad_value
    with path.open("a") as stream:
        stream.write(json.dumps(bad) + "\n")
    state = usage.snapshot("bad-metadata")
    assert state["requests"] == 1 and state["total_tokens"] == 3
    assert state["ledger_error"] and "PRIVATE" not in json.dumps(state)


def test_incremental_analysis_has_its_own_phase(ledger):
    one_call("incremental", phase="analysis_incremental", raw={"prompt_tokens": 6, "completion_tokens": 1, "total_tokens": 7})
    state = usage.snapshot("incremental")
    assert state["phases"][0]["phase"] == "analysis_incremental"


class BrowserAction(BaseModel):
    action: str


async def test_browser_schema_failure_still_records_provider_tokens(ledger, monkeypatch):
    import openai
    class Client:
        def __init__(self, **kwargs):
            self.chat = SimpleNamespace(completions=SimpleNamespace(create=self.create))
        async def __aenter__(self):
            return self
        async def __aexit__(self, *args):
            return False
        async def create(self, **kwargs):
            return SimpleNamespace(model="deepseek-flash", choices=[SimpleNamespace(message=SimpleNamespace(content=None, tool_calls=[
                SimpleNamespace(type="function", function=SimpleNamespace(name="return_structured_output", arguments='{"unexpected":1}'))
            ]), finish_reason="tool_calls")],
                                   usage=SimpleNamespace(prompt_tokens=20, completion_tokens=4, total_tokens=24,
                                                         prompt_cache_hit_tokens=17, prompt_cache_miss_tokens=3))
    monkeypatch.setattr(openai, "AsyncOpenAI", Client)
    from browser_use.llm.exceptions import ModelProviderError
    from browser_use.llm.messages import UserMessage
    with usage.usage_scope("invalid-action", "browser"):
        with pytest.raises(ModelProviderError, match="不符合"):
            await llm_factory.DeepSeekBrowserChat("deepseek-flash", "PRIVATE-TEST-KEY").ainvoke([UserMessage(content="PRIVATE-PROMPT")], BrowserAction)
    state = usage.snapshot("invalid-action")
    assert state["total_tokens"] == 24 and state["cache_hit_tokens"] == 17
    assert state["parse_failed_requests"] == state["failed_requests"] == 1
    assert "PRIVATE" not in (ledger / "invalid-action" / "model-usage.jsonl").read_text()


async def test_invalid_analysis_keeps_usage_and_parent_context(ledger, monkeypatch):
    import httpx
    calls = []
    body = {"model": "deepseek-flash", "choices": [{"message": {"content": '{"results":[],"summary":"none"}'}}],
            "usage": {"prompt_tokens": 14, "completion_tokens": 2, "total_tokens": 16}}
    class Client:
        def __init__(self, **kwargs):
            pass
        async def __aenter__(self):
            return self
        async def __aexit__(self, *args):
            return False
        async def post(self, url, **kwargs):
            calls.append(kwargs)
            return httpx.Response(200, json=body)
    monkeypatch.setattr(analysis.httpx, "AsyncClient", Client)
    monkeypatch.setattr(analysis, "get_llm_config", lambda: {"api_key": "PRIVATE-KEY", "model": "deepseek-flash", "base_url": "https://fixture", "provider": "deepseek"})
    note = {"source_id": "note", "source_type": "note", "source_url": "https://fixture/n", "content": "parent context"}
    reply = {"source_id": "reply", "source_type": "reply", "source_url": "https://fixture/n", "parent_source_id": "note", "content": "own words"}
    with usage.usage_scope("bad-analysis", "review"):
        with pytest.raises(ValueError, match="来源"):
            await analysis.analyze_sources(ProspectingBrief(intent="research"), [reply], context_sources=[note, reply])
    assert '"source_id": "note"' in calls[0]["json"]["messages"][0]["content"]
    state = usage.snapshot("bad-analysis")
    assert state["requests"] == 1 and state["total_tokens"] == 16 and state["parse_failed_requests"] == 1
    assert state["phases"][0]["phase"] == "review"


async def test_intent_validation_failure_is_metered(ledger, monkeypatch):
    import httpx
    class Client:
        def __init__(self, **kwargs):
            pass
        async def __aenter__(self):
            return self
        async def __aexit__(self, *args):
            return False
        async def post(self, url, **kwargs):
            return httpx.Response(200, json={"choices": [{"message": {"content": "not JSON"}}],
                                           "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}})
    monkeypatch.setattr(dialogue.httpx, "AsyncClient", Client)
    monkeypatch.setattr(dialogue, "get_llm_config", lambda: {"api_key": "PRIVATE-KEY", "model": "deepseek-flash", "base_url": "https://fixture", "provider": "deepseek"})
    with usage.usage_scope("bad-intent", "intent"):
        with pytest.raises(ValueError, match="可用计划"):
            await dialogue.interpret_message("找客户", ProspectingBrief(), [])
    state = usage.snapshot("bad-intent")
    assert state["requests"] == 1 and state["total_tokens"] == 11 and state["parse_failed_requests"] == 1


async def test_legacy_wrapper_preserves_request_and_only_counts_once(ledger):
    class Wrapped:
        provider, name = "fixture", "fixture-model"
        def __init__(self):
            self.calls = []
        def __repr__(self):
            return "PRIVATE-KEY"
        async def ainvoke(self, messages, output_format=None, **kwargs):
            self.calls.append((messages, output_format, kwargs))
            return SimpleNamespace(completion="same response", usage=SimpleNamespace(prompt_tokens=9, completion_tokens=2,
                                                                                      total_tokens=11, prompt_cached_tokens=7))
    wrapped = Wrapped()
    model = llm_factory.ObservedChatModel(wrapped, "fixture-model")
    with usage.usage_scope("legacy", "compression"):
        result = await model.ainvoke(["unchanged"], BrowserAction, temperature=0)
    assert result.completion == "same response" and wrapped.calls == [(["unchanged"], BrowserAction, {"temperature": 0})]
    state = usage.snapshot("legacy")
    assert state["requests"] == 1 and state["total_tokens"] == 11 and state["cache_hit_tokens"] == 7
    assert "PRIVATE-KEY" not in repr(model)


async def test_legacy_failed_adapter_records_unknown_without_retry(ledger):
    class Wrapped:
        provider, name = "fixture", "fixture-model"
        calls = 0
        async def ainvoke(self, *args, **kwargs):
            self.calls += 1
            raise ValueError("PRIVATE response content must never persist")
    wrapped = Wrapped()
    with usage.usage_scope("legacy-failed", "browser"):
        with pytest.raises(ValueError):
            await llm_factory.ObservedChatModel(wrapped, "fixture-model").ainvoke([])
    state = usage.snapshot("legacy-failed")
    assert wrapped.calls == state["requests"] == state["unknown_usage_requests"] == 1
    assert state["input_tokens"] is None and state["output_tokens"] is None
    assert "PRIVATE" not in (ledger / "legacy-failed" / "model-usage.jsonl").read_text()