import re
import datetime
import logging
#import config
import ollama

# Configurable Parameters
CONFIG = {
    'AGE_POINT_LOSS_PER_DAY': 0.01,
    'SIZE_POINTS_PER_BILLION_PARAMETERS': 10,
    'FAILURE_PENALTY_PER_FAILURE': 1.0,
    'SIZE_PENALTY_THRESHOLD': 30,  # Billion parameters
    'SIZE_PENALTY_FACTOR': 0.5,  # Penalty per billion parameters above threshold
    'QUANTIZATION_LEVEL_PENALTY': {
        'q4': -5,
        'q8': 5,
    }
}

# Assume the logger is set up at the top of the file
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("ai")

class Model:
    def __init__(self, model_info):
        """
        Initializes the Model instance.

        Args:
            model_info (dict): Dictionary containing model parameters and metadata.
        """
        self.model_info = model_info
        self.score_adjustments = []
        self.failures = 0

        logger.info(f"Initializing Model instance with model_info: {model_info}")

        self.base_score = self._calculate_base_score()

    def _calculate_base_score(self):
        """
        Calculates the base score based on model parameters.

        Returns:
            float: The calculated base score.
        """
        score = 0

        model_name = self.model_info.get('name', 'Unknown')
        logger.debug(f"Calculating base score for model: {model_name}")

        # Age factor: newer models get higher base score
        today = datetime.date.today()
        modified_at = self.model_info.get('modified_at', today)
        try:
            if not isinstance(modified_at, datetime.date):
                modified_at = datetime.datetime.strptime(modified_at.split('T')[0], '%Y-%m-%d').date()
            age_in_days = (today - modified_at).days
        except Exception as e:
            logger.error(f"Error parsing modified_at: {e}")
            age_in_days = 0  # Default to 0 if date parsing fails

        age_factor = max(0, 100 - age_in_days * CONFIG['AGE_POINT_LOSS_PER_DAY'])  # Models lose points per day
        score += age_factor
        logger.debug(f"Age factor ({age_in_days} days old): {age_factor}")

        # Model size factor
        parameter_size = self.model_info.get('details', {}).get('parameter_size', '0B').lower()
        try:
            if 'b' in parameter_size:
                size_value = float(parameter_size.replace('b', ''))
                if 't' in parameter_size:
                    size_value *= 1000  # Convert trillion to billion
            else:
                size_value = 0
        except ValueError:
            size_value = 0
        size_factor = size_value * CONFIG['SIZE_POINTS_PER_BILLION_PARAMETERS']  # Points per billion parameters
        logger.debug(f"Model size factor ({parameter_size} parameters): {size_factor}")

        # Apply penalty for models that are too large
        if size_value > CONFIG['SIZE_PENALTY_THRESHOLD']:
            excess_size = size_value - CONFIG['SIZE_PENALTY_THRESHOLD']
            size_penalty = excess_size * CONFIG['SIZE_PENALTY_FACTOR']
            logger.debug(f"Applying size penalty for model exceeding {CONFIG['SIZE_PENALTY_THRESHOLD']}B: {size_penalty}")
            score -= size_penalty

        score += size_factor

        # Quantization level factor
        quantization_level = self.model_info.get('details', {}).get('quantization_level', '').lower()
        quantization_factor = CONFIG['QUANTIZATION_LEVEL_PENALTY'].get(quantization_level, 0)
        score += quantization_factor
        logger.debug(f"Quantization factor ({quantization_level}): {quantization_factor}")

        logger.info(f"Calculated base score for model {model_name}: {score}")

        # Cap the base score if necessary
        return score

    def execute_prompt(self, prompt):
        """
        Executes a prompt using the LLM model and returns the text output.

        Args:
            prompt (str): The input prompt to send to the LLM.

        Returns:
            str: The text output generated by the LLM.
        """
        logger.info(f"Executing prompt: {prompt}")

        # Placeholder for actual LLM interaction
        output = self._simulate_llm_response(prompt)
        logger.debug(f"Received output: {output}")

        return output

    def parse_python_code(self, text_output, all_blocks=False):
        """
        Parses the LLM-generated text output to extract Python code.

        Args:
            text_output (str): The text output from the LLM.
            all_blocks (bool): If True, return all code blocks. Otherwise, return only the first one.

        Returns:
            str: The extracted Python code.
        """
        logger.debug("Parsing Python code from text output")

        code_blocks = re.findall(r'```python(.*?)```', text_output, re.DOTALL)
        if not code_blocks:
            logger.warning("No Python code blocks found in the text output")
            return ''

        if all_blocks:
            python_code = '\n'.join(code_blocks).strip()
        else:
            python_code = code_blocks[0].strip()
        
        logger.debug(f"Extracted Python code: {python_code}")

        return python_code

    def adjust_score(self, adjustment):
        """
        Adjusts the model's score based on external evaluation.

        Args:
            adjustment (float): The amount to adjust the score by.
        """
        self.score_adjustments.append(adjustment)
        logger.info(f"Adjusting score by {adjustment}. Total adjustments: {self.score_adjustments}")

    def evaluate_code(self, code):
        """
        Evaluates the extracted Python code for correctness.

        Args:
            code (str): The Python code to evaluate.

        Returns:
            bool: True if the code is syntactically correct, False otherwise.
        """
        logger.info("Evaluating code")

        try:
            compile(code, '<string>', 'exec')
            logger.debug("Code is syntactically correct")
            return True
        except Exception as e:
            self.failures += 1
            logger.warning(f"Code evaluation failed: {e}. Total failures: {self.failures}")
            return False

    def get_score(self):
        """
        Calculates and returns the current score of the model.

        Returns:
            float: The current score.
        """
        logger.debug("Calculating current score")

        # Start with base score
        current_score = self.base_score

        # Apply adjustments
        total_adjustments = sum(self.score_adjustments)
        current_score += total_adjustments

        # Deduct penalties for failures
        failure_penalty = self.failures * CONFIG['FAILURE_PENALTY_PER_FAILURE']  # Each failure deducts points
        current_score -= failure_penalty

        logger.debug(f"Base score: {self.base_score}")
        logger.debug(f"Total adjustments: {total_adjustments}")
        logger.debug(f"Failure penalty: {failure_penalty}")
        logger.info(f"Current score: {current_score}")

        return current_score

    def update_base_score(self):
        """
        Recalculates the base score based on current model parameters.
        """
        logger.info("Updating base score")

        self.base_score = self._calculate_base_score()

    def _simulate_llm_response(self, prompt):
        """
        Simulates an LLM response (placeholder for actual LLM interaction).

        Args:
            prompt (str): The input prompt.

        Returns:
            str: Simulated LLM output.
        """
        logger.debug("Simulating LLM response (placeholder)")

        # This method should be replaced with actual code to interface with the LLM.
        return """
        Here is your Python code:
        ```python
        print('Hello, World!')
        ```
        """

    def __repr__(self):
        model_name = self.model_info.get('name', 'Unknown')
        return f"Model(name={model_name}, base_score={self.base_score}, current_score={self.get_score()})"

class AI:
    def __init__(self, models=None):
        """
        Initializes the AI instance with a list of models.

        Args:
            models (list, optional): A list of Model instances.
        """
        self.models = models if models is not None else []
        logger.info(f"AI initialized with models: {[model.model_info.get('name') for model in self.models]}")

    def update(self, models):
        """
        Updates the models in the AI instance.

        Args:
            models (list): A list of Model instances to update or add.
        """
        for new_model in models:
            # Check if a model with the same name exists
            existing_model = next(
                (model for model in self.models if model.model_info.get('name') == new_model.model_info.get('name')),
                None
            )
            if existing_model:
                # Update existing model's information
                index = self.models.index(existing_model)
                self.models[index] = new_model
                logger.info(f"Updated model '{new_model.model_info.get('name')}'")
            else:
                # Add new model to the list
                self.models.append(new_model)
                logger.info(f"Added new model '{new_model.model_info.get('name')}'")

    def select_best_models(self, exclude_models=None):
        """
        Selects the models in order of their scores.

        Args:
            exclude_models (list, optional): A list of models to exclude from selection.

        Returns:
            list: A list of models ordered by their scores, from highest to lowest.
        """
        if exclude_models is None:
            exclude_models = []

        available_models = [model for model in self.models if model not in exclude_models]

        if not available_models:
            logger.error("No available models to select.")
            return []

        ordered_models = sorted(available_models, key=lambda m: m.get_score(), reverse=True)
        logger.info(f"Ordered models by score: {[model.model_info.get('name') for model in ordered_models]}")
        return ordered_models

    def _attempt_task(self, task_method, prompt, max_attempts=3):
        """
        Internal method to attempt a task with retries.

        Args:
            task_method (callable): The method of Model to execute (e.g., execute_prompt).
            prompt (str): The prompt to use.
            max_attempts (int): Maximum number of attempts per model.

        Returns:
            str or None: The result if successful, None otherwise.
        """
        attempted_models = []
        error_info = ''

        while len(attempted_models) < len(self.models):
            models_ordered = self.select_best_models(exclude_models=attempted_models)
            if not models_ordered:
                break

            model = models_ordered[0]
            attempt = 0
            while attempt < max_attempts:
                attempt += 1
                logger.info(f"Attempt {attempt} with model '{model.model_info.get('name')}'")
                modified_prompt = prompt + error_info
                output = getattr(model, task_method)(modified_prompt)
                code = model.parse_python_code(output)

                if model.evaluate_code(code):
                    logger.info(f"Successfully generated valid code on attempt {attempt} with model '{model.model_info.get('name')}'")
                    return code
                else:
                    logger.warning(f"Generated code is invalid on attempt {attempt} with model '{model.model_info.get('name')}'")
                    error_info = "\n# Note: Previous attempt failed due to syntax error."
                    model.adjust_score(-1)  # Penalize the model for failure

            attempted_models.append(model)
            logger.info(f"Model '{model.model_info.get('name')}' failed after {max_attempts} attempts. Trying next best model.")

        logger.error("Failed to generate valid code after trying all models.")
        return None

    def getCode(self, prompt):
        """
        Generates code for the given prompt using the best available model.

        Args:
            prompt (str): The prompt describing the code to generate.

        Returns:
            str or None: The generated code if successful, None otherwise.
        """
        logger.info(f"getCode called with prompt: {prompt}")
        return self._attempt_task(task_method='execute_prompt', prompt=prompt)

    def getTestCode(self, prompt):
        """
        Generates test code for the given prompt using the best available model.

        Args:
            prompt (str): The prompt describing the test code to generate.

        Returns:
            str or None: The generated test code if successful, None otherwise.
        """
        logger.info(f"getTestCode called with prompt: {prompt}")
        return self._attempt_task(task_method='execute_prompt', prompt=prompt)

if __name__ == "__main__":
    # Load models from Ollama
    model_list = ollama.list()['models']
    models = [Model(model_info) for model_info in model_list]

    # Initialize AI with the list of models
    ai = AI(models=models)

    # Menu for interactive session
    while True:
        print("\nAI Interactive Menu:")
        print("1. Execute a prompt")
        print("2. View current model scores")
        print("3. View last generated response and code")
        print("4. Exit")
        choice = input("Choose an option: ")

        if choice == "1":
            prompt = input("Enter your prompt: ")
            result = ai.getCode(prompt)
            if result:
                print("\nGenerated Code:")
                print(result)
            else:
                print("\nFailed to generate valid code.")

        elif choice == "2":
            print("\nCurrent Model Scores:")
            for model in ai.models:
                print(f"{model.model_info.get('name')}: {model.get_score()}")

        elif choice == "3":
            prompt = input("Enter your prompt to generate and view code blocks: ")
            response = ai.getCode(prompt)
            if response:
                print("\nLast Generated Code:")
                print(response)
            else:
                print("\nNo code was generated.")

        elif choice == "4":
            print("Exiting...")
            break

        else:
            print("Invalid choice. Please try again.")
