import io
import gc
import os
import traceback
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import onnxruntime as ort
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from huggingface_hub import hf_hub_download

app = FastAPI(title="Mind Into Code: Multi-Engine Studio Pro")

# HF_TOKEN for secure downloads
HF_TOKEN = os.getenv("HF_TOKEN", "").strip()

# Attribution Header Middleware
@app.middleware("http")
async def add_attribution_header(request, call_next):
    if request.method == "OPTIONS":
        return await call_next(request)
    response = await call_next(request)
    response.headers["X-Powered-By"] = "Mind Into Code"
    return response

# CORS Configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["Content-Type", "X-Api-Key", "Authorization", "Accept"],
    expose_headers=["*"],
)

# --- Task 1: Native PyTorch SAFMN Architecture ---
class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim))
        self.eps = eps
    def forward(self, x):
        u = x.mean(1, keepdim=True)
        s = (x - u).pow(2).mean(1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.eps)
        x = self.weight[:, None, None] * x + self.bias[:, None, None]
        return x

class SAFM(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = LayerNorm(dim)
        self.q = nn.Conv2d(dim, dim, 1)
        self.k = nn.Conv2d(dim, dim, 1)
        self.v = nn.Conv2d(dim, dim, 1)
        self.proj = nn.Conv2d(dim, dim, 1)
        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, groups=dim)
    def forward(self, x):
        x = self.norm(x)
        q = self.q(x)
        k = self.k(x)
        v = self.v(x)
        modulation = torch.sigmoid(self.dwconv(k))
        out = self.proj(q * modulation + v)
        return out

class CCM(nn.Module):
    def __init__(self, dim, ffn_scale=2.0):
        super().__init__()
        self.norm = LayerNorm(dim)
        self.conv1 = nn.Conv2d(dim, int(dim * ffn_scale), 1)
        self.conv2 = nn.Conv2d(int(dim * ffn_scale), dim, 1)
    def forward(self, x):
        x = self.norm(x)
        x = self.conv1(x)
        x = F.gelu(x)
        x = self.conv2(x)
        return x

class Block(nn.Module):
    def __init__(self, dim, ffn_scale=2.0):
        super().__init__()
        self.safm = SAFM(dim)
        self.ccm = CCM(dim, ffn_scale)
    def forward(self, x):
        x = x + self.safm(x)
        x = x + self.ccm(x)
        return x

class SAFMN(nn.Module):
    def __init__(self, dim=128, n_blocks=16, ffn_scale=2.0, upscale=4):
        super().__init__()
        self.to_feat = nn.Conv2d(3, dim, 3, 1, 1)
        self.blocks = nn.ModuleList([Block(dim, ffn_scale) for _ in range(n_blocks)])
        self.upsampler = nn.Sequential(
            nn.Conv2d(dim, 3 * (upscale ** 2), 3, 1, 1),
            nn.PixelShuffle(upscale)
        )
    def forward(self, x):
        x = self.to_feat(x)
        res = x
        for block in self.blocks:
            x = block(x)
        x = x + res
        x = self.upsampler(x)
        return x

# ONNX Model Registry
ONNX_MODELS = {
    "SCUNet-GAN": {
        "repo_id": "deepghs/image_restoration",
        "filename": "SCUNet-GAN.onnx",
        "scale": 1
    },
    "NAFNet-REDS": {
        "repo_id": "deepghs/image_restoration",
        "filename": "NAFNet-REDS-width64.onnx",
        "scale": 1 
    }
}

ALLOWED_MODELS = ["studio_pro", "safmn_v1"]

class UnifiedUpscaler:
    def __init__(self):
        self.sessions = {}
        self.models_dir = os.path.join(os.getcwd(), 'Models')
        os.makedirs(self.models_dir, exist_ok=True)
        
        print("[Mind Into Code] Starting Unified Pre-flight Initialization...")
        # Initialize ONNX Pro Pipeline
        for model_id in ONNX_MODELS:
            self._initialize_onnx_session(model_id)
            
        # Task 4: Initialize Native PyTorch SAFMN
        self._initialize_safmn_native()
        
        if "safmn_v1" in self.sessions:
            print('Backend Ready')

    def _initialize_onnx_session(self, model_id):
        config = ONNX_MODELS.get(model_id)
        try:
            local_path = hf_hub_download(
                repo_id=config["repo_id"],
                filename=config["filename"],
                local_dir=self.models_dir,
                token=HF_TOKEN if HF_TOKEN else None
            )
            sess_options = ort.SessionOptions()
            sess_options.intra_op_num_threads = 2
            session = ort.InferenceSession(local_path, sess_options=sess_options, providers=['CPUExecutionProvider'])
            self.sessions[model_id] = {
                "session": session,
                "input_name": session.get_inputs()[0].name,
                "scale": config["scale"]
            }
        except Exception as e:
            print(f"[Mind Into Code] ONNX failure for {model_id}: {e}")

    def _initialize_safmn_native(self):
        # Task 2: CPU-Only Loading
        model_filename = "safmn_l_x4.pth"
        local_path = os.path.join(self.models_dir, model_filename)
        try:
            model = SAFMN(dim=128, n_blocks=16, upscale=4)
            if os.path.exists(local_path):
                print(f"[Mind Into Code] Loading SAFMN-L weights from {local_path}...")
                model.load_state_dict(torch.load(local_path, map_location='cpu'))
            
            model.eval()
            self.sessions["safmn_v1"] = {"model": model, "scale": 4}
        except Exception as e:
            print(f"[Mind Into Code] SAFMN Native failure: {e}")

    def process_tile_onnx(self, tile, model_id):
        sess_data = self.sessions[model_id]
        h, w = tile.shape[:2]
        mod = 32
        pad_h = (mod - h % mod) % mod
        pad_w = (mod - w % mod) % mod
        if pad_h > 0 or pad_w > 0:
            tile = cv2.copyMakeBorder(tile, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)
        
        img = tile.astype(np.float32) / 255.0
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.transpose(img, (2, 0, 1))
        img = np.expand_dims(img, axis=0)
        
        output = sess_data["session"].run(None, {sess_data["input_name"]: img})[0]
        output = np.squeeze(output, axis=0)
        output = np.clip(output, 0, 1)
        output = np.transpose(output, (1, 2, 0))
        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        output = (output * 255.0).astype(np.uint8)
        
        if pad_h > 0 or pad_w > 0:
            output = output[0 : h, 0 : w]
        return output

    def process_tile_native(self, tile):
        # Task 3: Native Inference Logic
        model = self.sessions["safmn_v1"]["model"]
        img = tile.astype(np.float32) / 255.0
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_tensor = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
        
        with torch.no_grad(): # RAM Saver
            output_tensor = model(img_tensor)
        
        output = output_tensor.squeeze(0).permute(1, 2, 0).numpy()
        output = np.clip(output, 0, 1)
        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        output = (output * 255.0).astype(np.uint8)
        return output

    def studio_pro_pipeline(self, img, tile_size=64, padding=8):
        h, w, c = img.shape
        scale = 4
        canvas = np.zeros((h * scale, w * scale, c), dtype=np.uint8)
        for y in range(0, h, tile_size):
            for x in range(0, w, tile_size):
                y1, x1 = max(0, y - padding), max(0, x - padding)
                y2, x2 = min(h, y + tile_size + padding), min(w, x + tile_size + padding)
                tile = img[y1:y2, x1:x2]
                denoised = self.process_tile_onnx(tile, "SCUNet-GAN")
                upscaled = cv2.resize(denoised, (denoised.shape[1] * 4, denoised.shape[0] * 4), interpolation=cv2.INTER_CUBIC)
                restored = self.process_tile_onnx(upscaled, "NAFNet-REDS")
                
                t_h, t_w = min(tile_size, h - y) * 4, min(tile_size, w - x) * 4
                top, left = (y - y1) * 4, (x - x1) * 4
                canvas[y*4:y*4+t_h, x*4:x*4+t_w] = restored[top:top+t_h, left:left+t_w]
        return canvas

    def safmn_pipeline(self, img, tile_size=128, padding=8):
        h, w, c = img.shape
        canvas = np.zeros((h * 4, w * 4, c), dtype=np.uint8)
        for y in range(0, h, tile_size):
            for x in range(0, w, tile_size):
                y1, x1 = max(0, y - padding), max(0, x - padding)
                y2, x2 = min(h, y + tile_size + padding), min(w, x + tile_size + padding)
                tile = img[y1:y2, x1:x2]
                upscaled = self.process_tile_native(tile)
                
                t_h, t_w = min(tile_size, h - y) * 4, min(tile_size, w - x) * 4
                top, left = (y - y1) * 4, (x - x1) * 4
                canvas[y*4:y*4+t_h, x*4:x*4+t_w] = upscaled[top:top+t_h, left:left+t_w]
        return canvas

# Initialize Engine
engine = UnifiedUpscaler()

@app.get("/")
@app.get("/health")
async def health():
    return {"status": "ready", "engines": ALLOWED_MODELS}

@app.post("/upscale")
async def upscale(
    file: UploadFile = File(...),
    model_type: str = Form("studio_pro")
):
    try:
        contents = await file.read()
        input_image = Image.open(io.BytesIO(contents)).convert("RGB")
        img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
        
        if model_type == "safmn_v1":
            output = engine.safmn_pipeline(img)
        else:
            output = engine.studio_pro_pipeline(img)

        _, buffer = cv2.imencode(".png", output)
        return StreamingResponse(io.BytesIO(buffer), media_type="image/png")
    except Exception:
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"error": "Inference Error."})

print('Backend Ready')