from fastapi import FastAPI, UploadFile, File
from fastapi.responses import Response
import torch
import torch.nn.functional as F
import numpy as np
import cv2
from PIL import Image
import io
import torchvision.transforms.functional as TF
from transformers import SegformerForSemanticSegmentation

app = FastAPI(title="Unified Cartographer API", version="1.0")

# ==========================================
# 1. LOAD MODEL ON SERVER STARTUP
# ==========================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Server booting on: {DEVICE}")

class UnifiedCartographer(torch.nn.Module):
    def __init__(self, num_classes=5):
        super().__init__()
        self.model = SegformerForSemanticSegmentation.from_pretrained(
            "nvidia/segformer-b4-finetuned-cityscapes-1024-1024", 
            num_labels=num_classes, ignore_mismatched_sizes=True
        )
    def forward(self, x):
        outputs = self.model(pixel_values=x)
        return F.interpolate(outputs.logits, size=x.shape[-2:], mode="bilinear", align_corners=False)

model = UnifiedCartographer(num_classes=5).to(DEVICE)
# ⚠️ Ensure this filename matches your uploaded weights perfectly
model.load_state_dict(torch.load("unified_cartographer_andhra_v4.pth", map_location=DEVICE))
model.eval()

# ==========================================
# 2. THE RENDERING ENGINE (Now with Legend!)
# ==========================================
def draw_legend(canvas):
    """Draws a semi-transparent legend in the top-left corner"""
    legend_items = [
        ("Bldg (High Conf)", (255, 0, 0)),     # Red
        ("Bldg (Med Conf)", (255, 165, 0)),    # Orange
        ("Bldg (Low Conf)", (255, 0, 255)),    # Magenta
        ("Roads", (255, 255, 0)),              # Yellow
        ("Water", (0, 100, 255)),              # Blue
        ("Open Area", (50, 205, 50))           # Green
    ]
    
    # 1. Create a semi-transparent black background box for readability
    overlay = canvas.copy()
    cv2.rectangle(overlay, (10, 10), (210, 160), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.7, canvas, 0.3, 0, canvas)
    
    # 2. Add the title of the legend
    cv2.putText(canvas, "AI TACTICAL MAP", (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
    
    # 3. Draw the color boxes and labels
    y_offset = 55
    for label, color in legend_items:
        # Draw the little color square
        cv2.rectangle(canvas, (20, y_offset - 10), (35, y_offset + 5), color, -1)
        # Draw the text next to it
        cv2.putText(canvas, label, (45, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1)
        y_offset += 20
        
    return canvas

def draw_hollow_confidence_map(base_img, mask_pred, bldg_probs):
    # Darken the original satellite image for the background
    canvas = np.ascontiguousarray((base_img * 0.6).astype(np.uint8))
    is_bldg = (mask_pred == 1)
    
    # Calculate confidence tiers
    mask_high = np.uint8(is_bldg & (bldg_probs >= 0.75))
    mask_med  = np.uint8(is_bldg & (bldg_probs >= 0.25) & (bldg_probs < 0.75))
    mask_low  = np.uint8(is_bldg & (bldg_probs < 0.25))
    
    # Draw Hollow Buildings First
    cv2.drawContours(canvas, cv2.findContours(mask_high, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0], -1, (255, 0, 0), 3) 
    cv2.drawContours(canvas, cv2.findContours(mask_med, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0], -1, (255, 165, 0), 3) 
    cv2.drawContours(canvas, cv2.findContours(mask_low, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0], -1, (255, 0, 255), 3) 

    # Paint Solid Terrain Underneath
    road_mask = np.uint8(mask_pred == 2)
    canvas[road_mask == 1] = [255, 255, 0]   # Yellow Roads
    canvas[mask_pred == 3] = [0, 100, 255]   # Blue Water
    canvas[mask_pred == 4] = [50, 205, 50]   # Green Open Area
    
    # ✨ Add the Legend to the final canvas before returning
    canvas = draw_legend(canvas)
    
    return canvas

# ==========================================
# 3. THE API ENDPOINTS
# ==========================================
@app.get("/")
def health_check():
    return {"status": "online", "model": "SegFormer-B4", "device": DEVICE}

@app.post("/predict")
async def predict_segmentation(file: UploadFile = File(...)):
    # 1. Read the uploaded image into memory
    image_bytes = await file.read()
    raw_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB").resize((512, 512), Image.BILINEAR)
    img_np = np.array(raw_pil)
    
    # 2. Prepare for the AI
    tensor_img = TF.to_tensor(raw_pil)
    tensor_img = TF.normalize(tensor_img, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    tensor_img = tensor_img.unsqueeze(0).to(DEVICE)

    # 3. Run Inference & Extract Probabilities
    with torch.no_grad():
        logits = model(tensor_img)
        mask_pred = torch.argmax(logits, dim=1).squeeze().cpu().numpy()
        probs = F.softmax(logits, dim=1).squeeze().cpu().numpy()
        bldg_probs = probs[1, :, :] # Extract Class 1 (Buildings) probability map

    # 4. Generate the proper tactical map (now with the legend baked in)
    output_canvas = draw_hollow_confidence_map(img_np, mask_pred, bldg_probs)

    # 5. Compress back to a PNG and send to the user
    is_success, buffer = cv2.imencode(".png", cv2.cvtColor(output_canvas, cv2.COLOR_RGB2BGR))
    io_buf = io.BytesIO(buffer)
    
    return Response(content=io_buf.getvalue(), media_type="image/png")