import json
import os
import time
from pathlib import Path
from typing import Optional, Dict, Any, Callable, List
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: Optional[int] = None,
        base_url: Optional[str] = None,
        generation_kwargs: Optional[Dict[str, Any]] = None,
        llm_call: Optional[Callable[[List[Dict[str, str]], Dict[str, Any]], str]] = None,
        trajectory_formatter: Optional[Callable[[List[Dict[str, str]]], Dict]] = None
    ):
        self.system_prompt = system_prompt
        self.output_file = Path(output_file)

        # Max retries: use explicit value if given, else env var, else 4
        if max_retries is not None:
            self.max_retries = max_retries
        else:
            self.max_retries = int(os.getenv("AEGIS_MAX_RETRIES", "4"))

        self.success_count = 0

        # Build generation kwargs with environment-aware defaults
        default_temp = float(os.getenv("AEGIS_TEMPERATURE", "0.2"))
        self.generation_kwargs = generation_kwargs or {}
        if "temperature" not in self.generation_kwargs:
            self.generation_kwargs["temperature"] = default_temp

        # Retry delay after LLM failure
        self.retry_delay = float(os.getenv("AEGIS_RETRY_DELAY", "2"))

        self.trajectory_formatter = trajectory_formatter or self._default_trajectory

        # Set up LLM call mechanism
        if llm_call is not None:
            self.llm_call = llm_call
        else:
            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.llm_call = self._openai_call

    def _openai_call(self, messages: List[Dict[str, str]], kwargs: Dict[str, Any]) -> str:
        """Default LLM call using OpenAI client."""
        call_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }
        response = self.client.chat.completions.create(**call_kwargs)
        return response.choices[0].message.content

    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) -> bool:
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        for iteration in range(1, self.max_retries + 1):
            try:
                raw_output = self.llm_call(messages, self.generation_kwargs)
            except Exception as e:
                print(f"  ⚠️ LLM call failed on iteration {iteration}: {e}")
                time.sleep(self.retry_delay)
                continue

            parsed_code = self.parse_output(raw_output)
            critic_result = self.execute_critic(parsed_code)

            if critic_result.get("success", False):
                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.get('feedback', 'Unknown error')}\nFix the code."
                })

        print(f"  💀 ABORTED after {self.max_retries} retries.")
        return False

    def _save_trajectory(self, trajectory_messages: list):
        record = self.trajectory_formatter(trajectory_messages)
        with open(self.output_file, "a") as f:
            f.write(json.dumps(record) + "\n")

    @staticmethod
    def _default_trajectory(messages):
        return {
            "conversations": [
                {
                    "from": "system" if m["role"] == "system" else "human" if m["role"] == "user" else "gpt",
                    "value": m["content"]
                }
                for m in messages
            ]
        }

    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.")