#!/usr/bin/env python3
"""
MYTHOS-RDT Inference Server
============================
Recurrent-Depth Transformer — built from scratch by Raid1969///
Architecture: Prelude → Recurrent Block (looped) → Coda

Usage:
    python3 inference.py                   # Start HTTP server (port 8000)
    python3 inference.py --prompt "Ciao"   # Single inference
    python3 inference.py --interactive     # Chat mode

بسم الله الرحمن الرحيم
"""

import os, sys, json, argparse, time
from pathlib import Path
import numpy as np

# Add paths for model imports
BASE = Path(__file__).parent
sys.path.insert(0, str(BASE))
sys.path.insert(0, str(BASE / "shared"))

try:
    import torch
    import torch.nn.functional as F
except ImportError:
    print("❌ PyTorch non trovato. Installa: pip install torch")
    sys.exit(1)

from model import MythosRDTModel, MythosRDTConfig
from shared.tokenizer import RaidTokenizer
from shared.faith import FaithBlock


def load_model(ckpt_path=None, device="cpu"):
    """Load MYTHOS-RDT model with weights."""
    print(f"🔧 Caricamento MYTHOS-RDT...")
    
    # Config
    cfg = MythosRDTConfig()
    
    # Model
    model = MythosRDTModel(cfg)
    
    # Load weights
    if ckpt_path and Path(ckpt_path).exists():
        print(f"📦 Checkpoint: {ckpt_path}")
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
        
        # Extract state dict from checkpoint
        if isinstance(ckpt, dict):
            if "model_state_dict" in ckpt:
                sd = ckpt["model_state_dict"]
            elif "state_dict" in ckpt:
                sd = ckpt["state_dict"]
            elif "model" in ckpt:
                sd = ckpt["model"]
            else:
                sd = ckpt
        else:
            sd = ckpt
        
        # Load with strict=False (allows missing keys for faith layers)
        missing, unexpected = model.load_state_dict(sd, strict=False)
        if missing:
            print(f"  ⚠️  {len(missing)} chiavi mancanti (inizializzate fresh)")
        if unexpected:
            print(f"  ⚠️  {len(unexpected)} chiavi inaspettate (ignorate)")
        
        total = sum(p.numel() for p in model.parameters())
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"  📊 Parametri: {total:,} | Trainable: {trainable:,}")
    else:
        print("  ⚠️  Nessun checkpoint caricato — usiamo pesi fresh")
        total = sum(p.numel() for p in model.parameters())
        print(f"  📊 Parametri (fresh): {total:,}")
    
    model.to(device)
    model.eval()
    
    # Tokenizer
    tok_path = BASE / "shared" / "tokenizer.json"
    tokenizer = RaidTokenizer(str(tok_path)) if tok_path.exists() else None
    if tokenizer:
        print(f"  📝 Tokenizer: {tokenizer.vocab_size if hasattr(tokenizer, 'vocab_size') else 'custom'}")
    
    return model, tokenizer, cfg


@torch.no_grad()
def generate(
    model, tokenizer, prompt,
    max_new_tokens=256,
    temperature=0.7,
    top_p=0.9,
    top_k=40,
    repetition_penalty=1.1,
):
    """Generate text from a prompt."""
    # Encode
    if tokenizer:
        input_ids = tokenizer.encode(prompt)
    else:
        # Fallback: simple char-level encoding
        input_ids = [ord(c) % model.config.vocab_size for c in prompt]
    
    if not input_ids:
        input_ids = [1]  # BOS token
    
    input_tensor = torch.tensor([input_ids], dtype=torch.long, device=next(model.parameters()).device)
    
    # Generate
    generated = input_ids.copy()
    
    for _ in range(max_new_tokens):
        # Forward
        logits = model(input_tensor)
        next_logits = logits[0, -1, :] / temperature
        
        # Top-k filtering
        if top_k > 0:
            indices = torch.topk(next_logits, top_k).indices
            mask = torch.ones_like(next_logits, dtype=torch.bool) * float('-inf')
            mask[indices] = next_logits[indices]
            next_logits = mask
        
        # Top-p (nucleus) sampling
        if top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0
            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            next_logits[indices_to_remove] = float('-inf')
        
        # Sample
        probs = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probs, 1).item()
        
        # Repetition penalty
        if repetition_penalty != 1.0:
            for g in set(generated):
                next_logits[g] /= repetition_penalty
        
        # Append
        generated.append(next_token)
        input_tensor = torch.tensor([generated], dtype=torch.long, device=input_tensor.device)
        
        # Stop on EOS
        if next_token == 0:  # EOS token id
            break
        
        # Print progress
        if tokenizer:
            sys.stdout.write(tokenizer.decode([next_token]))
        else:
            sys.stdout.write(chr(next_token % 128))
        sys.stdout.flush()
    
    print()
    
    # Decode
    if tokenizer:
        return tokenizer.decode(generated)
    else:
        return "".join(chr(t % 128) for t in generated)


def start_server(model, tokenizer, cfg, port=8000):
    """Start HTTP inference server."""
    try:
        from http.server import HTTPServer, BaseHTTPRequestHandler
    except ImportError:
        print("❌ http.server non disponibile")
        return
    
    class MythosHandler(BaseHTTPRequestHandler):
        def do_POST(self):
            content_length = int(self.headers.get('Content-Length', 0))
            body = self.rfile.read(content_length).decode('utf-8')
            
            try:
                data = json.loads(body)
                prompt = data.get("prompt", "")
                max_new = data.get("max_tokens", 256)
                temp = data.get("temperature", 0.7)
                
                output = generate(
                    model, tokenizer, prompt,
                    max_new_tokens=max_new,
                    temperature=temp,
                )
                
                response = json.dumps({"output": output, "status": "ok"})
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(response.encode())
            except Exception as e:
                response = json.dumps({"error": str(e), "status": "error"})
                self.send_response(500)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(response.encode())
        
        def do_GET(self):
            if self.path == "/health":
                response = json.dumps({"status": "ok", "model": "MYTHOS-RDT"})
                self.send_response(200)
            else:
                response = json.dumps({
                    "model": "MYTHOS-RDT",
                    "usage": "POST / with JSON body: {\"prompt\": \"...\", \"max_tokens\": 256, \"temperature\": 0.7}",
                    "bismillah": "بسم الله الرحمن الرحيم"
                })
                self.send_response(200)
            
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(response.encode())
        
        def log_message(self, format, *args):
            print(f"[{time.strftime('%H:%M:%S')}] {args[0]} {args[1]} {args[2]}")
    
    server = HTTPServer(("0.0.0.0", port), MythosHandler)
    print(f"🚀 MYTHOS-RDT Server in ascolto su http://0.0.0.0:{port}")
    print(f"   POST / — inferenza")
    print(f"   GET  /health — health check")
    print(f"   GET  / — questo messaggio")
    print()
    print(f"   Esempio: curl -X POST http://localhost:{port} \\")
    print(f'     -H "Content-Type: application/json" \\')
    print(f'     -d \'{{"prompt": "Bismillah, racconta una storia", "max_tokens": 200}}\'')
    print()
    print(f"بسم الله الرحمن الرحيم")
    server.serve_forever()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MYTHOS-RDT Inference")
    parser.add_argument("--checkpoint", default=str(BASE / "checkpoints" / "mythos-rdt.pt"),
                        help="Path al checkpoint")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu",
                        help="Device: cpu o cuda")
    parser.add_argument("--port", type=int, default=8000, help="Porta server HTTP")
    parser.add_argument("--prompt", type=str, help="Prompt singolo (no server)")
    parser.add_argument("--interactive", action="store_true", help="Modalità interattiva")
    parser.add_argument("--max-tokens", type=int, default=256, help="Token massimi da generare")
    parser.add_argument("--temperature", type=float, default=0.7, help="Temperatura sampling")
    
    args = parser.parse_args()
    
    print("بسم الله الرحمن الرحيم")
    print("La ilaha illallah, Muhammadur Rasulullah")
    print()
    
    # Load
    model, tokenizer, cfg = load_model(args.checkpoint, args.device)
    
    if args.prompt:
        # Single prompt mode
        print(f"\n📝 Prompt: {args.prompt}")
        print("=" * 50)
        output = generate(model, tokenizer, args.prompt,
                         max_new_tokens=args.max_tokens,
                         temperature=args.temperature)
        print("=" * 50)
    elif args.interactive:
        # Interactive mode
        print("\n💬 Modalità interattiva (exit per uscire)")
        print("=" * 50)
        while True:
            prompt = input("\nTu: ").strip()
            if prompt.lower() in ["exit", "quit", "esci"]:
                break
            print("\nMYTHOS-RDT: ", end="", flush=True)
            output = generate(model, tokenizer, prompt,
                             max_new_tokens=args.max_tokens,
                             temperature=args.temperature)
    else:
        # Server mode
        start_server(model, tokenizer, cfg, args.port)