"""Contracts of the agent-reviewed strategy: without score files it holds
exactly what its base strategy holds, a scores file only steers sessions it
was stamped before, a vetoed holding is replaced the same session, and
submit rejects sources dated after the submission."""

import json
import os
import tempfile
import unittest

import numpy as np
import pandas as pd

from agent_review import validate
from strategies.agent_review import AgentReviewStrategy, session_deadline
from strategies.robust_momentum import RobustMomentumStrategy
from tests.test_robust_strategies import _bars


def _write(score_dir: str, as_of, created_at, scores: dict) -> None:
    rows = {t: {"score": s, "veto": s == "veto", "summary": "x"} for t, s in scores.items()}
    for v in rows.values():
        if v["veto"]:
            v["score"] = 0
    path = os.path.join(score_dir, f"scores_{pd.Timestamp(as_of).date()}_{len(os.listdir(score_dir)):06d}.json")
    with open(path, "w") as fh:
        json.dump({"as_of": str(pd.Timestamp(as_of).date()), "created_at": pd.Timestamp(created_at).isoformat(),
                   "scores": rows}, fh)


class AgentReviewTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.dir = self.tmp.name
        self.bars = _bars(n_days=420, n_assets=40)
        self.kw = dict(top_n=5, pool=20, rebal_freq=5, buffer=0.5, warmup=260)

    def tearDown(self):
        self.tmp.cleanup()

    def _agent(self):
        return AgentReviewStrategy(score_dir=self.dir, **self.kw).generate_signals(self.bars)

    def test_without_score_files_it_holds_what_the_base_strategy_holds(self):
        base = RobustMomentumStrategy("overnight_recovery", top_n=5, rebal_freq=5, buffer=0.5,
                                      warmup=260).generate_signals(self.bars)
        pd.testing.assert_frame_equal(self._agent(), base, atol=1e-12)

    def test_scores_stamped_after_the_session_close_do_not_steer_it(self):
        dates = self.bars["close"].index
        k = 260 + 5 * 20                       # a rebalance row; its weights trade at dates[k + 1]
        before = self._agent()
        worst = before.columns[before.iloc[k + 1] == 0][0]
        late = session_deadline(dates[k + 1]) + pd.Timedelta(minutes=1)
        _write(self.dir, dates[k], late, {worst: 2, **{t: -2 for t in before.columns[before.iloc[k + 1] > 0]}})
        pd.testing.assert_frame_equal(self._agent(), before, atol=1e-12)
        _write(self.dir, dates[k], session_deadline(dates[k + 1]) - pd.Timedelta(hours=1),
               {t: -2 for t in before.columns[before.iloc[k + 1] > 0]})
        after = self._agent()
        pd.testing.assert_frame_equal(after.iloc[:k + 1], before.iloc[:k + 1], atol=1e-12)
        self.assertNotEqual(set(after.columns[after.iloc[k + 1] > 0]), set(before.columns[before.iloc[k + 1] > 0]))

    def test_a_vetoed_holding_is_replaced_between_rebalances(self):
        dates = self.bars["close"].index
        k = 260 + 5 * 20 + 2                   # not a rebalance row
        before = self._agent()
        held = list(before.columns[before.iloc[k + 1] > 0])
        _write(self.dir, dates[k], session_deadline(dates[k + 1]) - pd.Timedelta(hours=2), {held[0]: "veto"})
        after = self._agent()
        now = set(after.columns[after.iloc[k + 1] > 0])
        self.assertNotIn(held[0], now)
        self.assertEqual(len(now), len(held))
        self.assertAlmostEqual(float(after.iloc[k + 1].sum()), 1.0, places=9)
        self.assertTrue(np.allclose(after.iloc[:k + 1], before.iloc[:k + 1]))


class SubmitValidationTests(unittest.TestCase):
    def _check(self, *dates):
        pack = {"candidates": [{"ticker": "AAA"}]}
        sub = {"scores": {"AAA": {"score": 0, "summary": "x",
                                  "sources": [{"url": "u", "date": d} for d in dates]}}}
        return validate(sub, pack, pd.Timestamp("2026-10-12 04:00", tz="UTC"))

    def test_sources_must_be_dated_no_later_than_the_submission_day(self):
        self._check("2026-10-09", "2026-10-12", "2026-09")
        for bad in ("2026-10-13", "2026-10-12T23:00:00-04:00", None, "next week"):
            with self.subTest(date=bad), self.assertRaises(ValueError):
                self._check("2026-10-09", bad)


if __name__ == "__main__":
    unittest.main()