from skydiscover.search.base_database import Program, ProgramDatabase
from skydiscover.config import DatabaseConfig
from typing import Optional, Tuple, List, Dict
import logging
from dataclasses import dataclass
import random

logger = logging.getLogger(__name__)

@dataclass
class SearchStrategy(Program):
    """Program entry for the search strategy database."""


class SearchStrategyDatabase(ProgramDatabase):
    """Database for storing and sampling evolved search strategy programs."""

    def __init__(self, name: str, config: DatabaseConfig):
        super().__init__(name, config)

        if config.random_seed is not None:
            random.seed(config.random_seed)
            logger.debug(f"Database: Set random seed to {config.random_seed}")

    def add(self, program: SearchStrategy, iteration: Optional[int] = None, **kwargs) -> str:
        """Add a program to the database."""
        self.programs[program.id] = program

        if iteration is not None:
            self.last_iteration = max(self.last_iteration, iteration)

        if self.config.db_path:
            self._save_program(program)

        self._update_best_program(program)

        logger.debug(f"Added program {program.id} to evolve database")
        return program.id

    def sample(
        self,
        num_inspirations: int = 5,
        **kwargs
    ) -> Tuple[Dict[str, SearchStrategy], Dict[str, List[SearchStrategy]]]:
        """
        Sample a search strategy to refine and inspiration strategies for evolution.
        """
        def safe_score(p):
            score = p.metrics.get("combined_score") if p.metrics else None
            if not isinstance(score, (int, float)):
                logger.warning(f"Program {p.id} has invalid combined_score: {score}, metrics: {p.metrics}")
            return float(score) if isinstance(score, (int, float)) else float('-inf')

        parent = max(self.programs.values(), key=safe_score)
        available_programs = list(self.programs.values())
        num_to_sample = max(0, min(num_inspirations, len(available_programs)))

        inspirations = random.sample(available_programs, num_to_sample) if num_to_sample > 0 else []
        inspirations = [p for p in inspirations if p.id != parent.id]
        return {"": parent}, {"": inspirations}
