#include "Insights/ViewModels/MinimalTimerExporter.h"
#include "TraceServices/Model/Threads.h"
#include "Misc/Paths.h"
#include "CborReader.h"

namespace TraceTimer
{
ANSICHAR Term = '\n';
FORCEINLINE FArchive& operator<<(FArchive& Ar, const ANSICHAR* Str)
{
    if (Str)
    {
        Ar.Serialize((void*)Str, FCStringAnsi::Strlen(Str));
    }
    return Ar;
}
FORCEINLINE FArchive& operator<<(FArchive& Ar, const TCHAR* Str)
{
    if (Str)
    {
        FTCHARToUTF8 Converter(Str);
        Ar.Serialize((void*)Converter.Get(), Converter.Length());
    }
    return Ar;
}

FORCEINLINE FArchive& operator<<(FArchive& Ar, uint32 Value)
{
    ANSICHAR Buffer[16];
    int32 Len = FCStringAnsi::Snprintf(Buffer, 16, "%u", Value);
    Ar.Serialize(Buffer, Len);
    return Ar;
}

FORCEINLINE FArchive& operator<<(FArchive& Ar, double Value)
{
    ANSICHAR Buffer[32];
    int32 Len = FCStringAnsi::Snprintf(Buffer, 32, "%.9f", Value);
    Ar.Serialize(Buffer, Len);
    return Ar;
}

// --- FExportContext 实现 ---
FExportContext::FExportContext(const Trace::IAnalysisSession& InSession): Session(InSession) {}

bool FExportContext::Initialize(const Trace::ITimingProfilerProvider* InProvider)
{
    if (!InProvider)
        return false;
    Provider = InProvider;

    Trace::FAnalysisSessionReadScope _(Session);
    Provider->ReadTimers([this](const Trace::ITimingProfilerTimerReader& Out)
                         {
                             TimerReader = &Out;
                         });

    const Trace::IThreadProvider& ThreadProvider = Trace::ReadThreadProvider(Session);
    ThreadProvider.EnumerateThreads([this](const Trace::FThreadInfo& Thread)
                                    {
                                        uint32 TimelineIndex = 0;
                                        if (Provider->GetCpuThreadTimelineIndex(Thread.Id, TimelineIndex))
                                            TimelineIndexToThreadId.Add(TimelineIndex, Thread.Id);
                                    });

    uint32 GpuIndex = 0;
    if (Provider->GetGpuTimelineIndex(GpuIndex))
        TimelineIndexToThreadId.Add(GpuIndex, 0xFFFFFFFF);
    if (Provider->GetGpu2TimelineIndex(GpuIndex))
        TimelineIndexToThreadId.Add(GpuIndex, 0xFFFFFFFE);

    return TimerReader != nullptr;
}

uint32 FExportContext::GetThreadId(uint32 TimelineIndex) const
{
    const uint32* Found = TimelineIndexToThreadId.Find(TimelineIndex);
    return Found ? *Found : 0;
}

void CsvEventWriter::WriteHeader(FArchive& Ar)
{
    if (EnumHasAnyFlags(Fields, EExportField::ThreadId)) Ar << TEXT("ThreadId,");
    if (EnumHasAnyFlags(Fields, EExportField::Name))     Ar << TEXT("Name,");
    if (EnumHasAnyFlags(Fields, EExportField::Start))    Ar << TEXT("Start,");
    if (EnumHasAnyFlags(Fields, EExportField::End))      Ar << TEXT("End,");
    if (EnumHasAnyFlags(Fields, EExportField::Duration)) Ar << TEXT("Duration,");
    if (EnumHasAnyFlags(Fields, EExportField::Depth))    Ar << TEXT("Depth"); 
    if (EnumHasAnyFlags(Fields, EExportField::Metadata)) Ar << TEXT("Metadata"); 
    
    Ar << Term;
}

void CsvEventWriter::ProcessEvent(FArchive& Ar, const FEventExportData& Data)
{
    // 技巧：定义一个控制变量处理逗号，虽然多一行，但能保证 CSV 格式完美
    bool bNeedComma = false;
    auto WriteDelimiter = [&]() { if (bNeedComma) Ar << TEXT(","); bNeedComma = true; };

    if (EnumHasAnyFlags(Fields, EExportField::ThreadId))
    {
        WriteDelimiter();
        Ar << Data.ThreadId;
    }

    if (EnumHasAnyFlags(Fields, EExportField::Name))
    {
        WriteDelimiter();
        // 链式写入引号和内容，简洁明了
        Ar << TEXT("\"") << (Data.Timer ? Data.Timer->Name : TEXT("Unknown")) << TEXT("\"");
    }

    // 浮点数部分直接内联转换，不给变量起名
    if (EnumHasAnyFlags(Fields, EExportField::Start))
    {
        WriteDelimiter();
        Ar << *FString::Printf(TEXT("%.9f"), Data.Start);
    }

    if (EnumHasAnyFlags(Fields, EExportField::End))
    {
        WriteDelimiter();
        Ar << *FString::Printf(TEXT("%.9f"), Data.End);
    }

    if (EnumHasAnyFlags(Fields, EExportField::Duration))
    {
        WriteDelimiter();
        Ar << *FString::Printf(TEXT("%.9f"), Data.End - Data.Start);
    }

    if (EnumHasAnyFlags(Fields, EExportField::Depth))
    {
        WriteDelimiter();
        Ar << Data.Depth;
    }

    if (EnumHasAnyFlags(Fields, EExportField::Metadata))
    {
        WriteDelimiter();
        FString OutMetadata = TEXT("");
        FTraceExportProcessor::DecodeMetadataToString(OutMetadata, Data.Metadata);
        Ar << *OutMetadata;
    }

    Ar << Term;
}

void JsonEventWriter::ProcessEvent(FArchive& Ar, const FEventExportData& Data)
{
    WriteNode(Ar, Data);
    Ar << Term;
}

void JsonEventWriter::WriteNode(FArchive& Ar, const FEventExportData& Node)
{
    Ar << TEXT("{");

    bool bNeedComma = false;
    auto WriteKey = [&](const TCHAR* Key) {
        if (bNeedComma) Ar << TEXT(",");
        Ar << TEXT("\"") << Key << TEXT("\":");
        bNeedComma = true;
    };

    if (EnumHasAnyFlags(Fields, EExportField::Name))
    {
        WriteKey(TEXT("name"));
        Ar << TEXT("\"") << (Node.Timer ? Node.Timer->Name : TEXT("Unknown")) << TEXT("\"");
    }

    if (EnumHasAnyFlags(Fields, EExportField::Start))
    {
        WriteKey(TEXT("start"));
        Ar << *FString::Printf(TEXT("%.9f"), Node.Start);
    }

    if (EnumHasAnyFlags(Fields, EExportField::End))
    {
        WriteKey(TEXT("end"));
        Ar << *FString::Printf(TEXT("%.9f"), Node.End);
    }

    if (EnumHasAnyFlags(Fields, EExportField::Duration))
    {
        WriteKey(TEXT("duration"));
        Ar << *FString::Printf(TEXT("%.9f"), Node.End - Node.Start);
    }

    if (EnumHasAnyFlags(Fields, EExportField::Depth))
    {
        WriteKey(TEXT("depth"));
        Ar << Node.Depth;
    }

    if (EnumHasAnyFlags(Fields, EExportField::ThreadId))
    {
        WriteKey(TEXT("threadId"));
        Ar << Node.ThreadId;
    }
    if (EnumHasAnyFlags(Fields, EExportField::Metadata) && Node.Metadata.Num() > 0)
    {
        WriteKey(TEXT("metadata"));
        FString OutMetadata = TEXT("");
        FTraceExportProcessor::DecodeMetadataToString(OutMetadata, Node.Metadata);
        Ar << *OutMetadata;
    }

    if (Node.Children.Num() > 0)
    {
        WriteKey(TEXT("children"));
        Ar << TEXT("[");
        for (int32 i = 0; i < Node.Children.Num(); ++i)
        {
            if (i > 0) Ar << TEXT(",");
            WriteNode(Ar, *Node.Children[i]);
        }
        Ar << TEXT("]");
    }

    Ar << TEXT("}");
}

FExportConfig::FExportConfig(const FString& Filename)
{
    Folder    = FPaths::GetPath(Filename);
    BaseName  = FPaths::GetBaseFilename(Filename);
    Extension = FPaths::GetExtension(Filename);
}

void ChunkingFileWriter::Initialize(const FExportConfig& InConfig)
{
    Config = InConfig;
    PersistentBuffer.Reset();
    MemAr = MakeUnique<FMemoryWriter>(PersistentBuffer, false);
    FileAr.Reset();
    ChunkIndex = 1;
    WrittenEventCount = 0;
    CurrentFilename.Empty();
}

FArchive* ChunkingFileWriter::GetCurrentArchive()
{
    return MemAr.Get();
}

void ChunkingFileWriter::CheckFlush(IEventWriter* Writer)
{
    if (PersistentBuffer.Num() >= (int32)Config.MaxMemSize)
    {
        FlushToDisk();
    }
}

void ChunkingFileWriter::FlushToDisk()
{
    if (!FileAr || PersistentBuffer.Num() == 0)
        return;
    
    FileAr->Serialize(PersistentBuffer.GetData(), PersistentBuffer.Num());
    PersistentBuffer.Reset();
    MemAr->Seek(0);
}

void ChunkingFileWriter::WriteEvent(IEventWriter* Writer, const FEventExportData& Data)
{
    if (!Writer)
        return;
    
    Writer->ProcessEvent(*MemAr, Data);
    WrittenEventCount++;
    
    CheckFlush(Writer);
    
    if (WrittenEventCount >= Config.EventsPerChunk)
    {
        FlushToDisk();
        SwitchToNextChunk(Writer);
    }
}

void ChunkingFileWriter::SwitchToNextChunk(IEventWriter* Writer)
{
    FileAr.Reset();
    
    if (!Writer)
        return;
    
    FString PartSuffix = ChunkIndex > 1 ? FString::Printf(TEXT("_part%u"), ChunkIndex) : TEXT("");
    CurrentFilename = FPaths::Combine(Config.Folder, FString::Printf(TEXT("%s%s.%s"), *Config.BaseName, *PartSuffix, *Config.Extension));
    
    FileAr.Reset(IFileManager::Get().CreateFileWriter(*CurrentFilename));
    if (FileAr)
    {
        PersistentBuffer.Reset();
        MemAr = MakeUnique<FMemoryWriter>(PersistentBuffer, false);
        Writer->WriteHeader(*MemAr);
    }
    
    WrittenEventCount = 0;
    ChunkIndex++;
}

void ChunkingFileWriter::Close()
{
    FlushToDisk();
    FileAr.Reset();
    PersistentBuffer.Reset();
    MemAr.Reset();
    ChunkIndex = 1;
    WrittenEventCount = 0;
    CurrentFilename.Empty();
}

ChunkingFileWriter::~ChunkingFileWriter()
{
    Close();
}
bool FExportFilter::ShouldFilterByEvent(FEventExportData& Node) const
{
    if (bOnlyMetadata && Node.Metadata.Num())
    {
        return false;
    }
    return true;
}
FTraceExportProcessor::FTraceExportProcessor(FExportContext& InContext, IEventWriter& InWriter, ChunkingFileWriter& InChunkWriter, const FExportFilter& InFilter)
    : Context(InContext), Writer(InWriter), ChunkWriter(InChunkWriter), Filter(InFilter) {}

void FTraceExportProcessor::Run()
{
    Trace::FAnalysisSessionReadScope _(Context.Session);
    uint32 TimelineCount = Context.Provider->GetTimelineCount();
    for (uint32 i = 0; i < TimelineCount; ++i)
    {
        if (Filter.ShouldFilterByTrack(i))
            continue;
        ProcessTimeline(i);
    }
}
void FTraceExportProcessor::ProcessTimeline(uint32 TimelineIndex)
{
    uint32 ThreadId = Context.GetThreadId(TimelineIndex);
    Context.Provider->ReadTimeline(TimelineIndex, [this, ThreadId](const Trace::ITimingProfilerProvider::Timeline& Timeline)
                                   {
                                       if (Writer.WantsStructuredPath())
                                           EnumerateStructured(Timeline, ThreadId);
                                       else
                                           EnumerateFlat(Timeline, ThreadId);
                                   });
}

void FTraceExportProcessor::EnumerateFlat(const Trace::ITimingProfilerProvider::Timeline& Timeline, uint32 ThreadId)
{
    Timeline.EnumerateEvents(Filter.GetStartTime(), Filter.GetEndTime(), [&](double S, double E, uint32 Depth, const Trace::FTimingProfilerEvent& Event)
                             {
                                 FEventExportData Data;
                                 Data.Timer    = Context.TimerReader->GetTimer(Event.TimerIndex);
                                 Data.Start    = S;
                                 Data.End      = E;
                                 Data.Depth    = Depth;
                                 Data.ThreadId = ThreadId;
                                 Data.Metadata = Context.TimerReader->GetMetadata(Event.TimerIndex);

                                 if (Filter.ShouldFilterByEvent(Data))
                                     return Trace::EEventEnumerate::Continue;

                                 ChunkWriter.WriteEvent(&Writer, Data);
                                 return Trace::EEventEnumerate::Continue;
                             });
}

void FTraceExportProcessor::EnumerateStructured(const Trace::ITimingProfilerProvider::Timeline& Timeline, uint32 ThreadId)
{
    TArray<TSharedPtr<FEventExportData>> Stack;
    FEventExportData StaticNode;
    Timeline.EnumerateEvents(Filter.GetStartTime(), Filter.GetEndTime(), [&](double S, double E, uint32 Depth, const Trace::FTimingProfilerEvent& Event)
                             {
                                 StaticNode.Timer    = Context.TimerReader->GetTimer(Event.TimerIndex);
                                 StaticNode.Start    = S;
                                 StaticNode.End      = E;
                                 StaticNode.Depth    = Depth;
                                 StaticNode.ThreadId = ThreadId;
                                 StaticNode.Metadata = Context.TimerReader->GetMetadata(Event.TimerIndex);

                                 if (Filter.ShouldFilterByEvent(StaticNode))
                                     return Trace::EEventEnumerate::Continue;

                                 ChunkWriter.NotifyEventWritten();
                                 TSharedPtr<FEventExportData> Node = MakeShared<FEventExportData>();
                                 StaticNode.MoveTo(*Node);

                                 if (Depth == 0)
                                 {
                                     if (Stack.Num() > 0)
                                     {
                                         ChunkWriter.WriteEvent(&Writer, *Stack[0]);
                                     }
                                     Stack.Empty();
                                     Stack.Add(Node);
                                 }
                                 else
                                 {
                                     while (Stack.Num() > (int32)Depth)
                                         Stack.Pop();
                                     if (Stack.Num() > 0)
                                         Stack.Last()->Children.Add(Node);
                                     Stack.Add(Node);
                                 }
                                 return Trace::EEventEnumerate::Continue;
                             });
    if (Stack.Num() > 0)
    {
        ChunkWriter.WriteEvent(&Writer, *Stack[0]);
    }
}
void FTraceExportProcessor::DecodeMetadataToString(FString& Str, const TArrayView<const uint8>& Metadata)
{
    FMemoryReaderView MemoryReader(Metadata);
    FCborReader CborReader(&MemoryReader, ECborEndianness::StandardCompliant);
    FCborContext Context;

    if (!CborReader.ReadNext(Context) || Context.MajorType() != ECborCode::Map)
    {
        return;
    }

    bool bFirst = true;

    while (true)
    {
        // Read key
        if (!CborReader.ReadNext(Context) || !Context.IsString())
        {
            break;
        }

        if (bFirst)
        {
            bFirst = false;
            // Str += TEXT(" - ");
        }
        else
        {
            Str += TEXT(", ");
        }

        FString Key(Context.AsLength(), Context.AsCString());
        Str += Key;
        Str += TEXT(":");

        if (!CborReader.ReadNext(Context))
        {
            break;
        }

        switch (Context.MajorType())
        {
            case ECborCode::Int:
            case ECborCode::Uint: {
                uint64 Value = Context.AsUInt();
                Str += FString::Printf(TEXT("%llu"), Value);
                continue;
            }

            case ECborCode::TextString: {
                Str += Context.AsString();
                continue;
            }

            case ECborCode::ByteString: {
                Str.AppendChars(Context.AsCString(), Context.AsLength());
                continue;
            }
        }

        if (Context.RawCode() == (ECborCode::Prim | ECborCode::Value_4Bytes))
        {
            float Value = Context.AsFloat();
            Str += FString::Printf(TEXT("%f"), Value);
            continue;
        }

        if (Context.RawCode() == (ECborCode::Prim | ECborCode::Value_8Bytes))
        {
            double Value = Context.AsDouble();
            Str += FString::Printf(TEXT("%g"), Value);
            continue;
        }

        if (Context.IsFiniteContainer())
        {
            CborReader.SkipContainer(ECborCode::Array);
        }
    }
}

}  // namespace TraceTimer

// --- FMinimalTimerExporter 静态入口实现 ---

bool FMinimalTimerExporter::ExportTimersToCSV(const Trace::IAnalysisSession& Session, const FString& Filename)
{
    return false;
}

bool FMinimalTimerExporter::ExportMetadataToCSV(const Trace::IAnalysisSession& Session, const FString& Filename, const TraceTimer::FExportFilter& InFilter)
{
    TraceTimer::FExportFilter Filter = InFilter;
    Filter.bOnlyMetadata             = true;

    TraceTimer::CsvEventWriter Writer(Filter.Fields);
    return ExecuteTimingExport(Session, Filename, Filter, Writer);
}

bool FMinimalTimerExporter::ExportTimingEventsToCSV(const Trace::IAnalysisSession& Session, const FString& Filename, const TraceTimer::FExportFilter& InFilter)
{
    TraceTimer::FExportFilter Filter = InFilter;
    TraceTimer::CsvEventWriter Writer(Filter.Fields);
    return ExecuteTimingExport(Session, Filename, Filter, Writer);
}

bool FMinimalTimerExporter::ExportTimingEventsToJSON(const Trace::IAnalysisSession& Session, const FString& Filename, const TraceTimer::FExportFilter& InFilter)
{
    TraceTimer::FExportFilter Filter = InFilter;
    TraceTimer::JsonEventWriter Writer(Filter.Fields);
    return ExecuteTimingExport(Session, Filename, Filter, Writer);
}

bool FMinimalTimerExporter::ExecuteTimingExport(const Trace::IAnalysisSession& Session, const FString& Filename, TraceTimer::FExportFilter& Filter, TraceTimer::IEventWriter& Writer)
{
    using namespace TraceTimer;
    FExportContext Context(Session);
    if (!Context.Initialize(Trace::ReadTimingProfilerProvider(Session)))
        return false;

    FExportConfig Config(Filename);
    
    ChunkingFileWriter ChunkWriter;
    ChunkWriter.Initialize(Config);
    ChunkWriter.SwitchToNextChunk(&Writer);

    FTraceExportProcessor Processor(Context, Writer, ChunkWriter, Filter);
    Processor.Run();
    return true;
}