import json
from typing import Union, Any, Dict
from dataclasses import dataclass


message_specs = {
    "RegistrationRequest": {
        "pi": "int",
        "T": "Point",
    },
    "UserCredentials": {
        "X3": "Point",
        "PI3": "ZKP",
        "pi": "int",
        "T": "Point",
    },
    "AuthInitRequest": {
        "X1": "Point",
        "X2": "Point",
        "PI1": "ZKP",
        "PI2": "ZKP",
    },
    "AuthInitialValues": {
        "T": "Point",
        "pi": "int",
        "x4": "int",
        "X1": "Point",
        "X2": "Point",
        "X3": "Point",
        "X4": "Point",
        "beta": "Point",
        "PI1": "ZKP",
        "PI2": "ZKP",
        "PI3": "ZKP",
        "PIBeta": "ZKP",
    },
    "AuthInitResponse": {
        "X3": "Point",
        "X4": "Point",
        "PI3": "ZKP",
        "PI4": "ZKP",
        "beta": "Point",
        "PIBeta": "ZKP",
    },
    "AuthFinishRequest": {
        "alpha": "Point",
        "PIAlpha": "ZKP",
        "r": "int",
    },
}


def indent(text: str, tabs: int) -> str:
    return "    " * tabs + text


def generate_messages():
    template = '''\
# THIS FILE IS GENERATED BY _messages_generator.py
# DO NOT MODIFY DIRECTLY AS IT WILL BE OVERWRITTEN

import json
from typing import Union, Any, Dict
from dataclasses import dataclass

from owl_common import Config, Curves, Point, ZKP
from extended_curves import FourQPoint


def get_curve(curve: Curves):
    
    from cryptography.hazmat.primitives.asymmetric import ec
    
    curve_map = {
        Curves.P256: ec.SECP256R1(),
        Curves.P384: ec.SECP384R1(),
        Curves.P521: ec.SECP521R1(),
    }
    return curve_map.get(curve)


def parse_num(x: Any) -> Union[int, None]:
    """Parse a number from hex string"""
    try:
        if isinstance(x, str):
            return int(x, 16)
        elif isinstance(x, int):
            return x
        return None
    except (ValueError, TypeError):
        return None


def parse_point(x: Any, curve: Curves) -> Union[Point, FourQPoint, None]:
    
    try:
        if curve == Curves.FOURQ:
            return FourQPoint.from_hex(x)
        else:
            curve_obj = get_curve(curve)
            if curve_obj is None:
                return None
            return Point.from_hex(x, curve_obj)
    except Exception:
        return None


def parse_zkp(x: Any) -> Union[ZKP, None]:
    """Parse a ZKP from dict"""
    try:
        if isinstance(x, dict):
            h = parse_num(x.get('h'))
            r = parse_num(x.get('r'))
            if h is not None and r is not None:
                return ZKP(h=h, r=r)
    except Exception:
        pass
    return None


class DeserializationError(Exception):
    def __init__(self, message: str):
        super().__init__(message)
        self.name = "DeserializationError"


'''

    for cls_name, attrs in message_specs.items():

        template += f"@dataclass\n"
        template += f"class {cls_name}:\n"
        
        for attr, attr_type in attrs.items():
            
            if attr_type == "Point":
                template += indent(f"{attr}: Union[Point, FourQPoint]\n", 1)
            else:
                template += indent(f"{attr}: {attr_type}\n", 1)
        
        template += "\n"
        
        template += indent("@classmethod\n", 1)
        template += indent("def deserialize(\n", 1)
        template += indent("cls,\n", 2)
        template += indent("x: Any,\n", 2)
        template += indent("cfg: Config,\n", 2)
        template += indent(f") -> Union['{cls_name}', DeserializationError]:\n", 1)
        
        template += indent('if isinstance(x, str):\n', 2)
        template += indent("try:\n", 3)
        template += indent("x = json.loads(x)\n", 4)
        template += indent("except json.JSONDecodeError:\n", 3)
        template += indent(f'return DeserializationError("Failed to parse JSON for {cls_name}")\n', 4)
        
        template += "\n"
        template += indent("if x is not None:\n", 2)
        
        parsed_attrs = []
        for attr, attr_type in attrs.items():
            if attr_type == "str":
                parse_expr = f"x.get('{attr}')"
            elif attr_type == "int":
                parse_expr = f"parse_num(x.get('{attr}'))"
            elif attr_type == "Point":
                parse_expr = f"parse_point(x.get('{attr}'), cfg.curve)"
            elif attr_type == "ZKP":
                parse_expr = f"parse_zkp(x.get('{attr}'))"
            else:
                parse_expr = f"x.get('{attr}')"
            
            template += indent(f"{attr} = {parse_expr}\n", 3)
            parsed_attrs.append(attr)
        
        template += "\n"
        template += indent("if ", 3)
        conditions = [f"{attr} is not None" for attr in parsed_attrs]
        template += " and ".join(conditions)
        template += ":\n"
        
        template += indent("return cls(\n", 4)
        for i, attr in enumerate(parsed_attrs):
            suffix = "," if i < len(parsed_attrs) - 1 else ""
            template += indent(f"{attr}={attr}{suffix}\n", 5)
        template += indent(")\n", 4)
        
        template += "\n"
        template += indent("return DeserializationError(\n", 2)
        template += indent(f'"Failed to deserialize {cls_name}: invalid data"\n', 3)
        template += indent(")\n", 2)
        
        template += "\n"
        
        template += indent("def serialize(self) -> Dict[str, Any]:\n", 1)
        template += indent('"""Serialize the object to a dict"""\n', 2)
        template += indent("return {\n", 2)
        
        for attr, attr_type in attrs.items():
            if attr_type == "int":
                serialize_expr = f"format(self.{attr}, 'x')"
            elif attr_type == "Point":
                serialize_expr = f"self.{attr}.to_hex()"
            elif attr_type == "ZKP":
                serialize_expr = f"{{'h': format(self.{attr}.h, 'x'), 'r': format(self.{attr}.r, 'x')}}"
            else:
                serialize_expr = f"self.{attr}"
            
            template += indent(f'"{attr}": {serialize_expr},\n', 3)
        
        template += indent("}\n", 2)
        
        template += "\n"
        
        template += indent("def to_json(self) -> str:\n", 1)
        template += indent('"""Serialize the object to JSON"""\n', 2)
        template += indent("return json.dumps(self.serialize())\n", 2)
        
        template += "\n\n"
    
    return template


if __name__ == "__main__":
    content = generate_messages()
    
    with open("messages.py", "w", encoding="utf-8") as f:
        f.write(content)
    
