"""
logos/server.py - Matroska Router Server
Protocol 25: Recursive Manifold Engine (RLM) w/ Harmonic Convergence

This server acts as the "Manifold Constraint," forcing all traffic through your Matroska logic.
It implements tiered token consumption, routing based on harmonic resonance, and recursive state refinement.
"""

from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_sock import Sock
from logos.agent_dispatcher import NeuralRouter, PERSONAS, LogosSwarm
import numpy as np
import logging
import sys
import asyncio
from logos.agents.video_atomizer import VideoAtomizer
import requests

# Force UTF-8 encoding for Windows consoles (Protocol 24: Charmap Resilience)
if sys.platform == 'win32':
    if hasattr(sys.stdout, 'reconfigure'):
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
        sys.stderr.reconfigure(encoding='utf-8', errors='replace')
    else:
        import codecs
        sys.stdout = codecs.getwriter("utf-8")(sys.stdout.detach())
        sys.stderr = codecs.getwriter("utf-8")(sys.stderr.detach())

# --- CONFIGURATION ---
from logos.config import SERVER_HOST, SERVER_PORT, LLM_ENDPOINT, UNIFIED_MODEL_ID

# --- CONFIGURATION ---
HOST = SERVER_HOST
PORT = SERVER_PORT 

# Initialize the Flask "Manifold"
app = Flask(__name__)
sock = Sock(app)
CORS(app, resources={r"/*": {"origins": "*"}}) # Full Permissive CORS for Local Swarm
# We use the existing NeuralRouter logic but adapted for this server
swarm_os = LogosSwarm(base_url=LLM_ENDPOINT)
v_node = VideoAtomizer()

# Global Client Manager for Broadcast Pulse
class ConnectionManager:
    def __init__(self):
        self.active_connections = []

    def connect(self, ws):
        self.active_connections.append(ws)

    def disconnect(self, ws):
        if ws in self.active_connections:
            self.active_connections.remove(ws)

    def broadcast(self, message):
        import json
        for connection in self.active_connections:
            try:
                connection.send(json.dumps(message))
            except:
                pass

manager = ConnectionManager()

@sock.route('/neural-link')
def neural_link(ws):
    """
    Protocol 19: WebSocket Neural Bridge for Realtime Telemetry.
    """
    manager.connect(ws)
    try:
        while True:
            data = ws.receive()
            if data:
                # Handle Command from GUI
                import json
                try:
                    payload = json.loads(data)
                    content = payload.get('content')
                    
                    if content:
                        logger.info(f"[GUI] Received Command: {content}")
                        # Execute via Swarm (Async run in thread)
                        # We use a simple non-blocking trigger here
                        loop = asyncio.new_event_loop()
                        asyncio.set_event_loop(loop)
                        res = loop.run_until_complete(swarm_os.process(content))
                        loop.close()
                        
                        # Broadcast Result Back
                        manager.broadcast({
                            "type": "TENSOR_UPDATE",
                            "node": res.get('node'),
                            "origin": swarm_os.state.get('last_node', 0),
                            "tensor": res.get('tensor'),
                            "status": res.get('status')
                        })
                except Exception as e:
                    logger.error(f"[GUI] WS Error: {e}")
    except:
        pass
    finally:
        manager.disconnect(ws)

# --- MANIFOLD STATE TRACKING ---
from logos.manifold_state import ManifoldState
manifold = ManifoldState()

# Set up Logging (Telemetry)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("LOGOS_Router")

# ==========================================
# PROTOCOL 25: RECURSIVE MANIFOLD ENGINE (RLM)
# ==========================================

from logos.mhc_router import execute_recursive_manifold, calculate_manifold_constraint, SHELL_CONFIG, RATE_LIMITS

# ==========================================
# PROTOCOL 26: GÖDEL-ZETA DATASTORE
# ==========================================
from logos.memory.prime_db import PrimeTokenDB
prime_db = PrimeTokenDB()

# ==========================================
# PROTOCOL 40: MTL INTERPRETER (Genesis Kernel)
# ==========================================
try:
    from logos.mtl.interpreter import MTLInterpreter
    from logos.kernel import GenesisKernel
    mtl_interpreter = MTLInterpreter()
    genesis_kernel = GenesisKernel()
    MTL_AVAILABLE = True
    logger.info("[SERVER] MTL Interpreter and Genesis Kernel loaded")
except ImportError as e:
    mtl_interpreter = None
    genesis_kernel = None
    MTL_AVAILABLE = False
    logger.warning(f"[SERVER] MTL not available: {e}")

# Simple in-memory index for the session (Simulating the Topology Graph)
# Map[composite_id] -> filepath
TOPOLOGY_INDEX = {}

# --- PROTOCOL 29: CONTEXT SERVICE ENDPOINTS ---

@app.route('/v1/context/neurons', methods=['POST'])
def upsert_neurons():
    """Batch Upsert Neurons."""
    data = request.json
    neurons = data.get('neurons', [])
    updated = []
    for n in neurons:
        updated.append(manifold.upsert_neuron(n))
    return jsonify({"status": "success", "upserted": len(updated), "neurons": updated})

@app.route('/v1/context/buffer', methods=['POST'])
def update_context_buffer():
    """
    Protocol 30: Context Injection.
    Updates the active swarm memory from external agents (like the CLI Video Atomizer).
    """
    data = request.json
    atoms = data.get('atoms', [])
    if atoms:
        swarm_os.state['context_buffer'] = atoms
        logger.info(f"[CONTEXT] Buffer Updated via API: {atoms}")
        return jsonify({"status": "UPDATED", "count": len(atoms)})
    return jsonify({"status": "NO_CHANGE"}), 400

@app.route('/v1/context/query', methods=['POST'])
def query_context():
    """Semantic/Topological Query."""
    data = request.json
    results = manifold.query_neurons(
        query_text=data.get('query_text'),
        filters=data.get('filters'),
        limit=data.get('limit', 10)
    )
    return jsonify({"results": results, "count": len(results)})

@app.route('/v1/context/neuron/<int:prime_index>', methods=['GET'])
def get_neuron_prime(prime_index):
    """Direct Access by Prime Index."""
    neuron = manifold.get_neuron_by_prime(prime_index)
    if neuron:
        return jsonify(neuron)
    return jsonify({"error": "Not Found"}), 404

# --- API ENDPOINTS ---

@app.route('/', methods=['GET'])
@app.route('/v1', methods=['GET'])
def health_check():
    summary = manifold.get_summary()
    return jsonify({
        "status": "online",
        "system": "LOGOS Matroska Router",
        "protocol": "Recursive Manifold (Protocol 25) + Gödel-Zeta (Protocol 26)",
        "shells": list(SHELL_CONFIG.keys()),
        "manifold_state": summary,
        "topology_size": len(TOPOLOGY_INDEX)
    })

@app.route('/index-module', methods=['POST'])
def index_module():
    """
    Encodes file content into a unique Composite Integer (Gödel Number).
    The file effectively becomes a number in the infinite prime field.
    """
    data = request.json
    filepath = data.get('filepath')
    content = data.get('content', '')
    
    if not filepath: return jsonify({"error": "filepath required"}), 400
    
    # 1. Tokenize (Extract keywords/atoms)
    # Simple heuristic: split by non-alphanumeric, filter small words
    import re
    words = re.findall(r'\b\w+\b', content.lower())
    significant_tokens = [w for w in words if len(w) > 3][:50] # Limit to top 50 for now
    
    # 2. Encode into Manifold
    composite_id, primes = prime_db.encode_state(significant_tokens)
    
    # 3. Store in Topology
    TOPOLOGY_INDEX[composite_id] = filepath
    
    logger.info(f"[INDEX] {filepath} -> Manifold ID: {composite_id}")
    
    return jsonify({
        "status": "INDEXED", 
        "manifold_id": composite_id,
        "prime_coordinates": primes,
        "token_count": len(significant_tokens)
    })

@app.route('/query-topology', methods=['GET'])
def query_topology():
    """
    Finds files that contain the Concept (Prime).
    Operation: O(1) Divisibility Check per node.
    """
    concept = request.args.get('concept')
    if not concept: return jsonify({"error": "concept required"}), 400
    
    # 1. Get the Prime for the concept
    target_prime = prime_db.get_token_prime(concept)
    
    # 2. "Scan the Manifold" (Divisibility Check)
    matches = []
    for comp_id, fpath in TOPOLOGY_INDEX.items():
        # THE GODEL CHECK: O(1) Divisibility
        if comp_id % target_prime == 0:
            matches.append({
                "file": fpath,
                "manifold_id": comp_id
            })
            
    return jsonify({
        "matches": matches, 
        "concept": concept,
        "concept_prime": target_prime,
        "total_nodes_scanned": len(TOPOLOGY_INDEX)
    })

@app.route('/ingest', methods=['POST'])
def ingest_signal():
    """
    PROTOCOL 25: MANIFOLD INGESTION (Zero-Loss)
    Strictly enforcing Prime Token DB. All data entering the graph must be an Integer.
    """
    data = request.json
    source_val = data.get('value') # Could be text, url, or json
    source_node = data.get('source', 1)
    tensor = data.get('tensor', {})
    
    if not source_val: 
        return jsonify({"error": "Null Signal"}), 400
        
    logger.info(f"[INGEST] Absorbing Signal from Node {source_node}...")
    
    # 1. NORM MINIMIZATION (Text -> Integer)
    # We strip the "Soft" text and keep only the "Hard" Prime Coordinate
    if isinstance(source_val, str):
        # Quick tokenization for the signal value itself if it's short, or use Tensor metadata
        tokens = [source_val[:50]] # Treat the value identity as a token for now
        if 'atoms' in tensor:
            tokens = [t['concept'] for t in tensor.get('atoms', [])]
            
        composite_id, primes = prime_db.encode_state(tokens)
    else:
        # Already integer/object?
        composite_id = 997 # Unknown artifact
        primes = []

    # 2. UPDATE MANIFOLD STATE
    # The signal is now just a number (composite_id) and its vector (primes)
    manifold.graph["nodes"][composite_id] = {
        "type": "SIGNAL_ARTIFACT",
        "prime": composite_id,
        "factors": primes,
        "source": source_node,
        "geometry": tensor.get("coords", {"x":0,"y":0,"z":0})
    }
    
    # Link Source -> Signal
    manifold.graph["edges"].append({
        "source": source_node, 
        "target": composite_id, 
        "weight": len(primes)
    })
    
    return jsonify({
        "status": "ABSORBED",
        "manifold_id": composite_id,
        "norm_minimized": True
    })

# ==========================================
# PROTOCOL 40: MTL EXECUTION ENDPOINT
# ==========================================
@app.route('/v1/mtl/execute', methods=['POST'])
def execute_mtl():
    """Execute MTL code via API."""
    if not MTL_AVAILABLE:
        return jsonify({"error": "MTL not available"}), 503
    
    data = request.json
    code = data.get('code', '')
    
    if not code:
        return jsonify({"error": "No code provided"}), 400
    
    logger.info(f"[MTL] Executing: {code[:100]}...")
    
    try:
        result = mtl_interpreter.execute(code)
        logger.info(f"[MTL] Result: {result}")
        return jsonify({
            "status": "success",
            "result": result,
            "code": code
        })
    except Exception as e:
        logger.error(f"[MTL] Error: {e}")
        return jsonify({"error": str(e)}), 400

@app.route('/v1/kernel/process', methods=['POST'])
def kernel_process():
    """Process a packet through the Genesis Kernel."""
    if not MTL_AVAILABLE or not genesis_kernel:
        return jsonify({"error": "Kernel not available"}), 503
    
    data = request.json
    packet = data.get('packet')
    source = data.get('source', 'API')
    
    if not packet:
        return jsonify({"error": "No packet provided"}), 400
    
    logger.info(f"[KERNEL] Processing packet {packet} from {source}")
    
    try:
        result = genesis_kernel.process_packet(int(packet), source=source)
        return jsonify({
            "status": "success",
            "result": result
        })
    except Exception as e:
        logger.error(f"[KERNEL] Error: {e}")
        return jsonify({"error": str(e)}), 400

@app.route('/favicon.ico', methods=['GET'])
def favicon():
    return "", 204

@app.route('/v1/chat/completions', methods=['GET'])
def chat_completions_probe():
    return jsonify({
        "error": "Method Not Allowed",
        "message": "This endpoint requires POST requests with a JSON body.",
        "geometry": "Matroska V1"
    }), 405

@app.route('/v1/models', methods=['GET'])
def list_models():
    return jsonify({
        "object": "list",
        "data": [
            {"id": "logos-matroska-router", "object": "model", "owned_by": "logos"},
            {"id": "dolphin-x1-8b", "object": "model", "owned_by": "local"},
            {"id": "essentialai/rnj-1", "object": "model", "owned_by": "local"},
            {"id": "google/gemma-3-4b", "object": "model", "owned_by": "local"}
        ]
    })

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    """
    OpenAI-Compatible Endpoint wrapping the LOGOS RLM.
    """
    data = request.json
    messages = data.get('messages', [])
    target_model = data.get('model', UNIFIED_MODEL_ID)
    
    # [FIX] VIRTUAL ID MAPPING
    # If the user/CLI requests the virtual router, map it to the underlying inference engine
    if target_model == "logos-matroska-router":
        target_model = UNIFIED_MODEL_ID
    
    if not messages: return jsonify({"error": "No messages provided"}), 400
    
    last_msg = next((m for m in reversed(messages) if m['role'] == 'user'), None)
    if not last_msg: return jsonify({"error": "No user message found"}), 400

    # Vision Handling Check
    last_prompt = ""
    request.is_vision = False
    if isinstance(last_msg['content'], list):
        request.is_vision = True
        for part in last_msg['content']:
             if part.get('type') == 'text': last_prompt += part.get('text', "") + " "
    else:
        last_prompt = last_msg['content']

    # --- EXECUTE PROTOCOL 25 (RLM) or SWARM DELEGATION ---
    
    # 1. Swarm Delegation (Protocols 17 & 27)
    if last_prompt.startswith("SWARM:") or last_prompt.startswith("RUN_FLOW:"):
        # Direct Handoff to the Neural Router / Swarm
        # Since swarm methods are async, we run them in a new event loop
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        if last_prompt.startswith("RUN_FLOW:"):
            flow_name = last_prompt.replace("RUN_FLOW:", "").strip()
            # Resolve path
            flow_path = os.path.join(os.getcwd(), ".agent", "flows", flow_name)
            if not flow_path.endswith(".json"): flow_path += ".json"
            
            logger.info(f"[SERVER] Delegating Flow to Swarm: {flow_name}")
            result = loop.run_until_complete(swarm_os.execute_flow(flow_path))
            final_state = f"FLOW_EXECUTION_COMPLETE\nResult: {result}"
        else:
             # SWARM: ...
            payload = last_prompt.replace("SWARM:", "").strip()
            logger.info(f"[SERVER] Delegating Task to Swarm: {payload}")
            result = loop.run_until_complete(swarm_os.process(payload))
            final_state = f"SWARM_OP_COMPLETE\nNode: {result.get('node')}\nAlignment: {result.get('alignment')}\nTensor: {result.get('tensor')}"
            
        loop.close()
        
        # Create a mock trajectory for the response format
        trajectory = [{"iter": 0, "shell": "SWARM_DELEGATE"}]
        
    else:
        # 2. Default Recursive Manifold (Protocol 25)
        final_state, trajectory, atomic_state_obj = execute_recursive_manifold(last_prompt, target_model)
        
        # [FIX] Merge transient Atomic Graph -> Global Persistence (Only for RLM)
        if hasattr(atomic_state_obj, "graph"):
            # Merge Nodes
            for nid, n_data in atomic_state_obj.graph["nodes"].items():
                manifold.graph["nodes"][nid] = n_data
                if "geometry" not in n_data:
                    prime_val = n_data.get("prime", 2)
                    heat_val = n_data.get("heat", 0)
                    shell = trajectory[-1]['shell'] if trajectory else "INNER_SHELL"
                    
                    domain_map = {"INNER_SHELL": 0, "PRIME_CHANNEL": 5, "OUTER_SHELL": 10}
                    z_depth = domain_map.get(shell, 5) + (prime_val % 5)
                    
                    n_data["geometry"] = {
                        "position": {"x": heat_val * 10, "y": prime_val % 100, "z": z_depth},
                        "domain": shell
                    }
            manifold.graph["edges"].extend(atomic_state_obj.graph["edges"])
            manifold.resonance_product = atomic_state_obj.resonance_product

    # Construct Token Usage
    prompt_tokens = len(last_prompt) // 4
    completion_tokens = len(final_state) // 4
    total_tokens = prompt_tokens + completion_tokens

    return jsonify({
        "id": f"chatcmpl-logos-{int(time.time())}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": target_model,
        "choices": [{
            "index": 0,
            "message": { "role": "assistant", "content": final_state },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens
        },
        "system_fingerprint": f"logos-rlm-v1-depth-{len(trajectory)}"
    })

if __name__ == '__main__':
    print(f"[SERVER] LOGOS Matroska Router Active on Port {PORT}")
    print(f"[SERVER] Connect Antigravity to: http://localhost:{PORT}/v1")
    app.run(host=HOST, port=PORT)