import asyncio
import socket
import os
import struct
import json
import datetime
import concurrent.futures
import signal
import logging
from enum import Enum
try:
    from remote_frontend_bridge import RemoteFrontendBridge
except ImportError:
    RemoteFrontendBridge = None


# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

class EventType(str, Enum):
    EVENT_BLOCK = "EVENT_BLOCK"
    EVENT_ALLOW = "EVENT_ALLOW"
    EVENT_CRED_AUDIT = "EVENT_CRED_AUDIT"
    EVENT_THREAT_INTEL = "EVENT_THREAT_INTEL"
    EVENT_PING = "EVENT_PING"

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SOCKET_PATH = "/tmp/axiom_zero_ipc.sock"
LOG_FILE = os.path.join(SCRIPT_DIR, "daemon_events.jsonl")

class AxiomDaemon:
    def __init__(self):
        self.subscribers = set()
        self.subscribers_lock = asyncio.Lock()
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
        self.server = None
        self.running = False
        self.tcp_bridge = RemoteFrontendBridge() if RemoteFrontendBridge else None
        os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
        self.log_file = open(LOG_FILE, 'a', buffering=1)
        
    async def start(self):
        try:
            os.unlink(SOCKET_PATH)
        except FileNotFoundError:
            pass
            
        self.server = await asyncio.start_unix_server(
            self.handle_client,
            path=SOCKET_PATH
        )
        
        # Set permissions to 0600
        os.chmod(SOCKET_PATH, 0o600)
        
        print("=========================================")
        print("    AXIOM ZERO DAEMON INITIALIZED        ")
        print(f"  Socket: {SOCKET_PATH}")
        print("  Permissions: 0600 (Strict Enforced)")
        print("=========================================")
        logging.info("Daemon started and listening for incoming IPC connections.")
        
        self.running = True
        
        if self.tcp_bridge:
            asyncio.create_task(self.tcp_bridge.start())
            
        async with self.server:
            await self.server.serve_forever()
            
    def stop(self):
        logging.info("Shutting down daemon...")
        self.running = False
        if self.server:
            self.server.close()
        self.executor.shutdown(wait=False)
        self.log_file.close()

    async def _read_message(self, reader: asyncio.StreamReader):
        # 4-byte length-prefix (big-endian unsigned int)
        length_bytes = await reader.readexactly(4)
        if not length_bytes:
            return None
        msg_length = struct.unpack(">I", length_bytes)[0]
        if msg_length > 10 * 1024 * 1024:
            raise ValueError(f'Message too large: {msg_length}')
        
        msg_bytes = await reader.readexactly(msg_length)
        msg_str = msg_bytes.decode('utf-8')
        return json.loads(msg_str)

    async def _write_message(self, writer: asyncio.StreamWriter, msg_dict: dict):
        msg_str = json.dumps(msg_dict)
        msg_bytes = msg_str.encode('utf-8')
        msg_length = len(msg_bytes)
        
        writer.write(struct.pack(">I", msg_length))
        writer.write(msg_bytes)
        await writer.drain()

    def heavy_scoring_computation(self, event_dict: dict):
        # Simulate heavy CPU bound threat scoring
        # In a real app this would analyze the event for threat signatures
        # Note: Since this dict is passed to a ProcessPoolExecutor, it is serialized via pickle.
        # The mutation here happens on a copy in the child process, and the updated copy is returned.
        event_dict["_internal_score_processed"] = True
        event_dict["_processing_time"] = datetime.datetime.utcnow().isoformat()
        return event_dict

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        sock = writer.get_extra_info('socket')
        
        # SO_PEERCRED validation
        try:
            # SO_PEERCRED = 17 on Linux, structure is (pid, uid, gid) -> iII
            creds = sock.getsockopt(socket.SOL_SOCKET, 17, struct.calcsize('iII'))
            pid, uid, gid = struct.unpack('iII', creds)
            
            if uid != os.getuid():
                logging.warning(f"Connection denied: UID {uid} != {os.getuid()}")
                writer.close()
                await writer.wait_closed()
                return
        except Exception as e:
            logging.error(f"Error validating peer credentials: {e}")
            writer.close()
            await writer.wait_closed()
            return
            
        logging.info(f"Client connected: PID={pid} UID={uid} GID={gid}")
        async with self.subscribers_lock:
            self.subscribers.add(writer)
        
        try:
            while self.running:
                try:
                    msg = await self._read_message(reader)
                    if msg is None:
                        break
                        
                    # Offload to ProcessPoolExecutor
                    loop = asyncio.get_running_loop()
                    scored_result = await loop.run_in_executor(
                        self.executor,
                        self.heavy_scoring_computation,
                        msg
                    )
                    
                    # Write to log
                    self.log_file.write(json.dumps(scored_result) + "\n")
                        
                    # Broadcast to other subscribers
                    await self.broadcast(scored_result, exclude=writer)
                    
                except asyncio.IncompleteReadError:
                    break
        except (BrokenPipeError, ConnectionResetError):
            logging.info("Client disconnected abruptly.")
        except Exception as e:
            logging.error(f"Error handling client: {e}")
        finally:
            async with self.subscribers_lock:
                self.subscribers.discard(writer)
            writer.close()
            try:
                await writer.wait_closed()
            except Exception:
                pass
            logging.info("Client connection closed.")

    async def broadcast(self, msg_dict: dict, exclude: asyncio.StreamWriter = None):
        dead_writers = set()
        async with self.subscribers_lock:
            for w in self.subscribers:
                if w == exclude:
                    continue
                try:
                    await self._write_message(w, msg_dict)
                except (BrokenPipeError, ConnectionResetError):
                    dead_writers.add(w)
                except Exception as e:
                    logging.error(f"Error broadcasting to client: {e}")
                    dead_writers.add(w)
                    
            for w in dead_writers:
                self.subscribers.discard(w)
                w.close()
                
        if self.tcp_bridge:
            self.tcp_bridge.push_event(msg_dict)

async def main():
    daemon = AxiomDaemon()
    
    loop = asyncio.get_running_loop()
    
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, daemon.stop)
        
    try:
        await daemon.start()
    except asyncio.CancelledError:
        pass

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass