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

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


class ForgeDatabase:
    """
    A comprehensive Supabase client for the Aegis Forge pipeline.
    Stores trajectories with all fields shown in the example, plus metadata.
    """

    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).")
        except Exception as e:
            logger.error(f"❌ Database connection failed: {e}")
            raise

        # Metadata from environment
        self.model_name = os.getenv("MODEL_NAME", "unknown")
        self.environment = os.getenv("RUN_ENV") or os.getenv("AEGIS_ENVIRONMENT", "unknown")
        self.aegis_version = os.getenv("AEGIS_VERSION", "0.1.0")

    def fetch_pending_tasks(self, table_name: str, limit: int = 10) -> list:
        """Retrieve a batch of unprocessed tasks (expects 'status' column)."""
        try:
            response = self.client.table(table_name)\
                .select("*")\
                .eq("status", "pending")\
                .limit(limit)\
                .execute()
            return response.data
        except Exception as e:
            logger.error(f"Error fetching pending tasks from {table_name}: {e}")
            return []

    def update_task_status(self, table_name: str, record_id: str,
                           status: str, error_log: str = None):
        """Update the status of a specific task."""
        try:
            data = {"status": status}
            if error_log:
                data["error_log"] = error_log
            self.client.table(table_name)\
                .update(data)\
                .eq("id", record_id)\
                .execute()
            logger.debug(f"Updated record {record_id} in {table_name} to {status}.")
        except Exception as e:
            logger.error(f"Error updating status for {record_id} in {table_name}: {e}")

    def insert_trajectory(self, table_name: str, payload: dict):
        """
        Insert a completed trajectory with the full set of fields.
        Expected payload keys (matching your example):
          - prompt_id: UUID (string) – the original task ID
          - prompt: the user prompt (string)
          - generated_code: the code output after parsing (string)
          - validation_log: the error log or success message (string)
          - success: boolean
          - exit_code: integer (from subprocess)
          - iterations: number of attempts used (integer)
          - created_at: ISO timestamp (string) – if missing, auto‑filled
          - plan_output: (string) – empty in example, but can be filled
          - phase: integer (training phase)
        Additional metadata (model_name, environment, aegis_version) are added automatically.
        """
        # Build the record with all available fields
        record = {
            "prompt_id": payload.get("prompt_id"),
            "prompt": payload.get("prompt"),
            "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"),
            "created_at": payload.get("created_at", datetime.now(timezone.utc).isoformat()),
            "plan_output": payload.get("plan_output", ""),
            "phase": payload.get("phase"),
            "model_name": self.model_name,
            "environment": self.environment,
            "aegis_version": self.aegis_version,
        }

        # Remove any None values to avoid SQL errors (columns may be nullable)
        record = {k: v for k, v in record.items() if v is not None}

        try:
            self.client.table(table_name).insert(record).execute()
            logger.debug(f"Successfully inserted trajectory into {table_name}.")
        except Exception as e:
            logger.error(f"Error inserting 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