using System.Text.Json;

namespace ClaudeDo.Worker.Runner;

public sealed class StreamResult
{
    public string? ResultMarkdown { get; set; }
    public string? StructuredOutputJson { get; set; }
    public string? SessionId { get; set; }
    public int TurnCount { get; set; }
    public int TokensIn { get; set; }
    public int TokensOut { get; set; }
    public int ApiRetryCount { get; set; }
}

public sealed class StreamAnalyzer
{
    private string? _resultMarkdown;
    private string? _structuredOutputJson;
    private string? _sessionId;
    private int _turnCount;
    private int _tokensIn;
    private int _tokensOut;
    private int _apiRetryCount;

    public void ProcessLine(string ndjsonLine)
    {
        if (string.IsNullOrWhiteSpace(ndjsonLine)) return;

        try
        {
            using var doc = JsonDocument.Parse(ndjsonLine);
            var root = doc.RootElement;

            if (!root.TryGetProperty("type", out var typeProp)) return;
            var type = typeProp.GetString();

            switch (type)
            {
                case "result":
                    if (root.TryGetProperty("result", out var resultProp))
                        _resultMarkdown = resultProp.GetString();
                    if (root.TryGetProperty("structured_output", out var structuredProp))
                        _structuredOutputJson = structuredProp.ToString();
                    if (root.TryGetProperty("session_id", out var sessionProp))
                        _sessionId = sessionProp.GetString();
                    // Authoritative token totals live on the result event.
                    if (root.TryGetProperty("usage", out var resultUsage))
                    {
                        if (resultUsage.TryGetProperty("input_tokens", out var inp))
                            _tokensIn = inp.GetInt32();
                        if (resultUsage.TryGetProperty("output_tokens", out var outp))
                            _tokensOut = outp.GetInt32();
                    }
                    break;

                case "assistant":
                    _turnCount++;
                    break;

                case "system":
                    if (root.TryGetProperty("subtype", out var subtypeProp) &&
                        subtypeProp.GetString() == "api_retry")
                        _apiRetryCount++;
                    break;

                case "stream_event":
                    TryAccumulateUsage(root);
                    break;
            }
        }
        catch (JsonException) { /* Malformed JSON — skip */ }
    }

    public StreamResult GetResult() => new()
    {
        ResultMarkdown = FallbackResult(),
        StructuredOutputJson = _structuredOutputJson,
        SessionId = _sessionId,
        TurnCount = _turnCount,
        TokensIn = _tokensIn,
        TokensOut = _tokensOut,
        ApiRetryCount = _apiRetryCount,
    };

    private string? FallbackResult()
    {
        if (!string.IsNullOrEmpty(_resultMarkdown)) return _resultMarkdown;
        if (_structuredOutputJson is null) return _resultMarkdown;
        try
        {
            using var doc = JsonDocument.Parse(_structuredOutputJson);
            if (doc.RootElement.TryGetProperty("summary", out var s))
                return s.GetString();
        }
        catch { }
        return _structuredOutputJson;
    }

    private void TryAccumulateUsage(JsonElement root)
    {
        if (!root.TryGetProperty("event", out var eventProp)) return;
        if (eventProp.TryGetProperty("message", out var msgProp) &&
            msgProp.TryGetProperty("usage", out var usageProp))
        {
            if (usageProp.TryGetProperty("input_tokens", out var inp))
                _tokensIn += inp.GetInt32();
            if (usageProp.TryGetProperty("output_tokens", out var outp))
                _tokensOut += outp.GetInt32();
        }
    }
}