from dataclasses import dataclass
from enum import Enum
from typing import List, Dict, Any, Optional
import re

class ModelType(Enum):
    O1_MINI = "o1-mini"
    # O1_MINI = "gpt-4o-mini"
    GPT_4O_MINI = "gpt-4o-mini"
    
class UIComponentType(Enum):
    TEXTBOX = "textbox"
    MARKDOWN = "markdown"
    DATAFRAME = "dataframe"

@dataclass
class UIConfig:
    component_type: UIComponentType
    label: str
    default_value: Any = None
    visible: Optional[bool] = True
    interactive: Optional[bool] = True
    lines: Optional[int] = None
    description: str = ""
    show_copy_button: Optional[bool] = False
    elem_classes: Optional[List[str]] = None

@dataclass
class PromptConfig:
    prompt: str
    inputs: List[str]
    outputs: List[str]
    model: ModelType
    description: str = ""
    step: Optional[str] = None
    sub_step: Optional[str] = None
    ui: Dict[str, UIConfig] = None
    thoughts: Optional[List[str]] = None

safe_globals = {
"PromptConfig": PromptConfig,
"ModelType": ModelType,
"UIConfig": UIConfig,
"UIComponentType": UIComponentType,
}

class ConfigParser:
    def __init__(self, content: str):
        self.content = content

    def clean_enum_values(self, text: str) -> str:
        """Clean enum values in the text to make them evaluatable"""
        # Replace enum patterns like <TEXTBOX: 'textbox'> with just the enum name
        enum_pattern = r'<(\w+):\s*\'[^\']+\'>'
        return re.sub(enum_pattern, r'UIComponentType.\1', text)

    def parse_config(self) -> Dict[str, PromptConfig]:
        """Parse the config content and return a dictionary of PromptConfig objects"""
        # Extract each block wrapped by $$ markers
        block_pattern = r"\$\$(.*?)\$\$"
        blocks = re.findall(block_pattern, self.content, re.DOTALL)
        print(f"parse_config: Processing {len(self.content)} characters of content\nFound {len(blocks)} config blocks.")

        configs = {}
        errors = []
        for i, block in enumerate(blocks, 1):
            block = block.strip()
            dict_str = "{" + block + "}"
            try:
                # Clean the enum values before evaluation
                cleaned_dict_str = self.clean_enum_values(dict_str)
                config_dict = eval(cleaned_dict_str, safe_globals)
                configs.update(config_dict)
                key = list(config_dict.keys())[0]
                # print(f"Block {i}: Successfully parsed: {key} - prompt: {config_dict[key].prompt[:100]}")
                print(f"Block {i}: Successfully parsed: {key}")
            except Exception as e:
                errors.append({
                    'block': i,
                    'error': str(e),
                    'content': dict_str
                })
                print(f"Block {i}: Error evaluating block: {e}")

        if errors:
            print("\nParsing errors occurred:")
            for error in errors:
                print(f"\nBlock {error['block']}: {error['error']}")
                print("Content:")
                print(error['content'])

        # print(f"\nFinished parsing {len(configs)} configs")
        return configs