import os
import dspy
from dotenv import load_dotenv
from dspy import InputField, OutputField, Signature, Module
import re
load_dotenv()
class AnalyzeAndFixVLMOutput(Signature):
    """
    You are a professional environmental analyst. Using the given visual features of the image, generate structured report addressing the following:

    1. Safety Overview: Evaluate whether the water poses immediate health or environmental risks based on visible features like color, turbidity, oil sheens, and transparency. Focus on the observed characteristics.

    2. Key Features: Briefly identify significant visual aspects such as water clarity, sediment presence, organic material, or surface conditions. Avoid assumptions beyond the visible evidence.

    3. Physical Appearance: Describe the water's appearance (e.g., color, transparency, surface reflection), and explain how these may indicate the water's general quality, without making overreaching conclusions.

    4. Broad Classification: Provide a general classification (Clean, Polluted, Requires Further Testing) based solely on the visible evidence. Clearly justify the classification with specific reference to the observed features.

    5. Environmental Impact: Discuss how the water's condition may affect surrounding ecosystems or human activity, staying within the scope of the visual observations and avoiding premature conclusions.

    6. Economic Considerations: Mention any potential economic impacts (e.g., treatment needs, impact on local industries), emphasizing that further testing is required to confirm these impacts, based on the extracted features.

    7. Recommendations: Offer limited, relevant recommendations based on the water's visible conditions, without over-prescribing actions. Highlight the need for further investigation if necessary. Ensure the report remains objective and based strictly on the observed features."
    
    ---
    Input:
    - Visual features: Color, Turbidity, Oil Sheens, Transparency, etc.
    
    Output:
    - A professional report, ensuring the content is structured, objective, and provides actionable insights.

    MOST IMPORTANT: IF THE IMAGE IS NOT OF WATER, PLEASE RETURN 'NO WATER FOUND'
    """
    visual_features = InputField(type=str, desc="Visual features: color, turbidity, oil sheens, water transparency, etc.")
    additional_info = InputField(type=str, desc="Additional information: Numerical values, etc.")
    corrected_report = OutputField(type=str, desc="Generate a structured, and professionally written water quality report. Ensure that the report is fact-checked, free from overreaching conclusions, and strictly based on visible evidence.")

class TextToMarkdown(Signature):
    """
    Convert the report text to professional Markdown format.
    ---
    Input:
    - Text report
    
    Output:
    - A well-structured Markdown report with clear headings, bullet points, and proper formatting to enhance readability.
    """
    report = InputField(type= str, desc="Text report.")
    markdown = OutputField(type=str,desc="Professionally formatted Markdown report with clear headings and structure.")
    
class English2Urdu(Signature):
    """
    Translate the provided English text into professional and accurate Urdu.
    ---
    Input:
    - English text
    
    Output:
    - Urdu translation of the text, maintaining the original meaning, tone, and professionalism.
    """
    english_text = InputField(type=str, desc="English text to be translated.")
    urdu_text = OutputField(type=str, desc="Urdu translation of the text with accurate meaning and tone.")

class ReportGenerator(Module):
    def __init__(self, model_names):
        super().__init__()
        # Convert single model name to list if needed
        self.model_names = model_names if isinstance(model_names, list) else [model_names]
        self.current_model_index = 0
        self.initialize_model()

    def initialize_model(self):
        """Initialize the model with the current model name"""
        model_name = self.model_names[self.current_model_index]
        lm = dspy.GROQ(model=model_name, api_key=os.getenv("GROQ_API_KEY"), max_tokens=4096)
        dspy.settings.configure(lm=lm)
        self.parser = dspy.ChainOfThought(AnalyzeAndFixVLMOutput)

    def forward(self, vlm_output, additional_info=None):
        """Generate report with fallback to other models if one fails"""
        last_error = None
        for i in range(len(self.model_names)):
            try:
                # Try with current model
                corrected_report = self.parser(visual_features=str(vlm_output), additional_info=additional_info)
                # Update current model index for next request
                self.current_model_index = (self.current_model_index + 1) % len(self.model_names)
                return corrected_report
            except Exception as e:
                last_error = e
                # Try next model
                self.current_model_index = (self.current_model_index + 1) % len(self.model_names)
                self.initialize_model()
                continue

        # If all models fail, raise the last error
        raise last_error or Exception("All models failed to generate a response")