import json
import os
import time
import logging
import threading
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional, Dict, Any, Callable, List
from concurrent.futures import ThreadPoolExecutor, as_completed
from openai import OpenAI

# Attempt to import ForgeDatabase safely for cloud mode
try:
    from aegis_forge.cloud_db import ForgeDatabase
except ImportError:
    ForgeDatabase = None

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class UniversalAegisForge(ABC):
    """
    The Base Class for Execution-Guided Synthesis.
    Handles the LLM, the retry loop, context management, Cloud DB sync, and Parallel Execution.
    """

    def __init__(
        self,
        system_prompt: str,
        output_file: Optional[str] = None,
        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,
        db_config: Optional[Dict[str, str]] = None,
        max_workers: Optional[int] = None
    ):
        self.system_prompt = system_prompt
        self.output_file = Path(output_file) if output_file else None

        # Thread safety locks
        self.file_lock = threading.Lock()
        self.db_lock = threading.Lock()

        # Database Initialization
        self.db_config = db_config
        self.db = None
        self.read_table = None
        self.write_table = None
        if self.db_config:
            if ForgeDatabase is None:
                raise ImportError("cloud_db.py module not found. Cannot initialize cloud mode.")
            logger.info("📡 Cloud Configuration detected. Initializing ForgeDatabase...")
            self.db = ForgeDatabase()

            # Require write_table; read_table is optional (used only in cloud mode)
            if "write_table" not in self.db_config:
                raise ValueError("db_config must contain 'write_table' key.")
            self.write_table = self.db_config["write_table"]
            self.read_table = self.db_config.get("read_table")  # may be None

        # Concurrency & Retry setup
        self.max_workers = max_workers or int(os.getenv("AEGIS_WORKERS", "3"))
        self.max_retries = max_retries if max_retries is not None else int(os.getenv("AEGIS_MAX_RETRIES", "10"))
        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

        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:
            # Default OpenAI client with timeout and minimal internal retries
            self.client = OpenAI(
                api_key=os.getenv("OPENAI_API_KEY", "dummy"),
                base_url=base_url,
                timeout=6000.0,          # seconds per request
                max_retries=0           # disable client‑side retries
            )
            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:
        # Use timeout from kwargs if provided, otherwise fallback to 120s
        timeout = kwargs.get("timeout", 6000.0)
        call_kwargs = {
            "model": self.model,
            "messages": messages,
            "timeout": timeout,
            **kwargs
        }
        response = self.client.chat.completions.create(**call_kwargs)
        return response.choices[0].message.content

    @abstractmethod
    def generate_prompt(self) -> str:
        """Override this: Generate a random combinatorial problem."""
        pass

    @abstractmethod
    def parse_output(self, raw_llm_response: str) -> str:
        """Override this: Extract the code/JSON from the raw LLM response."""
        pass

    @abstractmethod
    def execute_critic(self, parsed_output: str) -> dict:
        """
        Override this: The Deterministic Validator.
        Must return: {"success": bool, "feedback": str}
        """
        pass

    def run_synthesis_loop(self, user_prompt: str, task_id: Optional[str] = None) -> bool:
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        final_error = "Unknown Error"
        last_exception = None
        parsed_code = None
        exit_code = None
        plan_output = None

        for iteration in range(1, self.max_retries + 1):
            try:
                raw_output = self.llm_call(messages, self.generation_kwargs)
                last_exception = None
            except Exception as e:
                logger.warning(f"[Task {task_id or 'Local'}] ⚠️ LLM call failed on iteration {iteration}: {e}")
                last_exception = e
                time.sleep(self.retry_delay)
                continue

            parsed_code = self.parse_output(raw_output)

            try:
                critic_result = self.execute_critic(parsed_code)
            except Exception as e:
                logger.error(f"[Task {task_id or 'Local'}] ⚠️ Critic execution crashed unexpectedly: {e}")
                critic_result = {"success": False, "feedback": f"Critic System Crash: {str(e)}"}

            exit_code = critic_result.get("exit_code")
            plan_output = critic_result.get("plan_output")

            if critic_result.get("success", False):
                logger.info(f"[Task {task_id or 'Local'}] ✅ SUCCESS on iteration {iteration}.")
                trajectory = messages + [{"role": "assistant", "content": raw_output}]

                self._save_trajectory(
                    trajectory_messages=trajectory,
                    task_id=task_id,
                    is_success=True,
                    user_prompt=user_prompt,
                    generated_code=parsed_code,
                    exit_code=exit_code,
                    iterations=iteration,
                    plan_output=plan_output,
                    phase=self.phase,
                    reasoning=""
                )

                with self.file_lock:
                    self.success_count += 1
                return True
            else:
                logger.warning(f"[Task {task_id or 'Local'}] ❌ FAILED on iteration {iteration}. Retrying...")
                final_error = critic_result.get('feedback', 'Unknown error')
                messages.append({"role": "assistant", "content": raw_output})
                messages.append({
                    "role": "user",
                    "content": f"Execution failed with error:\n{final_error}\nFix the code."
                })

        # After all retries
        logger.error(f"[Task {task_id or 'Local'}] 💀 ABORTED after {self.max_retries} retries.")
        if last_exception:
            final_error = f"LLM Call Failed repeatedly. Last exception: {str(last_exception)}"

        self._save_trajectory(
            trajectory_messages=messages,
            task_id=task_id,
            is_success=False,
            error_log=final_error,
            user_prompt=user_prompt,
            generated_code=parsed_code,
            exit_code=exit_code,
            iterations=self.max_retries,
            plan_output=plan_output,
            phase=self.phase,
            reasoning="" 
        )
        return False

    def _save_trajectory(
        self,
        trajectory_messages: list,
        task_id: Optional[str],
        is_success: bool,
        error_log: str = None,
        user_prompt: str = None,
        generated_code: str = None,
        exit_code: int = None,
        iterations: int = None,
        plan_output: str = None,
        phase: int = None,
        reasoning: str = ""   # new parameter
    ):
        """Routes the saved data to Cloud DB or a Local File. Thread-safe."""
        # For local file backup, keep the conversation format
        record = self.trajectory_formatter(trajectory_messages)

        # 1. Cloud Save (if configured)
        if self.db and self.db_config and self.write_table:
            # Build payload matching the old trajectories table
            db_payload = {
                "id": task_id,
                "prompt": user_prompt,
                "reasoning": reasoning,
                "generated_code": generated_code,
                "validation_log": error_log if error_log else ("Validation and plan passed." if is_success else ""),
                "success": is_success,
                "exit_code": exit_code,
                "iterations": iterations,
                "plan_output": plan_output,
                "phase": phase,
            }
            # Remove None values to avoid SQL errors for nullable columns
            db_payload = {k: v for k, v in db_payload.items() if v is not None}

            with self.db_lock:
                self.db.insert_trajectory(self.write_table, db_payload)

        # 2. Local File Save (unchanged)
        if self.output_file:
            with self.file_lock:
                with open(self.output_file, "a") as f:
                    f.write(json.dumps(record) + "\n")

    @staticmethod
    def _default_trajectory(messages):
        """
        Default formatter: converts messages to a list of dicts with 'from' and 'value'.
        This is the format expected by many fine‑tuning libraries (e.g., ShareGPT).
        """
        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, batch_size: int = 10, run_mode: str = "local"):
        """
        Ignites the Forge in Parallel using self.max_workers.
        run_mode = "local" (uses generate_prompt()) or "cloud" (pulls from read_table).
        """
        logger.info(f"🚀 Igniting Forge in {run_mode.upper()} mode with {self.max_workers} PARALLEL WORKERS.")

        if run_mode == "cloud":
            if not self.db:
                raise ValueError("Cannot run in cloud mode without passing db_config during initialization.")
            if not self.read_table:
                raise ValueError("Cloud mode requires 'read_table' in db_config.")

            read_table = self.read_table
            tasks = self.db.fetch_pending_tasks(read_table, limit=batch_size)

            if not tasks:
                logger.info("📭 No pending tasks found in the database. Exiting.")
                return

            logger.info(f"📥 Pulled {len(tasks)} tasks from cloud. Starting parallel processing...")

            # Mark all fetched tasks as processing safely
            with self.db_lock:
                for task in tasks:
                    self.db.update_task_status(read_table, task.get("id"), "processing")

            # Execute Cloud Tasks in Parallel with strict crash recovery mapping
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = {
                    executor.submit(self.run_synthesis_loop,
                                    task.get("prompt", "MISSING PROMPT"),
                                    task.get("id")): task.get("id")
                    for task in tasks
                }

                for future in as_completed(futures):
                    task_id = futures[future]
                    try:
                        future.result()  # Raises exception if the thread crashed
                    except Exception as e:
                        logger.error(f"⚠️ Thread crashed unexpectedly for Task {task_id}: {e}")
                        # Stranded Task Recovery: Revert stuck "processing" task to "failed"
                        with self.db_lock:
                            self.db.update_task_status(read_table, task_id,
                                                        "failed", f"Thread Fatal Crash: {str(e)}")

        else:  # local mode
            logger.info(f"Generating {batch_size} local tasks...")
            local_prompts = [self.generate_prompt() for _ in range(batch_size)]

            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = {
                    executor.submit(self.run_synthesis_loop, prompt, f"LOCAL-{i}"): f"LOCAL-{i}"
                    for i, prompt in enumerate(local_prompts)
                }

                for future in as_completed(futures):
                    task_id = futures[future]
                    try:
                        future.result()
                    except Exception as e:
                        logger.error(f"⚠️ Local Thread {task_id} crashed unexpectedly: {e}")

        logger.info(f"\n🏁 Complete. Yield: {self.success_count} Gold Trajectories.")