"""
Clean interfaces for ProgramDatabase classes.

This module defines abstract base classes that capture the essential operations
needed for program database operations.
"""

from abc import ABC, abstractmethod
from dataclasses import asdict, dataclass, field, fields
from typing import Any, Dict, List, Optional, Tuple
import logging
from skydiscover.config import DatabaseConfig
from skydiscover.utils.metrics_utils import get_fitness_score

import os
import time

logger = logging.getLogger(__name__)


@dataclass
class Program:
    """Represents a program in the database"""

    # Program identification
    id: str
    solution: str
    language: str = "python"

    # Performance 
    metrics: Dict[str, Any] = field(default_factory=dict)

    # Tracking information
    iteration_found: int = 0  
    parent_id: Optional[str] = None # Parent program ID it mutates from 
    other_context_ids: Optional[List[str]] = None # other program IDs to provide as context to generate this program
    parent_info: Optional[Tuple[str, str]] = None  # information about the parent program
    context_info: Optional[List[Tuple[str, str]]] = None  # information about the context programs

    timestamp: float = field(default_factory=time.time)

    # Metadata
    metadata: Dict[str, Any] = field(default_factory=dict)
    artifacts: Dict[str, Any] = field(default_factory=dict)

    # Prompts
    prompts: Optional[Dict[str, Any]] = None
    generation: int = 0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation"""
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Program":
        """Create from dictionary representation"""
        # Get the valid field names for the Program dataclass
        valid_fields = {f.name for f in fields(cls)}

        # Filter the data to only include valid fields
        filtered_data = {k: v for k, v in data.items() if k in valid_fields}

        # Log if we're filtering out any fields
        if len(filtered_data) != len(data):
            filtered_out = set(data.keys()) - set(filtered_data.keys())
            logger.debug(f"Filtered out unsupported fields when loading Program: {filtered_out}")

        return cls(**filtered_data)


class ProgramDatabase(ABC):
    """
    Abstract base class for program storage and sampling.

    This interface captures the essential operations needed for any discovery process:
    - Add a program to the database
    - Sample a program and inspiration programs to learn from past experiences for the next discovery step
    """
    
    def __init__(self, name: str, config: DatabaseConfig, **kwargs: Any):
        self.name = name
        self.config = config

        # In-memory program storage
        self.programs: Dict[str, Program] = {}

        # Track the last iteration number (for resuming)
        self.last_iteration: int = 0

        # Optionally track initial program info (set by controller on first add)
        self.initial_program_id: Optional[str] = None
        self.initial_program_score: Optional[float] = None

        # Best program tracking
        self.best_program_id: Optional[str] = None

        # Prompt log
        self.prompts_by_program: Dict[str, Dict[str, Dict[str, str]]] = None

        # Initialize checkpoint manager (imported here to avoid circular imports)
        from skydiscover.search.utils.checkpoint_manager import CheckpointManager
        self.checkpoint_manager = CheckpointManager(self.config)

        # Load database from disk if path is provided
        if config.db_path and os.path.exists(config.db_path):
            self.load(config.db_path)

    ########################## MAIN METHODS TO IMPLEMENT ##########################
    @abstractmethod
    def add(self, program: Program, iteration: Optional[int] = None, **kwargs: Any) -> str:
        """
        Add a program to the database.

        Args:
            program: Program to add
            iteration: Current iteration (for tracking)

        Returns:
            Program ID
        """
        raise NotImplementedError("Subclasses must implement this method")

    @abstractmethod
    def sample(
        self, 
        num_inspirations: Optional[int] = 5,
        **kwargs: Any
    ) -> Tuple[Dict[str, Program], Dict[str, List[Program]]]:
        """
        Sample a parent program and inspiration programs for solution discovery.

        Args:
            num_inspirations: Number of inspiration programs to sample (defaults to 5)
            **kwargs: Search-specific parameters

        Returns:
            Tuple of (parent_program_dict, inspiration_programs_dict)
            - parent_program_dict: Dict with label as key and Program as value (e.g., {"Parent": program})
            - inspiration_programs_dict: Dict with labels as keys and List[Program] as values
        """
        raise NotImplementedError("Subclasses must implement this method")

    ########################## SAVE AND LOAD ##########################
    def save(self, path: Optional[str] = None, iteration: int = 0) -> None:
        """
        Save the database to disk

        Args:
            path: Path to save to (uses config.db_path if None)
            iteration: Current iteration number
        """
        self.checkpoint_manager.save(
            programs=self.programs,
            prompts_by_program=self.prompts_by_program,
            best_program_id=self.best_program_id,
            last_iteration=iteration if iteration is not None else self.last_iteration,
            path=path,
        )

    def load(self, path: str) -> None:
        """
        Load the database from disk

        Args:
            path: Path to load from
        """
        programs, best_id, last_iter = self.checkpoint_manager.load(path)
        self.programs = programs
        self.best_program_id = best_id
        self.last_iteration = last_iter

        self.log_status()

    def _save_program(
        self,
        program: Program,
        base_path: Optional[str] = None,
        prompts: Optional[Dict[str, Dict[str, str]]] = None,
    ) -> None:
        """
        Save a single program to disk.
        
        This is a convenience method that delegates to CheckpointManager.
        Subclasses should use this method when they need to save individual programs
        (e.g., during add() operations).
        
        Args:
            program: Program to save
            base_path: Base path to save to (uses config.db_path if None)
            prompts: Optional prompts to save with the program
        """
        self.checkpoint_manager._save_program(program, base_path, prompts)
        
    ########################## BEST PROGRAM TRACKING ##########################

    def _is_better(self, program1: Program, program2: Program) -> bool:
        """Determine if program1 has better fitness than program2."""
        if not program1.metrics and not program2.metrics:
            return program1.timestamp > program2.timestamp
        if program1.metrics and not program2.metrics:
            return True
        if not program1.metrics and program2.metrics:
            return False
        return get_fitness_score(program1.metrics) > get_fitness_score(program2.metrics)

    def _update_best_program(self, program: Program) -> None:
        """Update the best program tracking after a new program is added."""
        if self.best_program_id is None:
            self.best_program_id = program.id
            logger.debug(f"Set initial best program to {program.id}")
            return
        # If the best program is not in the database, set it to the new program
        if self.best_program_id not in self.programs:
            self.best_program_id = program.id
            return

        current_best = self.programs[self.best_program_id]

        # If the new program is better than the current best, set it to the new program
        if self._is_better(program, current_best):
            self.best_program_id = program.id

    def get_best_program(self, metric: Optional[str] = None) -> Optional[Program]:
        """Get the best program, optionally by a specific metric."""
        if not self.programs:
            return None

        if metric is None and self.best_program_id:
            if self.best_program_id in self.programs:
                return self.programs[self.best_program_id]
            else:
                logger.warning(
                    f"Tracked best program {self.best_program_id} no longer exists, will recalculate"
                )
                self.best_program_id = None

        if metric:
            sorted_programs = sorted(
                [p for p in self.programs.values() if metric in p.metrics],
                key=lambda p: p.metrics[metric],
                reverse=True,
            )
        else:
            sorted_programs = sorted(
                self.programs.values(),
                key=lambda p: get_fitness_score(p.metrics),
                reverse=True,
            )

        if sorted_programs and (
            self.best_program_id is None or sorted_programs[0].id != self.best_program_id
        ):
            self.best_program_id = sorted_programs[0].id

        return sorted_programs[0] if sorted_programs else None

    def get_top_programs(self, n: int = 10, metric: Optional[str] = None) -> List[Program]:
        """Get the top N programs, optionally by a specific metric."""
        if not self.programs:
            return []

        if metric:
            sorted_programs = sorted(
                [p for p in self.programs.values() if metric in p.metrics],
                key=lambda p: p.metrics[metric],
                reverse=True,
            )
        else:
            sorted_programs = sorted(
                self.programs.values(),
                key=lambda p: get_fitness_score(p.metrics),
                reverse=True,
            )

        return sorted_programs[:n]

    def get(self, program_id: str) -> Optional[Program]:
        """ Get a program by ID"""
        return self.programs.get(program_id)

    ########################## PROMPT LOGGING ##########################
    def log_prompt(
        self,
        program_id: str,
        template_key: str,
        prompt: Dict[str, str],
        responses: Optional[List[str]] = None,
    ) -> None:
        """
        Log a prompt for a program.
        Only logs if self.config.log_prompts is True.

        Args:
        program_id: ID of the program to log the prompt for
        template_key: Key for the prompt template
        prompt: Prompts in the format {template_key: { 'system': str, 'user': str }}.
        responses: Optional list of responses to the prompt, if available.
        """

        if not self.config.log_prompts:
            return

        if responses is None:
            responses = []
        prompt["responses"] = responses

        if self.prompts_by_program is None:
            self.prompts_by_program = {}

        if program_id not in self.prompts_by_program:
            self.prompts_by_program[program_id] = {}
        self.prompts_by_program[program_id][template_key] = prompt
    
    def log_status(self) -> None:
        """Log the status of the database"""
        best_program = self.get_best_program()
        best_score = None
        if best_program and best_program.metrics:
            best_score = best_program.metrics.get("combined_score")
        
        score_str = f"{best_score:.6f}" if best_score is not None else "N/A"
        logger.info(
            f"Database has {len(self.programs)} programs, best program score is {score_str}"
        )

    def get_statistics(self, num_recent_iterations: int = 100, k: int = 20, improvement_threshold: float = 0.10) -> Dict[str, Any]:
        """
        Get statistics about the database population.

        Args:
            num_recent_iterations: Number of recent iterations to include in trajectory stats
            k: Number of top scores to return
            improvement_threshold: Minimum score improvement to count as meaningful improvement.
        """
        import statistics

        population_size = len(self.programs)

        if self.programs:
            actual_last_iter = max(
                (p.iteration_found for p in self.programs.values() 
                 if isinstance(p.iteration_found, int)),
                default=0
            )
            self.last_iteration = max(self.last_iteration, actual_last_iter)
        else:
            actual_last_iter = 0

        scores = [
            p.metrics.get("combined_score")
            for p in self.programs.values()
            if p.metrics.get("combined_score") is not None
        ]

        if scores:
            scores_sorted = sorted(scores, reverse=True)
            n = len(scores_sorted)
            if n >= 4:
                quartiles = statistics.quantiles(scores, n=4)  
                q25, q50, q75 = quartiles[0], quartiles[1], quartiles[2]
            else:
                q25 = q50 = q75 = statistics.median(scores)
            score_stats = {
                "best": scores_sorted[0],
                "q75": q75,
                "q50": q50,
                "q25": q25,
                "worst": scores_sorted[-1],
            }

            top_scores = scores_sorted[:k]            
            n = len(scores)

            pct_top = sum(1 for s in scores if s >= q75) / n * 100
            pct_upper_mid = sum(1 for s in scores if q50 <= s < q75) / n * 100
            pct_lower_mid = sum(1 for s in scores if q25 <= s < q50) / n * 100
            pct_bottom = sum(1 for s in scores if s < q25) / n * 100
            
            unique_scores = len(set(round(s, 4) for s in scores))

            score_stats["score_tiers"] = {
                "top": {"threshold": f"score >= {q75:.4f}", "pct_programs": pct_top},
                "upper_mid": {"threshold": f"{q50:.4f} <= score < {q75:.4f}", "pct_programs": pct_upper_mid},
                "lower_mid": {"threshold": f"{q25:.4f} <= score < {q50:.4f}", "pct_programs": pct_lower_mid},
                "bottom": {"threshold": f"score < {q25:.4f}", "pct_programs": pct_bottom},
            }
            score_stats["unique_scores"] = unique_scores
        else:
            score_stats = {"best": None, "q75": None, "q50": None, "q25": None, "worst": None}
            top_scores = []

        programs_with_parents = [p for p in self.programs.values() if p.parent_id is not None]
        unique_parents = len({p.parent_id for p in programs_with_parents})
        avg_solutions_per_parent = len(programs_with_parents) / unique_parents if unique_parents > 0 else 0.0

        iterations_without_improvement = 0
        if scores:
            best_score = max(scores)
            near_best_programs = [
                p for p in self.programs.values()
                if p.metrics.get("combined_score") is not None
                and p.metrics.get("combined_score") >= best_score - improvement_threshold
            ]
            if near_best_programs:
                iteration_near_best_achieved = min(
                    p.iteration_found for p in near_best_programs 
                    if isinstance(p.iteration_found, int)
                )
                iterations_without_improvement = actual_last_iter - iteration_near_best_achieved

        if not self.programs:
            recent_stats = {}
            recent_programs = []
        else:
            recent_programs = [
                p for p in self.programs.values()
                if p.metrics.get("combined_score") is not None
                and isinstance(p.iteration_found, int)
                and p.iteration_found > actual_last_iter - num_recent_iterations
            ]
            recent_programs.sort(key=lambda p: p.iteration_found)
            
            execution_trace = []
            recent_scores = []
            parent_scores = []
            
            for p in recent_programs:
                prog_id = p.id
                prog_score = p.metrics.get("combined_score")
                recent_scores.append(prog_score)
                
                if p.parent_id is not None:
                    parent_id = p.parent_id
                    parent_label = None
                    if p.parent_info is not None:
                        parent_label = p.parent_info[0]
                    if parent_id in self.programs:
                        parent_program = self.programs[parent_id]
                        parent_score = parent_program.metrics.get("combined_score")
                        parent_scores.append(parent_score)
                    else:
                        parent_score = None
                        parent_scores.append(None)
                    parent_tuple = (parent_label, parent_id, parent_score)
                else:
                    parent_tuple = None
                    parent_scores.append(None)
                
                context_tuples = []
                if p.other_context_ids:
                    context_label_map = {}
                    if p.context_info is not None:
                        for label, ctx_id in p.context_info:
                            context_label_map[ctx_id] = label

                    for other_context_id in p.other_context_ids:
                        ctx_label = context_label_map.get(other_context_id)
                        if other_context_id in self.programs:
                            ctx_score = self.programs[other_context_id].metrics.get("combined_score")
                            context_tuples.append((ctx_label, other_context_id, ctx_score))
                        else:
                            context_tuples.append((ctx_label, other_context_id, None))
                
                trace_entry = {
                    "iteration": p.iteration_found,
                    "program": (prog_id, prog_score),
                    "parent": parent_tuple,
                    "context": context_tuples if context_tuples else None,
                }
                execution_trace.append(trace_entry)

            recent_programs_with_parents = [p for p in recent_programs if p.parent_id is not None]

            most_reused_parent_ratio = 0.0
            most_reused_parent_score = None
            most_reused_context_ratio = 0.0
            most_reused_context_score = None

            if recent_programs_with_parents:
                from collections import Counter
                parent_counts = Counter(p.parent_id for p in recent_programs_with_parents if p.parent_id)
                if parent_counts:
                    top_parent_id, parent_count = parent_counts.most_common(1)[0]
                    most_reused_parent_ratio = parent_count / len(recent_programs_with_parents)
                    if top_parent_id in self.programs:
                        most_reused_parent_score = self.programs[top_parent_id].metrics.get("combined_score")

            if recent_programs:
                from collections import Counter
                context_counts = Counter()
                programs_with_context = 0
                for p in recent_programs:
                    if p.other_context_ids:
                        programs_with_context += 1
                        for ctx_id in p.other_context_ids:
                            context_counts[ctx_id] += 1
                if context_counts and programs_with_context > 0:
                    top_ctx_id, context_count = context_counts.most_common(1)[0]
                    most_reused_context_ratio = context_count / programs_with_context
                    if top_ctx_id in self.programs:
                        most_reused_context_score = self.programs[top_ctx_id].metrics.get("combined_score")
                
            recent_stats = {
                "num_recent_iterations": min(num_recent_iterations, len(recent_scores)),
                "execution_trace": execution_trace,
                "score_trajectory": recent_scores,
                "parent_scores": parent_scores,
                "iterations_without_improvement": iterations_without_improvement,
                "improvement_threshold": improvement_threshold,
                "most_reused_parent_ratio": most_reused_parent_ratio,
                "most_reused_parent_score": most_reused_parent_score,
                "most_reused_context_ratio": most_reused_context_ratio,
                "most_reused_context_score": most_reused_context_score,
            }

        return {
            "previous_programs": recent_programs,
            "population_size": population_size,
            "solution_score_summary": score_stats,
            "avg_solutions_per_parent": avg_solutions_per_parent,
            "top_solution_scores": top_scores,
            "recent_solution_stats": recent_stats,
        }