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

PROMPT_AGENT2 = """You are a COBOL test data designer. Given a FieldTree JSON, generate test data covering boundary values.
Output: {"test_cases": [{"id": "TC-001", "fields": {"FIELD_NAME": value, ...}, "coverage_targets": ["DP-001-TRUE"]}]}
For each field, generate 1-3 test cases covering: zero, boundary (MAX), typical value. Return valid JSON only."""


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

    def design(self, tree: FieldTree, coverage_target: str = "boundary",
               spark_mode: bool = False) -> TestSuite:
        tree_json = {"fields": [{
            "name": f.name, "level": f.level, "pic": f.pic,
            "usage": f.usage, "length": f.length, "decimal": f.decimal,
            "signed": f.signed, "redefines": f.redefines, "occurs": f.occurs
        } for f in tree.flatten().values()]}

        messages = [
            {"role": "system", "content": PROMPT_AGENT2},
            {"role": "user", "content": json.dumps(tree_json)}
        ]
        raw = self.llm.call(messages)
        test_cases = self._parse(raw)
        suite = TestSuite(test_cases=test_cases)
        if spark_mode:
            suite.spark_config = SparkConfig(num_records=1000)
        return suite

    def _parse(self, raw: str) -> list[TestCase]:
        try:
            data = json.loads(raw)
            return [TestCase(**tc) for tc in data.get("test_cases", [])]
        except (json.JSONDecodeError, KeyError):
            return [TestCase(id="TC-FALLBACK", fields={"BR-AMT": 0})]