"""
Nexus-Nano Inference API - Path Fixed
Model: /app/models/nexus-nano.onnx
Ultra-lightweight single-file engine
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import onnxruntime as ort
import numpy as np
import chess
import time
import logging
import os
from typing import Optional, Tuple

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# ==================== NANO ENGINE ====================

class NexusNanoEngine:
    """Ultra-lightweight chess engine"""
    
    PIECE_VALUES = {
        chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3,
        chess.ROOK: 5, chess.QUEEN: 9, chess.KING: 0
    }
    
    def __init__(self, model_path: str):
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model not found: {model_path}")
        
        logger.info(f"📦 Loading model: {model_path}")
        logger.info(f"💾 Size: {os.path.getsize(model_path)/(1024*1024):.2f} MB")
        
        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = 2
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        self.session = ort.InferenceSession(
            model_path,
            sess_options=sess_options,
            providers=['CPUExecutionProvider']
        )
        
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.nodes = 0
        
        logger.info("✅ Engine ready!")
    
    def fen_to_tensor(self, fen: str) -> np.ndarray:
        board = chess.Board(fen)
        tensor = np.zeros((1, 12, 8, 8), dtype=np.float32)
        
        piece_map = {
            chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
            chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5
        }
        
        for sq, piece in board.piece_map().items():
            r, f = divmod(sq, 8)
            ch = piece_map[piece.piece_type] + (6 if piece.color == chess.BLACK else 0)
            tensor[0, ch, r, f] = 1.0
        
        return tensor
    
    def evaluate(self, board: chess.Board) -> float:
        self.nodes += 1
        tensor = self.fen_to_tensor(board.fen())
        output = self.session.run([self.output_name], {self.input_name: tensor})
        score = float(output[0][0][0]) * 400.0
        return -score if board.turn == chess.BLACK else score
    
    def order_moves(self, board: chess.Board, moves):
        scored = []
        for m in moves:
            s = 0
            if board.is_capture(m):
                v = board.piece_at(m.to_square)
                a = board.piece_at(m.from_square)
                if v and a:
                    s = self.PIECE_VALUES.get(v.piece_type, 0) * 10
                    s -= self.PIECE_VALUES.get(a.piece_type, 0)
            if m.promotion == chess.QUEEN:
                s += 90
            scored.append((s, m))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [m for _, m in scored]
    
    def alpha_beta(
        self, 
        board: chess.Board, 
        depth: int, 
        alpha: float, 
        beta: float
    ) -> Tuple[float, Optional[chess.Move]]:
        
        if board.is_game_over():
            return (-10000 if board.is_checkmate() else 0), None
        
        if depth == 0:
            return self.evaluate(board), None
        
        moves = list(board.legal_moves)
        if not moves:
            return 0, None
        
        moves = self.order_moves(board, moves)
        
        best_move = moves[0]
        best_score = float('-inf')
        
        for move in moves:
            board.push(move)
            score, _ = self.alpha_beta(board, depth - 1, -beta, -alpha)
            score = -score
            board.pop()
            
            if score > best_score:
                best_score = score
                best_move = move
            
            alpha = max(alpha, score)
            if alpha >= beta:
                break
        
        return best_score, best_move
    
    def search(self, fen: str, depth: int = 3):
        board = chess.Board(fen)
        self.nodes = 0
        
        moves = list(board.legal_moves)
        if len(moves) == 0:
            return {'best_move': '0000', 'evaluation': 0.0, 'nodes': 0, 'depth': 0}
        
        if len(moves) == 1:
            return {
                'best_move': moves[0].uci(),
                'evaluation': round(self.evaluate(board) / 100.0, 2),
                'nodes': 1,
                'depth': 0
            }
        
        best_move = moves[0]
        best_score = float('-inf')
        current_depth = 1
        
        for d in range(1, depth + 1):
            try:
                score, move = self.alpha_beta(board, d, float('-inf'), float('inf'))
                if move:
                    best_move = move
                    best_score = score
                    current_depth = d
            except:
                break
        
        return {
            'best_move': best_move.uci(),
            'evaluation': round(best_score / 100.0, 2),
            'depth': current_depth,
            'nodes': self.nodes
        }


# ==================== FASTAPI APP ====================

app = FastAPI(
    title="Nexus-Nano Inference API",
    description="Ultra-lightweight chess engine",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

engine = None


class MoveRequest(BaseModel):
    fen: str
    depth: Optional[int] = Field(3, ge=1, le=5)


class MoveResponse(BaseModel):
    best_move: str
    evaluation: float
    depth_searched: int
    nodes_evaluated: int
    time_taken: int


@app.on_event("startup")
async def startup():
    global engine
    logger.info("🚀 Starting Nexus-Nano API...")
    
    # FIXED: Correct path with hyphen
    model_path = "/app/models/nexus-nano.onnx"
    
    logger.info(f"🔍 Looking for: {model_path}")
    
    if os.path.exists("/app/models"):
        logger.info("📂 Files in /app/models/:")
        for f in os.listdir("/app/models"):
            full_path = os.path.join("/app/models", f)
            if os.path.isfile(full_path):
                size = os.path.getsize(full_path) / (1024*1024)
                logger.info(f"   ✓ {f} ({size:.2f} MB)")
    else:
        logger.error("❌ /app/models/ not found!")
        raise FileNotFoundError("/app/models/ directory missing")
    
    if not os.path.exists(model_path):
        logger.error(f"❌ Model not found: {model_path}")
        logger.error("💡 Available:", os.listdir("/app/models"))
        raise FileNotFoundError(f"Missing: {model_path}")
    
    try:
        engine = NexusNanoEngine(model_path)
        logger.info("🎉 Nexus-Nano ready!")
    except Exception as e:
        logger.error(f"❌ Load failed: {e}", exc_info=True)
        raise


@app.get("/health")
async def health():
    return {
        "status": "healthy" if engine else "unhealthy",
        "model": "nexus-nano",
        "version": "1.0.0",
        "model_loaded": engine is not None,
        "model_path": "/app/models/nexus-nano.onnx"
    }


@app.post("/get-move", response_model=MoveResponse)
async def get_move(req: MoveRequest):
    if not engine:
        raise HTTPException(status_code=503, detail="Engine not loaded")
    
    try:
        chess.Board(req.fen)
    except:
        raise HTTPException(status_code=400, detail="Invalid FEN")
    
    start = time.time()
    
    try:
        result = engine.search(req.fen, req.depth)
        elapsed = int((time.time() - start) * 1000)
        
        logger.info(
            f"✓ Move: {result['best_move']} | "
            f"Eval: {result['evaluation']:+.2f} | "
            f"Depth: {result['depth']} | "
            f"Nodes: {result['nodes']} | "
            f"Time: {elapsed}ms"
        )
        
        return MoveResponse(
            best_move=result['best_move'],
            evaluation=result['evaluation'],
            depth_searched=result['depth'],
            nodes_evaluated=result['nodes'],
            time_taken=elapsed
        )
        
    except Exception as e:
        logger.error(f"❌ Search error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/")
async def root():
    return {
        "name": "Nexus-Nano Inference API",
        "version": "1.0.0",
        "model": "2.8M parameters",
        "architecture": "Compact ResNet",
        "speed": "0.2-0.5s per move @ depth 3",
        "status": "online" if engine else "starting",
        "endpoints": {
            "POST /get-move": "Get best move",
            "GET /health": "Health check",
            "GET /docs": "API docs"
        }
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,
        log_level="info",
        access_log=True
    )