import json
from data.field_tree import FieldTree
from data.test_case import TestCase, TestSuite, SparkConfig
from agents.llm import LLMClient

P2 = "You are a COBOL test data designer. Given a FieldTree, generate boundary test cases. Output: {\"test_cases\":[{\"id\":\"TC-001\",\"fields\":{\"FIELD\":value},\"coverage_targets\":[\"DP-001\"]}]} JSON only."


class Agent2Data:
    def __init__(self, llm: LLMClient):
        self.llm = llm

    def design(self, tree: FieldTree, target="boundary", spark_mode=False) -> TestSuite:
        tree_d = {"fields": [{"name": f.name, "pic": f.pic, "usage": f.usage, "length": f.length,
                              "decimal": f.decimal, "signed": f.signed} for f in tree.flatten().values()]}
        r = self.llm.call([{"role": "system", "content": P2}, {"role": "user", "content": json.dumps(tree_d)}])
        try:
            tcs = [TestCase(**tc) for tc in json.loads(r).get("test_cases", [])]
        except:
            tcs = [TestCase(id="TC-FALLBACK", fields={"BR-AMT": 0})]
        s = TestSuite(test_cases=tcs)
        if spark_mode:
            s.spark_config = SparkConfig(num_records=1000)
        return s