from __future__ import annotations

import unittest

from docforge.python_logic import PythonLogicOwner, analyze_python_source


class PythonLogicTests(unittest.TestCase):
    def projection(self, source: str, qualified_name: str, line: int = 1):
        return analyze_python_source(
            source,
            source_id="source.example",
            owners=(PythonLogicOwner("py.symbol.example", qualified_name, line),),
            filename="example.py",
        )[0]

    def test_boolean_short_circuit_branches_and_terminals(self) -> None:
        projection = self.projection(
            """\
def decide(enabled, cached, stale):
    if enabled and (cached is None or stale):
        return "fetch"
    raise RuntimeError("disabled")
""",
            "decide",
        )
        kinds = [node.kind for node in projection.nodes]
        labels = [node.label for node in projection.nodes]
        edge_labels = [edge.label for edge in projection.edges]

        self.assertEqual(1, kinds.count("entry"))
        self.assertEqual(1, kinds.count("exit"))
        self.assertEqual(3, kinds.count("condition"))
        self.assertIn("enabled", labels)
        self.assertIn("cached is None", labels)
        self.assertIn("stale", labels)
        self.assertIn("TRUE", edge_labels)
        self.assertIn("FALSE", edge_labels)
        self.assertIn("RETURN", edge_labels)
        self.assertIn("RAISE", edge_labels)

    def test_loops_match_try_and_control_transfers_are_explicit(self) -> None:
        projection = self.projection(
            """\
def process(items, mode):
    for item in items:
        if item.skip:
            continue
        if item.stop:
            break
        consume(item)
    else:
        finish()
    match mode:
        case "safe":
            value = safe()
        case _:
            value = fallback()
    try:
        return value
    except ValueError:
        raise
    finally:
        cleanup()
""",
            "process",
        )
        kinds = {node.kind for node in projection.nodes}
        relations = {edge.relation for edge in projection.edges}
        labels = {edge.label for edge in projection.edges}

        self.assertTrue(
            {"loop", "continue", "break", "case", "try", "except", "finally"}.issubset(kinds)
        )
        self.assertTrue(
            {"loop", "continue", "break", "case", "exception", "return", "raise"}.issubset(
                relations
            )
        )
        self.assertIn("EXHAUSTED", labels)
        self.assertIn("NEXT ITEM", labels)
        self.assertIn("NEXT CASE", labels)

    def test_class_methods_nested_functions_and_async_functions_use_explicit_owners(self) -> None:
        source = """\
class Worker:
    async def run(self):
        async with self.session():
            await self.step()

        if self.enabled:
            def nested():
                return True

        return nested()
"""
        projections = analyze_python_source(
            source,
            source_id="source.worker",
            owners=(
                PythonLogicOwner("py.symbol.worker.run", "Worker.run", 2),
                PythonLogicOwner("py.symbol.worker.nested", "Worker.run.nested", 7),
            ),
            filename="worker.py",
        )

        self.assertEqual(
            ("py.symbol.worker.nested", "py.symbol.worker.run"),
            tuple(projection.owner_node_id for projection in projections),
        )
        run = next(item for item in projections if item.owner_node_id.endswith(".run"))
        self.assertIn("action", {node.kind for node in run.nodes})
        self.assertIn("call", {node.kind for node in run.nodes})

    def test_projection_is_deterministic(self) -> None:
        source = """\
def choose(first, second):
    return first if first else second
"""
        first = self.projection(source, "choose")
        second = self.projection(source, "choose")
        self.assertEqual(first, second)


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