"""Tests for the typed tools and the agent loop (no model — ScriptedClient).

Run from the project root:  python -m unittest discover -s tests
"""

import json
import unittest

from src.agent import (
    Agent,
    AssistantTurn,
    ScriptedClient,
    ToolCall,
    ToolContext,
    build_tools,
)
from src.ledger import Ledger


def seeded_ledger():
    lg = Ledger(":memory:")
    lg.record_income("2024-05-10", "Diseño", 45000, iva_rate="0.16")
    lg.record_expense("2024-05-15", "Adobe", 12000, iva_rate="0.16")
    return lg


class TestTools(unittest.TestCase):
    def setUp(self):
        self.ctx = ToolContext(seeded_ledger())
        self.tools = build_tools(self.ctx)

    def tearDown(self):
        self.ctx.ledger.close()

    def test_record_then_summary(self):
        self.tools["record_income"].handler(
            date="2024-05-25", description="Extra", amount=5000, iva_rate=0.16)
        summary = self.tools["month_summary"].handler(year=2024, month=5)
        self.assertEqual(summary["income"], "50000.00")  # 45,000 + 5,000

    def test_compute_isr_resico(self):
        r = self.tools["compute_isr_resico"].handler(year=2024, month=5)
        self.assertEqual(r["amount"], "495.00")
        self.assertIn("113-E", r["source"])  # provenance carried through

    def test_compare_regimes(self):
        r = self.tools["compare_regimes"].handler(year=2024, month=5)
        self.assertEqual(r["recommended"], "RESICO")
        self.assertEqual(r["resico_isr"], "495.00")
        self.assertEqual(r["general_isr"], "5418.90")

    def test_us_tax_estimate(self):
        r = self.tools["us_tax_estimate"].handler(gross_income=100000, expenses=40000)
        self.assertEqual(r["net_profit"]["amount"], "60000.00")
        self.assertEqual(r["self_employment_tax"]["amount"], "8477.73")

    def test_every_tool_has_schema(self):
        for name, tool in self.tools.items():
            s = tool.schema()
            self.assertEqual(s["function"]["name"], name)
            self.assertIn("parameters", s["function"])


class TestAgentLoop(unittest.TestCase):
    def setUp(self):
        self.ctx = ToolContext(seeded_ledger())
        self.tools = build_tools(self.ctx)

    def tearDown(self):
        self.ctx.ledger.close()

    def test_two_step_run(self):
        client = ScriptedClient([
            AssistantTurn(tool_calls=[ToolCall("compare_regimes", {"year": 2024, "month": 5})]),
            AssistantTurn(text="En mayo te conviene RESICO: pagas 495 vs 5,418.90."),
        ])
        agent = Agent(client, self.tools)
        trace = agent.run("¿Qué régimen me conviene en mayo?")

        self.assertEqual(len(trace.steps), 1)
        self.assertEqual(trace.steps[0].tool, "compare_regimes")
        self.assertTrue(trace.steps[0].ok)
        self.assertEqual(trace.steps[0].result["recommended"], "RESICO")
        self.assertIn("RESICO", trace.final_answer)

    def test_tool_error_is_surfaced_not_raised(self):
        # Missing required 'month' → handler raises → loop records ok=False, continues.
        client = ScriptedClient([
            AssistantTurn(tool_calls=[ToolCall("compute_isr_resico", {"year": 2024})]),
            AssistantTurn(text="Necesito el mes para calcular tu ISR."),
        ])
        agent = Agent(client, self.tools)
        trace = agent.run("¿Cuánto ISR debo?")
        self.assertFalse(trace.steps[0].ok)
        self.assertIn("error", trace.steps[0].result)
        self.assertIn("mes", trace.final_answer)

    def test_unknown_tool(self):
        client = ScriptedClient([
            AssistantTurn(tool_calls=[ToolCall("does_not_exist", {})]),
            AssistantTurn(text="done"),
        ])
        trace = Agent(client, self.tools).run("hi")
        self.assertFalse(trace.steps[0].ok)
        self.assertIn("unknown tool", trace.steps[0].result["error"])

    def test_trace_serializes(self):
        client = ScriptedClient([
            AssistantTurn(tool_calls=[ToolCall("compute_iva", {"year": 2024, "month": 5})]),
            AssistantTurn(text="Tu IVA a cargo es 5,280.00."),
        ])
        trace = Agent(client, self.tools).run("¿Cuánto IVA debo en mayo?")
        blob = trace.to_json()
        parsed = json.loads(blob)
        self.assertEqual(parsed["steps"][0]["tool"], "compute_iva")
        self.assertEqual(parsed["steps"][0]["result"]["amount"], "5280.00")

    def test_step_budget_terminates(self):
        # Client always asks for a tool call; the loop must stop at max_steps.
        loop_turn = AssistantTurn(tool_calls=[ToolCall("month_summary", {"year": 2024, "month": 5})])
        client = ScriptedClient([loop_turn] * 20)
        trace = Agent(client, self.tools, max_steps=3).run("loop forever")
        self.assertEqual(len(trace.steps), 3)
        self.assertIn("step limit", trace.final_answer)


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