from pyee import AsyncIOEventEmitter
from jsonschema import validate, ValidationError
from typing import Any, Dict, Optional, Union
#import json
from promptrix.promptrixTypes import PromptMemory, PromptFunctions, Tokenizer
from dataclasses import dataclass, asdict
import traceback

@dataclass
class CommandSchema:
    def __init__(self, type: str, title: str='', description: str='', returns: Optional[str] = None):
        self.type = type
        self.title = title
        self.description = description
        self.returns = returns

class SchemaBasedCommand(AsyncIOEventEmitter):
    def __init__(self, schema: CommandSchema, title: Optional[str] = None, description: Optional[str] = None):
        super().__init__()
        self._schema = schema
        self._title = title
        self._description = description

    @property
    def description(self) -> str:
        return self._description or self._schema.description

    @property
    def inputs(self) -> Optional[str]:
        if self._schema.properties:
            properties = self._schema.properties or {}
            inputs = [f'"{key}":"<{property.get("description", property.get("type", "any"))} value>"' for key, property in properties.items()]
            return ",".join(inputs)
        else:
            return None

    @property
    def output(self) -> Optional[str]:
        return self._schema.returns

    @property
    def schema(self) -> CommandSchema:
        return self._schema

    @property
    def title(self) -> str:
        return self._title or self._schema.title

    async def execute(self, input: Dict[str, Any], memory: 'PromptMemory', functions: 'PromptFunctions', tokenizer: 'Tokenizer') -> Any:
        raise NotImplementedError

    async def validate(self, input: Dict[str, Any], memory: 'PromptMemory', functions: 'PromptFunctions', tokenizer: 'Tokenizer') -> 'Validation':
        try:
            cleaned = self.clean_input(input)
            if self._schema is not None:
                validate(cleaned, asdict(self._schema))
            return {
                'type': 'Validation',
                'valid': True,
                'value': cleaned
            }
        except ValidationError as e:
            errors = f'"{e.path[0] if e.path else "input"}": {e.message}'
            message = errors
            return {
                'type': 'Validation',
                'valid': False,
                'feedback': f"The ['command']['input'] field has errors:\n{message}\n\nTry again."
            }

    def clean_input(self, input: Dict[str, Any]) -> Dict[str, Any]:
        cleaned = {}
        properties = self._schema.properties
        if not properties:
            return cleaned

        for key, property in properties.items():
            type = property.get('type', 'any')
            value = input.get(key)

            if value is None or (isinstance(value, str) and value.startswith("<") and value.endswith(">")):
                continue

            if not isinstance(value, str):
                value = json.dumps(value) if isinstance(value, dict) else str(value)

            try:
                if type == 'string':
                    cleaned[key] = str(value)
                elif type == 'number':
                    cleaned[key] = float(value)
                elif type == 'boolean':
                    cleaned[key] = bool(value)
                elif type in ['array', 'object']:
                    cleaned[key] = json.loads(value)
            except Exception:
                pass

        return cleaned
