using ClawdDotNet.Core.Accounting;
using ClawdDotNet.Core.Audit;
using ClawdDotNet.Core.Budget;
using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Engine;
using ClawdDotNet.Core.Storage;
using ClawdDotNet.Core.Tasks;
using ClawdDotNet.Core.Tests.Infrastructure;
using Microsoft.Extensions.Logging.Abstractions;
using Shouldly;

namespace ClawdDotNet.Core.Tests.Budget;

/// <summary>
/// #717: Large tasks and special budgets. A task too large for the day is not simply started —
/// its run proposes a split or a special budget. A special budget is a pot of its own: its
/// runs stay outside the daily limits and are visible apart.
/// </summary>
public sealed class SpecialBudgetTests : IDisposable
{
    private readonly string _directory = Path.Combine(Path.GetTempPath(), "clawd-tests", Guid.NewGuid().ToString("N"));
    private readonly SqliteStorage _storage;
    private readonly SqliteTaskRepository _tasks;
    private readonly SqliteSpecialBudgetRepository _pots;
    private readonly SqliteUsageRepository _usage;

    public SpecialBudgetTests()
    {
        _storage = new SqliteStorage(Path.Combine(_directory, "state.db"));
        _tasks = new SqliteTaskRepository(_storage);
        _pots = new SqliteSpecialBudgetRepository(_storage);
        _usage = new SqliteUsageRepository(_storage);
    }

    public void Dispose()
    {
        SqliteStorage.ReleasePooledConnections(_directory);
        try { Directory.Delete(_directory, recursive: true); }
        catch { /* cleanup is secondary */ }
    }

    private Task<TaskItem> Seed(string id, string parent = "", decimal? estimate = null, string assignee = "@worker")
        => _tasks.UpsertAsync(new TaskItem { Id = id, Title = id, Assignee = assignee, ParentId = parent, EstimateUsd = estimate }, default);

    private Task Pot(string taskId, decimal amount)
        => _pots.SetAsync(new SpecialBudget { TaskId = taskId, AmountUsd = amount, Reason = "one big job", RequestedBy = "worker", ApprovedAt = DateTime.UtcNow }, default);

    // ═══════════════════ the pot ═══════════════════

    [Fact]
    public async Task A_part_draws_on_the_pot_of_the_task_above_it()
    {
        await Seed("t-big");
        await Seed("t-part", parent: "t-big");
        await Seed("t-other");
        await Pot("t-big", 5m);

        (await SpecialBudgets.FindAsync("t-part", _tasks, _pots, default))!.TaskId.ShouldBe("t-big");
        (await SpecialBudgets.FindAsync("t-other", _tasks, _pots, default)).ShouldBeNull("outside the tree it does not exist");
    }

    [Fact]
    public async Task A_circle_in_the_tree_does_not_hang_the_lookup()
    {
        await Seed("t-a", parent: "t-b");
        await Seed("t-b", parent: "t-a");

        (await SpecialBudgets.FindAsync("t-a", _tasks, _pots, default)).ShouldBeNull();
    }

    [Fact]
    public async Task Runs_on_the_pot_count_for_the_pot_and_not_against_the_day()
    {
        var today = DateOnly.FromDateTime(DateTime.Now);
        await _usage.RecordAsync(new RunUsage { AgentId = "worker", PromptTokens = 1_000, CostUsd = 0.10m, OccurredAt = DateTime.Now }, default);
        await _usage.RecordAsync(new RunUsage { AgentId = "worker", PromptTokens = 9_000, CostUsd = 2.50m, OccurredAt = DateTime.Now, BudgetPot = "t-big" }, default);

        (await _pots.SpentAsync("t-big", default)).ShouldBe(2.50m);
        (await _usage.GetDailyForBudgetAsync(today, "worker", default)).CostUsd.ShouldBe(0.10m);
        (await _usage.GetDailyAsync(today, "worker", default)).CostUsd.ShouldBe(2.60m, "what was used is still shown in full");
        (await _usage.GetRecentForAgentAsync("worker", 10, default)).ShouldHaveSingleItem()
            .CostUsd.ShouldBe(0.10m, "a pot run would inflate the estimate of an ordinary run");
    }

    [Fact]
    public async Task The_pot_refuses_once_it_is_used_up_and_counts_runs_in_flight()
    {
        var guard = new BudgetGuard(_usage);
        await _usage.RecordAsync(new RunUsage { AgentId = "worker", PromptTokens = 1_000, CostUsd = 0.40m, OccurredAt = DateTime.Now }, default);
        var pot = new SpecialBudget { TaskId = "t-big", AmountUsd = 1.00m };

        // The agent's run averages 0.40. Used 0.50: the first run finds 0.50, the second 0.90 —
        // both below 1.00. The third finds 1.30 and is refused: runs in flight count.
        (await guard.ReservePotAsync(pot, spent: 0.50m, "worker", default)).ShouldNotBeNull();
        (await guard.ReservePotAsync(pot, spent: 0.50m, "worker", default)).ShouldNotBeNull();
        (await guard.ReservePotAsync(pot, spent: 0.50m, "worker", default)).ShouldBeNull();

        (await guard.ReservePotAsync(pot, spent: 1.00m, "worker", default) is null).ShouldBeTrue("used up is used up");
    }

    [Fact]
    public async Task A_run_in_the_tree_takes_place_although_the_day_is_used_up()
    {
        await Seed("t-big");
        await Seed("t-part", parent: "t-big");
        await Seed("t-other");
        await Pot("t-big", 5m);
        await _usage.RecordAsync(new RunUsage { AgentId = "worker", PromptTokens = 5_000, OccurredAt = DateTime.Now }, default);

        var audit = new InMemoryAuditRepository();
        var fixture = new EngineFixture(audit, usage: _usage, tasks: _tasks);
        fixture.Engine.SpecialBudgets = _pots;
        var agent = fixture.AddAgent("worker");
        agent.Budget.DailyTokens = 1_000; // used up
        fixture.Client.RespondsWithText("part done").RespondsWithText("part done").RespondsWithText("part done");

        var funded = await fixture.Engine.RunAsync(agent, "Do part 1", "inst", default, taskId: "t-part");
        var other = await fixture.Engine.RunAsync(agent, "Do something else", "inst", default, taskId: "t-other");

        funded.Status.ShouldBe(AgentRunStatus.Completed, funded.Error?.Message);
        other.Status.ShouldBe(AgentRunStatus.BudgetExceeded, "the pot does not raise the daily limit");
        audit.Receipts.First(r => r.TaskId == "t-part").BudgetPot.ShouldBe("t-big");
        audit.Receipts.First(r => r.TaskId == "t-other").BudgetPot.ShouldBeNull();
        (await _usage.GetDailyForBudgetAsync(DateOnly.FromDateTime(DateTime.Now), "worker", default)).TotalTokens
            .ShouldBe(5_000, "the pot run stayed outside the day");
    }

    [Fact]
    public async Task The_cost_analysis_shows_the_pot_apart()
    {
        var audit = new InMemoryAuditRepository();
        audit.Receipts.Enqueue(new RunReceipt { RunId = "r1", AgentId = "worker", CostUsd = 0.10m, OccurredAt = DateTime.UtcNow });
        audit.Receipts.Enqueue(new RunReceipt { RunId = "r2", AgentId = "worker", CostUsd = 2.00m, OccurredAt = DateTime.UtcNow, BudgetPot = "t-big" });
        var analytics = new Core.Analytics.UsageAnalytics(audit, null);

        var points = await analytics.SeriesAsync(
            new Core.Analytics.RunQuery { FromUtc = DateTime.UtcNow.AddDays(-1), ToUtc = DateTime.UtcNow.AddDays(1) },
            "day", "budget", TimeSpan.Zero, default);

        points.Select(p => p.Group).OrderBy(g => g).ShouldBe(["daily", "special:t-big"]);
    }

    // ═══════════════════ large tasks ═══════════════════

    private LargeTaskCheck Check(decimal dailyLimit = 1.00m)
    {
        var guard = new BudgetGuard(_usage);
        return new LargeTaskCheck(new TaskCostEstimator(new InMemoryAuditRepository(), _tasks, guard), _tasks, _pots,
            () => [new AgentConfig { AgentId = "worker", Budget = new AgentBudget { DailyCostUsd = dailyLimit } }],
            () => new FillWorkSettings());
    }

    [Fact]
    public async Task A_task_estimated_above_half_the_day_gets_a_planning_run()
    {
        var task = await Seed("t-big", estimate: 0.80m);

        var instruction = await Check().PlanningInstructionAsync(task, default);

        instruction.ShouldNotBeNull();
        instruction.ShouldContain("Do not start the work");
        instruction.ShouldContain("parent='t-big'");
        instruction.ShouldContain("request_budget");
        instruction.ShouldContain("estimated by its author, not measured");
    }

    [Theory]
    [InlineData(0.40, 1.00)]   // below half
    [InlineData(0.80, 0.00)]   // no cost limit — nothing to protect
    public async Task Otherwise_it_simply_runs(double estimate, double limit)
    {
        var task = await Seed("t-1", estimate: (decimal)estimate);

        (await Check((decimal)limit).PlanningInstructionAsync(task, default)).ShouldBeNull();
    }

    [Fact]
    public async Task The_agents_average_run_does_not_make_a_task_large()
    {
        await _usage.RecordAsync(new RunUsage { AgentId = "worker", PromptTokens = 1_000, CostUsd = 0.90m, OccurredAt = DateTime.Now }, default);
        var task = await Seed("t-1");

        (await Check().PlanningInstructionAsync(task, default)).ShouldBeNull("that says nothing about this task");
    }

    [Fact]
    public async Task A_split_or_funded_task_is_not_planned_again()
    {
        var split = await Seed("t-split", estimate: 0.80m);
        await Seed("t-split-1", parent: "t-split");
        var funded = await Seed("t-funded", estimate: 0.80m);
        await Pot("t-funded", 3m);

        (await Check().PlanningInstructionAsync(split, default)).ShouldBeNull();
        (await Check().PlanningInstructionAsync(funded, default)).ShouldBeNull();
    }

    // ═══════════════════ the scanner ═══════════════════

    private TaskScanner Scanner(PlanningDispatcher dispatcher, List<string>? notes = null)
        => new(_tasks, dispatcher, NullLoggerFactory.Instance, largeTaskCheck: Check(),
            noteOnTask: (id, note, _) => { notes?.Add($"{id}: {note}"); return Task.CompletedTask; });

    [Fact]
    public async Task The_planning_run_is_told_so_and_the_task_waits_for_its_parts()
    {
        await Seed("t-big", estimate: 0.80m);
        var notes = new List<string>();
        var dispatcher = new PlanningDispatcher(async task =>
        {
            // The agent splits it.
            await Seed("t-big-1", parent: "t-big");
            await Seed("t-big-2", parent: "t-big");
        });

        await Scanner(dispatcher, notes).ScanOnceAsync(default);

        dispatcher.Instructions.ShouldHaveSingleItem().ShouldContain("Planning run");
        var big = (await _tasks.GetAsync("t-big", default))!;
        big.Status.ShouldBe(TaskItemStatus.Blocked);
        big.LargePlannedAt.ShouldNotBeNull();
        notes.ShouldHaveSingleItem().ShouldContain("split into 2 part(s)");
    }

    [Fact]
    public async Task When_the_last_part_is_done_the_large_task_runs_again_once()
    {
        await Seed("t-big", estimate: 0.80m);
        var dispatcher = new PlanningDispatcher(async _ => await Seed("t-big-1", parent: "t-big"));
        var scanner = Scanner(dispatcher);

        await scanner.ScanOnceAsync(default);   // planning; the part is created
        await scanner.ScanOnceAsync(default);   // the part runs and is done
        (await _tasks.GetAsync("t-big", default))!.Status.ShouldBe(TaskItemStatus.Todo);

        await scanner.ScanOnceAsync(default);   // the large task brings its parts together
        (await _tasks.GetAsync("t-big", default))!.Status.ShouldBe(TaskItemStatus.Done);
        dispatcher.Instructions.Count.ShouldBe(1, "planned once, not at every run");
        dispatcher.Started.ShouldBe(["t-big", "t-big-1", "t-big"]);
    }

    [Fact]
    public async Task Without_a_split_it_waits_for_the_budget_decision_and_says_so()
    {
        await Seed("t-big", estimate: 0.80m);
        var notes = new List<string>();

        await Scanner(new PlanningDispatcher(_ => Task.CompletedTask), notes).ScanOnceAsync(default);

        (await _tasks.GetAsync("t-big", default))!.Status.ShouldBe(TaskItemStatus.Blocked);
        notes.ShouldHaveSingleItem().ShouldContain("waits for a special budget");
    }

    [Fact]
    public async Task Approved_during_planning_it_runs_on_the_pot_next()
    {
        await Seed("t-big", estimate: 0.80m);
        var dispatcher = new PlanningDispatcher(_ => Pot("t-big", 3m));
        var scanner = Scanner(dispatcher);

        await scanner.ScanOnceAsync(default);
        (await _tasks.GetAsync("t-big", default))!.Status.ShouldBe(TaskItemStatus.Todo);

        (await scanner.ScanOnceAsync(default)).ShouldBe(1, "its occurrence was given back — the planning run did not do the work");
        (await _tasks.GetAsync("t-big", default))!.Status.ShouldBe(TaskItemStatus.Done);
    }

    private sealed class PlanningDispatcher(Func<TaskItem, Task> onPlanning) : ITaskDispatcher
    {
        public List<string> Instructions { get; } = [];
        public List<string> Started { get; } = [];

        public Task<TaskLauf> DispatchAsync(TaskItem task, CancellationToken ct)
        {
            Started.Add(task.Id);
            return Task.FromResult(TaskLauf.Erledigt);
        }

        public async Task<TaskLauf> DispatchAsync(TaskItem task, string? instruction, CancellationToken ct)
        {
            Started.Add(task.Id);
            if (instruction is not null)
            {
                Instructions.Add(instruction);
                await onPlanning(task);
            }
            return TaskLauf.Erledigt;
        }
    }
}