import asyncio
import uuid
from typing import Generic, Callable, Awaitable, List, Any, Optional

from .db import CatDB
from .types import InputType, OutputType, DatabaseLocation, ExperimentConfig


class OutputValidator(Generic[OutputType]):
    def __init__(self, validator_id: str, name: str, confidence_level: float, predicate: Callable[[OutputType], bool]) -> None:
        self.id = validator_id
        self.name = name
        self.confidence_level = confidence_level
        self.predicate = predicate


class InputOutputValidator(Generic[InputType, OutputType]):
    def __init__(self, validator_id: str, name: str, confidence_level: float, predicate: Callable[[InputType, OutputType], bool]) -> None:
        self.id = validator_id
        self.name = name
        self.confidence_level = confidence_level
        self.predicate = predicate


class ExperimentRunner(Generic[InputType, OutputType]):
    """
    A class for managing experiment executions and recordings.

    This class interacts with a database (`CatDB`) to insert experiment values
    and execute them asynchronously.

    Args:
        db (CatDB): The database instance for storing experiment results.

    Attributes:
        db (CatDB): Stores the database connection for experiment execution.
    """
    def __init__(self, db: CatDB) -> None:
        """
        Initializes the ExperimentRunner with a database.

        Args:
            db (CatDB): The database instance used for storing experiment data.
        """
        self.db = db

    async def execute_experiment_run(self, config: ExperimentConfig[InputType, OutputType], run_id: str) -> List[Any]:
        insert_value_sets = await self.db.insert_values(config, run_id)
        return await asyncio.gather(*[self.db.execute_db_insert(values) for values in insert_value_sets])

    async def record_experiment(self, config: ExperimentConfig[InputType, OutputType], run_id: str) -> None:
        try:
            await self.execute_experiment_run(config, run_id)
        except Exception as error:
            print('Error in experiment:', error)
            raise


class BaseCat:
    """Base class for CAT testing providing common functionality."""

    db_path = "data/cat-tests.db"

    @staticmethod
    async def run_base_experiment(cats, experiment_config: ExperimentConfig):
        experiment_run_id = uuid.uuid4()
        for _ in range(experiment_config.times_to_repeat):
            await cats.record_experiment(config=experiment_config, run_id=experiment_run_id)

        rows = await cats.db.execute_query(
            f"""
            SELECT predicate_result, output, validator_id
            FROM {cats.db.table_name}
            WHERE experiment_id = ?
            ORDER BY timestamp DESC
            """,
            (experiment_config.experiment_id,),
        )

        return rows

    def validate_single_response(self, experiment_rows, validator, message_context: str) -> None:
        """Validate a single response from experiment rows."""
        if not self._has_failed_validation(experiment_rows):
            return

        response_output = experiment_rows[0][1]  # Get output from first row
        self.validate_response(response_output, validator, message_context)

    def validate_multiple_responses(self, experiment_rows, validator, inputs: List[str]) -> None:
        """Validate multiple responses from experiment rows."""
        if not self._has_failed_validation(experiment_rows):
            return

        response_outputs = self._parse_response_outputs(experiment_rows[0][1])
        for idx, (input_msg, response) in enumerate(zip(inputs, response_outputs), 1):
            context = f"Response {idx} to input: '{input_msg}'"
            self.validate_response(response, validator, context)

    @staticmethod
    def _has_failed_validation(rows) -> bool:
        """Check if experiment rows indicate a failed validation."""
        return bool(rows) and not rows[0][0]  # rows[0][0] is predicate_result

    @staticmethod
    def _parse_response_outputs(output_str: str) -> List[str]:
        """Parse the string representation of response outputs into a list."""
        return eval(output_str)
            


async def with_cat(location: DatabaseLocation, callback: Callable[[ExperimentRunner[Any, Any]], Awaitable[Any]]) -> Any:
    cats = ExperimentRunner[Any, Any](CatDB(location))
    return await callback(cats)
