import json
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Literal

from openai import OpenAI
from pydantic import BaseModel

IN_PATH, OUT_PATH, FAIL_PATH = "data_01/data.json", "data_02/data_processed.json", "data_02/failed.json"
MAX_TRIES = 2
WORKERS = 16
MODEL = "gpt-6-luna"


class Step(BaseModel):
    id: str
    expr: str


class Validation(BaseModel):
    input: dict[str, float]
    process: list[Step]
    output: str


# Structured Outputs(strict)는 dict 스키마를 지원하지 않아서 input을 리스트로 받은 뒤 dict로 변환
class Var(BaseModel):
    name: str
    value: float


class LLMValidation(BaseModel):
    input: list[Var]
    process: list[Step]
    output: str


class Result(BaseModel):
    status: Literal["success", "unsupported"]
    validation: LLMValidation | None


def evaluate(v: Validation) -> float:
    env = dict(v.input)
    for step in v.process:
        env[step.id] = eval(step.expr, {"__builtins__": {}}, env)
    return env[v.output]


client = OpenAI(
    api_key="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"
)
with open("system_prompt.md") as f:
    system_prompt = f.read()
with open(IN_PATH) as f:
    data = json.load(f)


def convert(item: dict) -> Validation | None:
    response = client.responses.parse(
        model=MODEL,
        reasoning={"effort": "none"},
        instructions=system_prompt,
        input=json.dumps(item, ensure_ascii=False),
        text_format=Result,
    )
    result = response.output_parsed
    if result is None or result.validation is None:
        return None

    raw = result.validation
    v = Validation(
        input={x.name: x.value for x in raw.input},
        process=raw.process,
        output=raw.output,
    )
    try:
        ok = abs(evaluate(v) - item["answer"]) < 1e-6
    except Exception:
        ok = False
    return v if ok else None


def process(item: dict) -> bool:
    for _ in range(MAX_TRIES):
        try:
            v = convert(item)
        except Exception as e:
            print(f"error: {e}")
            continue
        if v:
            item["validation"] = v.model_dump()
            return True
    return False


start = time.time()
with ThreadPoolExecutor(max_workers=WORKERS) as pool:
    results = list(pool.map(process, data))

processed = [item for item, ok in zip(data, results) if ok]
failed = [item for item, ok in zip(data, results) if not ok]
print(f"OK {len(processed)} / FAILED {len(failed)} ({time.time() - start:.1f}s)")

with open(OUT_PATH, "w") as f:
    json.dump(processed, f, indent=4, ensure_ascii=False)
with open(FAIL_PATH, "w") as f:
    json.dump(failed, f, indent=4, ensure_ascii=False)