import os
import logging
from supabase import create_client, Client
from dotenv import load_dotenv

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

class ForgeDatabase:
    """
    Supabase client for the Aegis Forge pipeline.
    Strictly acts as a Data Lake for storing completed trajectories.
    (Prompt queue management is handled locally on the worker pod).
    """

    def __init__(self):
        load_dotenv()
        self.url = os.getenv("SUPABASE_URL")
        self.key = os.getenv("SUPABASE_KEY")
        if not self.url or not self.key:
            raise ValueError("❌ SUPABASE_URL and SUPABASE_KEY must be set in .env")

        try:
            self.client: Client = create_client(self.url, self.key)
            logger.info("✅ Successfully connected to Forge Database (Supabase Data Lake).")
        except Exception as e:
            logger.error(f"❌ Database connection failed: {e}")
            raise

    def insert_trajectory(self, table_name: str, payload: dict):
        """
        Upserts a completed trajectory into the 'trajectories' table.
        Uses upsert to gracefully handle race conditions or script restarts.
        """
        record = {
            "id": payload.get("id"),
            "prompt": payload.get("prompt"),
            "reasoning": payload.get("reasoning", ""),
            "generated_code": payload.get("generated_code"),
            "validation_log": payload.get("validation_log"),
            "success": payload.get("success"),
            "exit_code": payload.get("exit_code"),
            "iterations": payload.get("iterations"),
            "plan_output": payload.get("plan_output", ""),
            "phase": payload.get("phase"),
        }

        record = {k: v for k, v in record.items() if v is not None}

        try:
            self.client.table(table_name).upsert(record).execute()
            logger.debug(f"Successfully upserted trajectory into {table_name}.")
        except Exception as e:
            logger.error(f"Error upserting trajectory into {table_name}: {e}")

    def export_gold_dataset(self, table_name: str, batch_size: int = 1000) -> list:
        """Retrieve all successful trajectories (success = true) with pagination."""
        all_records = []
        start = 0
        try:
            while True:
                end = start + batch_size - 1
                response = self.client.table(table_name)\
                    .select("*")\
                    .eq("success", True)\
                    .range(start, end)\
                    .execute()
                records = response.data
                if not records:
                    break
                all_records.extend(records)
                logger.info(f"📥 Fetched {len(records)} records from {table_name} "
                            f"(Total so far: {len(all_records)})...")
                if len(records) < batch_size:
                    break
                start += batch_size
            logger.info(f"✅ Successfully exported all {len(all_records)} gold trajectories.")
            return all_records
        except Exception as e:
            logger.error(f"Error exporting dataset from {table_name}: {e}")
            return all_records