import json
import os
from pathlib import Path
from openai import OpenAI

class UniversalAegisForge:
    """
    The Base Class for Execution-Guided Synthesis.
    Handles the LLM, the retry loop, context management, and saving data.
    """
    def __init__(self, system_prompt: str, output_file: str, max_retries: int = 4, base_url: str = None):
        self.system_prompt = system_prompt
        self.output_file = Path(output_file)
        self.max_retries = max_retries
        self.client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY", "dummy"),
            base_url=base_url
        )
        self.model = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
        self.success_count = 0

    def generate_prompt(self) -> str:
        raise NotImplementedError("You must write a prompt generator!")

    def parse_output(self, raw_llm_response: str) -> str:
        raise NotImplementedError("You must write a parser!")

    def execute_critic(self, parsed_output: str) -> dict:
        raise NotImplementedError("You must write the Critic logic!")

    def run_synthesis_loop(self, user_prompt: str):
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        for iteration in range(1, self.max_retries + 1):
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.2
            )
            raw_output = response.choices[0].message.content
            parsed_code = self.parse_output(raw_output)
            critic_result = self.execute_critic(parsed_code)
            
            if critic_result["success"]:
                print(f"  ✅ SUCCESS on iteration {iteration}.")
                trajectory = messages + [{"role": "assistant", "content": raw_output}]
                self._save_trajectory(trajectory)
                self.success_count += 1
                return True
            else:
                print(f"  ❌ FAILED on iteration {iteration}. Retrying...")
                messages.append({"role": "assistant", "content": raw_output})
                messages.append({
                    "role": "user", 
                    "content": f"Execution failed with error:\n{critic_result['feedback']}\nFix the code."
                })
                
        print(f"  💀 ABORTED after {self.max_retries} retries.")
        return False

    def _save_trajectory(self, trajectory_messages: list):
        record = {
            "conversations": [{"from": "system" if m["role"] == "system" else "human" if m["role"] == "user" else "gpt", "value": m["content"]} for m in trajectory_messages]
        }
        with open(self.output_file, "a") as f:
            f.write(json.dumps(record) + "\n")

    def start_factory(self, target_runs: int = 100):
        print(f"🚀 Igniting Forge. Target: {target_runs} runs.")
        for i in range(target_runs):
            print(f"\n--- Task {i+1}/{target_runs} ---")
            prompt = self.generate_prompt()
            self.run_synthesis_loop(prompt)
        print(f"\n🏁 Complete. Yield: {self.success_count}/{target_runs} Gold Trajectories.")