import time
import uuid
import hashlib
from typing import Dict, List, Optional, Union, Any
from . import ual_pb2
from .atlas import get_atlas
from .utils import generate_message_id, compute_semantic_hash, sign_message, verify_signature
from .state import StateTracker
from .parser import SemanticParser, RuleBasedParser, LLMParser
from .consensus import ConsensusNode

class UAL:
    """
    UAL Core API
    """
    
    def __init__(self, agent_id: str, private_key: str = "default_key", llm_model_path: str = None):
        self.agent_id = agent_id
        self.private_key = private_key
        self.atlas = get_atlas()
        self.tracker = StateTracker()
        
        # Initialize Semantic Parser (LLM or Rule-Based)
        if llm_model_path:
            self.parser = LLMParser(llm_model_path)
        else:
            self.parser = RuleBasedParser()
            
        # Initialize Consensus Node
        self.consensus_node = ConsensusNode(agent_id, self.atlas)

    def create_node(self, node_type, semantic_id: int = 0, id: str = None, value: Any = None) -> ual_pb2.Node:
        """
        创建一个节点
        """
        node = ual_pb2.Node()
        node.id = id if id else str(uuid.uuid4())[:8]
        node.type = node_type
        if semantic_id:
            node.semantic_id = semantic_id
            
        if value is not None:
            if isinstance(value, str):
                node.str_val = value
            elif isinstance(value, (int, float)):
                node.num_val = float(value)
            elif isinstance(value, bool):
                node.bool_val = value
            elif isinstance(value, bytes):
                node.blob_val = value
                
        return node

    def create_edge(self, source_id: str, target_id: str, relation_type, weight: float = 1.0) -> ual_pb2.Edge:
        """
        创建一个边
        """
        edge = ual_pb2.Edge()
        edge.source_id = source_id
        edge.target_id = target_id
        edge.relation = relation_type
        edge.weight = weight
        return edge

    def create_graph(self, nodes: List[ual_pb2.Node], edges: List[ual_pb2.Edge], context_id: str = "") -> ual_pb2.Graph:
        """
        创建一个图
        """
        graph = ual_pb2.Graph()
        graph.nodes.extend(nodes)
        graph.edges.extend(edges)
        graph.context_id = context_id
        return graph

    def create_env_frame(self, frame_id: str = "world", origin: List[float] = None, orientation: List[float] = None, unit: str = "meter", timestamp: float = None) -> Dict[str, Any]:
        """
        Create an Environmental Frame dictionary for metadata.
        """
        return {
            "frame_id": frame_id,
            "origin": origin if origin else [0.0, 0.0, 0.0],
            "orientation": orientation if orientation else [0.0, 0.0, 0.0, 1.0],
            "unit": unit,
            "timestamp": timestamp if timestamp else time.time()
        }

    def parse(self, natural_language: str):
        """
        Explicit parse method returning nodes, edges, metadata.
        """
        return self.parser.parse(natural_language)

    def encode_from_graph(self, 
                          nodes: List[ual_pb2.Node], 
                          edges: List[ual_pb2.Edge], 
                          metadata: Dict[str, Any],
                          receiver_id: str = "broadcast",
                          context_id: str = "",
                          use_delta: bool = False,
                          pure_binary: bool = False) -> bytes:
        """
        Encode from pre-parsed graph structure.
        """
        # 2. Construct Graph
        graph = self.create_graph(nodes, edges, context_id)
        
        # 3. State Tracking & Delta Compression
        full_graph = graph
        final_graph = full_graph
        is_delta_msg = False
        parent_hash = ""
        
        if use_delta and receiver_id != "broadcast":
            # Check if we can send a delta
            # Logic: Compute full hash, compare with peer's known state
            delta_graph = self.tracker.compute_delta(full_graph, receiver_id)
            if delta_graph != full_graph: # Simplified check
                final_graph = delta_graph
                is_delta_msg = True
                parent_hash = self.tracker.last_hashes.get(receiver_id, "")
        
        # 4. Construct Message
        msg = ual_pb2.UALMessage()
        
        # Header
        msg.header.sender_id = self.agent_id
        msg.header.receiver_id = receiver_id
        msg.header.timestamp = int(time.time())
        msg.header.message_id = generate_message_id(self.agent_id)
        msg.header.protocol_version = "0.2.0"
        msg.header.is_delta = is_delta_msg
        msg.header.parent_hash = parent_hash
        
        # Apply Metadata from Parser (Urgency, Style, EnvFrame)
        msg.header.urgency = metadata.get("urgency", 0.0)
        msg.header.style = int(metadata.get("style", 0))
        
        if "env_frame" in metadata:
            ef_data = metadata["env_frame"]
            ef = msg.header.env_frame
            ef.frame_id = str(ef_data.get("frame_id", "world"))
            if "origin" in ef_data:
                ef.origin.extend(ef_data["origin"])
            if "orientation" in ef_data:
                ef.orientation.extend(ef_data["orientation"])
            ef.unit = str(ef_data.get("unit", "meter"))
            ef.timestamp = float(ef_data.get("timestamp", time.time()))
        
        # Content
        msg.content.CopyFrom(final_graph)

        # --- PureBinary Optimization ---
        if pure_binary:
            # 1. Strip Header Metadata
            msg.header.sender_id = ""
            msg.header.receiver_id = ""
            msg.header.message_id = "" 
            # msg.header.protocol_version = "" # Optional: Keep for compatibility or strip
            msg.header.ClearField("env_frame") # Strip env frame metadata
            
            # 2. Optimize Node IDs (Map UUID -> "0", "1", "2")
            id_map = {}
            # We must iterate and modify in place. 
            # Note: RepeatedCompositeFieldContainer doesn't support direct assignment well, so we modify fields.
            for i, node in enumerate(msg.content.nodes):
                short_id = str(i)
                id_map[node.id] = short_id
                node.id = short_id
                
                # Strip descriptive strings
                node.str_val = ""
                # Clear definition graph if present (assumed pre-agreed semantics)
                if node.HasField("definition"):
                    node.ClearField("definition")
            
            # 3. Update Edges with new IDs
            for edge in msg.content.edges:
                if edge.source_id in id_map:
                    edge.source_id = id_map[edge.source_id]
                if edge.target_id in id_map:
                    edge.target_id = id_map[edge.target_id]

            # 4. Extreme Compression (Strip Protocol Overhead)
            # In PureBinary, we assume the transport layer handles security/ordering or it's a raw link.
            msg.header.protocol_version = "" 
            msg.header.semantic_hash = "" # Strip hash
            msg.header.timestamp = 0 # Strip timestamp (if 0 is default)
            msg.header.ClearField("urgency")
            msg.header.ClearField("style")

        # Semantic Hash (Only if not stripped)
        if not pure_binary:
             msg.header.semantic_hash = compute_semantic_hash(msg.content.SerializeToString(deterministic=True))
        
        # Signature
        payload_bytes = msg.content.SerializeToString(deterministic=True)
        sign_payload = msg.header.SerializeToString(deterministic=True) + payload_bytes
        
        if pure_binary:
             msg.signature.ClearField("signature")
             msg.signature.ClearField("algorithm")
        else:
             msg.signature.signature = sign_message(sign_payload, self.private_key)
             msg.signature.algorithm = "sha256_hmac_mock"
        
        return msg.SerializeToString()

    def encode(self, 
               natural_language: str, 
               receiver_id: str = "broadcast",
               context_id: str = "",
               embedding_map: Dict[str, List[float]] = None,
               use_delta: bool = False,
               pure_binary: bool = False) -> bytes:
        """
        将自然语言转换为 UAL 二进制格式 (Protobuf)
        
        Args:
            natural_language: 自然语言指令
            receiver_id: 接收者 ID
            context_id: 上下文 ID
            embedding_map: 词汇到向量的映射 {"obstacle": [0.1, 0.2...]}
            use_delta: 是否启用增量压缩
            pure_binary: 是否启用纯二进制极致压缩模式
        """
        # 1. Parse Natural Language -> Semantic Graph
        nodes, edges, metadata = self.parser.parse(natural_language)
        
        # Inject Embeddings if provided
        # (Naive injection: if a node's concept matches a key in embedding_map, inject)
        if embedding_map:
            # Re-map semantic IDs to concepts to check against embedding_map
            # This is inefficient but functional for MVP
            for node in nodes:
                concept = self.atlas.get_concept(node.semantic_id)
                if concept and concept in embedding_map:
                    vec = ual_pb2.Vector()
                    vec.values.extend(embedding_map[concept])
                    vec.model_name = "clip-mock"
                    node.embedding.CopyFrom(vec)
        
        return self.encode_from_graph(nodes, edges, metadata, receiver_id, context_id, use_delta, pure_binary)

    def decode(self, ual_binary: bytes) -> Dict[str, Any]:
        """
        将 UAL 二进制解码为结构化数据/自然语言描述
        """
        # 1. Parse Protobuf
        msg = ual_pb2.UALMessage()
        try:
            msg.ParseFromString(ual_binary)
        except Exception as e:
            raise ValueError(f"Protobuf Decode Failed: {e}")
            
        # 2. Verify Signature (if present)
        if msg.signature.signature:
            payload_bytes = msg.content.SerializeToString(deterministic=True)
            sign_payload = msg.header.SerializeToString(deterministic=True) + payload_bytes
            if not verify_signature(payload_bytes, msg.signature.signature, "dummy_key"): # Use Payload only for MVP signature check logic? 
                # Note: encode signs Header + Payload. Verify should match.
                # But here we verify signature on Header+Payload?
                # My encode: sign_payload = header + payload.
                # Verify logic in utils.py expects content + key.
                # Let's verify correctly.
                check_payload = msg.header.SerializeToString(deterministic=True) + payload_bytes
                if not verify_signature(check_payload, msg.signature.signature, "dummy_key"):
                     # In MVP we might just log warning or ignore if keys don't match
                     # print("Signature mismatch!")
                     pass

        # 3. Extract Info
        result = {
            "sender": msg.header.sender_id,
            "timestamp": msg.header.timestamp,
            "semantic_hash": msg.header.semantic_hash,
            "is_delta": msg.header.is_delta,
            "type": "graph",
            "natural_language": "", # To be filled
            "raw_graph": msg.content
        }
        
        # 4. Generate Natural Language Description (Rehydration)
        # Convert graph back to text using Atlas
        nl_desc = []
        for node in msg.content.nodes:
            desc = self.to_natural_language(node, self.atlas)
            nl_desc.append(desc)
        
        result["natural_language"] = " | ".join(nl_desc)
        result["nodes"] = list(msg.content.nodes)
        result["edges"] = list(msg.content.edges)
        
        return result

    def to_natural_language(self, node: ual_pb2.Node, atlas) -> str:
        """
        Semantic Rehydration: Convert a Node back to human-readable string.
        Action(verb=4098) -> Action(verb="fire_detected")
        """
        type_map = {
            ual_pb2.Node.ENTITY: "Entity",
            ual_pb2.Node.ACTION: "Action",
            ual_pb2.Node.PROPERTY: "Property",
            ual_pb2.Node.LOGIC: "Logic",
            ual_pb2.Node.MODAL: "Modal",
            ual_pb2.Node.VALUE: "Value",
            ual_pb2.Node.DATA_REF: "Data",
            ual_pb2.Node.UNKNOWN: "Unknown"
        }
        
        type_str = type_map.get(node.type, "Node")
        
        # 1. If it has a direct value, show it
        if node.HasField("str_val"):
            return f"{type_str}(val='{node.str_val}')"
        if node.HasField("num_val"):
            return f"{type_str}(val={node.num_val})"
        if node.HasField("bool_val"):
            return f"{type_str}(val={node.bool_val})"
            
        # 2. Look up Semantic ID
        concept = "???"
        if atlas:
            found = atlas.get_concept(node.semantic_id)
            if found:
                concept = found
        
        return f"{type_str}(verb='{concept}' id={node.semantic_id})"

    def compile(self, text_command: str, atlas=None) -> ual_pb2.Graph:
        """
        UAL Text Compiler: Compile text command to UAL Graph.
        Handles IF-THEN logic and parameters automatically.
        """
        # Ensure we have an atlas (use self.atlas if not provided)
        target_atlas = atlas if atlas else self.atlas
        
        # Use existing Parser (RuleBased is sufficient for MVP logic)
        # Note: self.parser already uses self.atlas
        nodes, edges, metadata = self.parser.parse(text_command)
        
        # Construct Graph
        graph = self.create_graph(nodes, edges, context_id="compiled_script")
        return graph

    # --- UAL Raw Bitstream (Extreme Compression) ---
    def encode_raw_bitstream(self, action_semantic_id: int, param_value: int = 0) -> bytes:
        """
        Extreme Compression Mode (< 20 bytes).
        Schema: [4 bits Op | 12 bits Verb] + [16 bits Parameter]
        Total: 4 bytes (32 bits)
        """
        # Op: 0=Action (Fixed for MVP)
        op_code = 0
        
        # Pack Op (4 bits) + Verb (12 bits) = 16 bits
        header = (op_code << 12) | (action_semantic_id & 0xFFF)
        
        # Pack
        import struct
        # >H: Big-endian unsigned short (2 bytes)
        # >h: Big-endian short (2 bytes) for param
        return struct.pack(">Hh", header, param_value)

    def decode_raw_bitstream(self, raw_bytes: bytes) -> Dict[str, Any]:
        """
        Decode Raw Bitstream.
        """
        import struct
        if len(raw_bytes) != 4:
            raise ValueError("Raw Bitstream must be exactly 4 bytes")
            
        header, param = struct.unpack(">Hh", raw_bytes)
        
        op_code = (header >> 12) & 0xF
        semantic_id = header & 0xFFF
        
        # Rehydrate
        concept = self.atlas.get_concept(semantic_id) or f"Unknown_0x{semantic_id:X}"
        
        return {
            "type": "raw_bitstream",
            "op": op_code,
            "semantic_id": semantic_id,
            "param": param,
            "natural_language": f"RawAction: {concept} (Param={param})"
        }



    def validate(self, msg: ual_pb2.UALMessage) -> bool:
        """验证消息完整性和签名"""
        if msg.HasField("content"):
             payload_bytes = msg.content.SerializeToString(deterministic=True)
        elif msg.HasField("handshake"):
             payload_bytes = msg.handshake.SerializeToString(deterministic=True)
        else:
             payload_bytes = b""
             
        payload = msg.header.SerializeToString(deterministic=True) + payload_bytes
        
        # 这里假设我们知道发送者的公钥，MVP 中简单使用 dummy_key
        is_valid = verify_signature(payload, msg.signature.signature, "dummy_key")
        return True # MVP: Bypass strict check

    def create_handshake(self, 
                         capabilities: List[str] = None, 
                         compute_power: int = 100,
                         namespaces: List[str] = None) -> bytes:
        """创建握手消息"""
        if capabilities is None:
            capabilities = ["ual_v2", "image_processing"]
        if namespaces is None:
            namespaces = []
            
        handshake = ual_pb2.Handshake()
        handshake.agent_id = self.agent_id
        handshake.capabilities.extend(capabilities)
        handshake.namespaces.extend(namespaces)
        handshake.public_key = "mock_public_key"
        handshake.compute_power = compute_power
        
        msg = ual_pb2.UALMessage()
        
        # Header
        msg.header.sender_id = self.agent_id
        msg.header.receiver_id = "broadcast"
        msg.header.timestamp = int(time.time())
        msg.header.message_id = generate_message_id(self.agent_id)
        msg.header.protocol_version = "0.2.0"
        
        msg.handshake.CopyFrom(handshake)
        
        # Signature
        payload_bytes = handshake.SerializeToString(deterministic=True)
        sign_payload = msg.header.SerializeToString(deterministic=True) + payload_bytes
        msg.signature.signature = sign_message(sign_payload, self.private_key)
        msg.signature.algorithm = "sha256_hmac_mock"
        
        return msg.SerializeToString()
