import json
from data.field_tree import FieldTree, Field
from agents.llm import LLMClient

PROMPT_AGENT1 = """You are a COBOL COPYBOOK parser. Given COPYBOOK text, output a JSON object:
{"fields": [{"name": "...", "level": N, "pic": "...", "usage": "DISPLAY|COMP-3|COMP|COMP-5", "offset": N, "length": N, "decimal": N, "signed": bool, "occurs": N|null, "redefines": "..."|null, "conditions": [{"name": "...", "value": "..."}], "children": [...]}]}
Return valid JSON only. No explanation."""


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

    def parse(self, copybook_text: str) -> FieldTree:
        messages = [
            {"role": "system", "content": PROMPT_AGENT1},
            {"role": "user", "content": copybook_text}
        ]
        raw = self.llm.call(messages)
        return self._parse_response(raw)

    def _parse_response(self, raw: str) -> FieldTree:
        try:
            data = json.loads(raw)
            fields = self._to_fields(data.get("fields", []), offset=0)
            return FieldTree(fields=fields, copybook_name="")
        except (json.JSONDecodeError, KeyError):
            return FieldTree(fields=[], copybook_name="parse_error")

    def _to_fields(self, raw_fields: list, offset: int = 0) -> list[Field]:
        result = []
        current_offset = offset
        for rf in raw_fields:
            f = Field(
                name=rf.get("name", ""),
                level=rf.get("level", 0),
                pic=rf.get("pic", ""),
                usage=rf.get("usage", "DISPLAY"),
                offset=current_offset,
                length=rf.get("length", 0),
                decimal=rf.get("decimal", 0),
                signed=rf.get("signed", False),
                occurs=rf.get("occurs"),
                redefines=rf.get("redefines"),
                conditions=rf.get("conditions", []))
            children = rf.get("children", [])
            f.children = self._to_fields(children, current_offset)
            current_offset += f.length
            result.append(f)
        return result