import dspy
from typing import List, Optional
from .models import EvaluationResult, CategoryEvaluation
from .constants import (
    TWEET_MAX_LENGTH,
    TWEET_TRUNCATION_SUFFIX,
    DEFAULT_SCORE,
    ERROR_PARSING,
    ERROR_VALIDATION,
    ERROR_GENERATION,
    ERROR_EVALUATION,
    MIN_SCORE,
    MAX_SCORE
)
from .helpers import format_evaluation_for_generator, truncate_tweet

class TweetGenerator(dspy.Signature):
    """Generate or improve a tweet based on input text and detailed evaluation feedback with reasoning."""
    
    input_text: str = dspy.InputField(desc="Original text or current tweet to improve")
    current_tweet: str = dspy.InputField(desc="Current best tweet version (empty for first generation)")
    previous_evaluation: str = dspy.InputField(desc="Previous evaluation with category-by-category reasoning and scores (empty for first generation)")
    improved_tweet: str = dspy.OutputField(desc=f"Generated or improved tweet text (max {TWEET_MAX_LENGTH} characters)")

class TweetEvaluator(dspy.Signature):
    """Evaluate a tweet across multiple custom categories. For each category, provide detailed reasoning explaining the score, then assign a score. Ensure the tweet maintains the same meaning as the original text."""
    
    original_text: str = dspy.InputField(desc="Original input text that started the optimization")
    current_best_tweet: str = dspy.InputField(desc="Current best tweet version for comparison (empty for first evaluation)")
    tweet_text: str = dspy.InputField(desc="Tweet text to evaluate")
    categories: str = dspy.InputField(desc="Comma-separated list of evaluation category descriptions")
    evaluations: List[CategoryEvaluation] = dspy.OutputField(
        desc=f"List of evaluations with category name, detailed reasoning, and score ({MIN_SCORE}-{MAX_SCORE}) for each category. Ensure the tweet conveys the same meaning as the original text."
    )

class TweetGeneratorModule(dspy.Module):
    """DSPy module for generating and improving tweets."""
    
    def __init__(self):
        super().__init__()
        self.generate = dspy.ChainOfThought(TweetGenerator)
    
    def forward(self, input_text: str, current_tweet: str = "", previous_evaluation: Optional[EvaluationResult] = None) -> str:
        """Generate or improve a tweet."""
        try:
            # Format previous evaluation as text
            eval_text = format_evaluation_for_generator(previous_evaluation)
            
            result = self.generate(
                input_text=input_text,
                current_tweet=current_tweet,
                previous_evaluation=eval_text
            )
            
            # Ensure tweet doesn't exceed character limit
            tweet = truncate_tweet(result.improved_tweet, TWEET_MAX_LENGTH, TWEET_TRUNCATION_SUFFIX)
            
            return tweet
        except Exception as e:
            raise Exception(f"{ERROR_GENERATION}: {str(e)}")

class TweetEvaluatorModule(dspy.Module):
    """DSPy module for evaluating tweets across custom categories."""
    
    def __init__(self):
        super().__init__()
        self.evaluate = dspy.ChainOfThought(TweetEvaluator)
    
    def forward(self, tweet_text: str, categories: List[str], original_text: str = "", current_best_tweet: str = "") -> EvaluationResult:
        """Evaluate a tweet across specified categories."""
        try:
            # Join categories into comma-separated string
            categories_str = ", ".join(categories)
            
            result = self.evaluate(
                original_text=original_text,
                current_best_tweet=current_best_tweet,
                tweet_text=tweet_text,
                categories=categories_str
            )
            
            # Extract and validate evaluations
            evaluations = result.evaluations
            
            # Ensure we have the right number of evaluations
            if len(evaluations) != len(categories):
                # Create default evaluations if mismatch
                evaluations = [
                    CategoryEvaluation(
                        category=cat,
                        reasoning=ERROR_PARSING,
                        score=DEFAULT_SCORE
                    ) for cat in categories
                ]
            else:
                # Validate each evaluation
                validated_evals = []
                for i, eval in enumerate(evaluations):
                    try:
                        # Ensure score is valid
                        score = max(MIN_SCORE, min(MAX_SCORE, int(eval.score)))
                        validated_evals.append(CategoryEvaluation(
                            category=categories[i] if i < len(categories) else eval.category,
                            reasoning=eval.reasoning if eval.reasoning else "No reasoning provided",
                            score=score
                        ))
                    except (ValueError, TypeError, AttributeError):
                        validated_evals.append(CategoryEvaluation(
                            category=categories[i] if i < len(categories) else "Unknown",
                            reasoning=ERROR_VALIDATION,
                            score=DEFAULT_SCORE
                        ))
                evaluations = validated_evals
            
            # Create validated result
            validated_result = EvaluationResult(evaluations=evaluations)
            
            return validated_result
        except Exception as e:
            # Return default evaluations on error
            default_evals = [
                CategoryEvaluation(
                    category=cat,
                    reasoning=f"{ERROR_EVALUATION}: {str(e)}",
                    score=DEFAULT_SCORE
                ) for cat in categories
            ]
            return EvaluationResult(evaluations=default_evals)