import io
import torch
import torchvision
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import HTMLResponse
from huggingface_hub import hf_hub_download
from PIL import Image

app = FastAPI()

# --- REPLICATING EXACT ARCHITECTURE FROM TRAIN.PY ---
class ResNet18(torch.nn.Module):
    def __init__(self):
        super().__init__()
        model = torchvision.models.resnet18()

        model.layer3[0].conv1.stride = (2, 1)
        model.layer3[0].downsample[0].stride = (2, 1)

        model.layer4[0].conv1.stride = (1, 1)
        model.layer4[0].downsample[0].stride = (1, 1)

        self.conv1 = model.conv1
        self.bn1 = model.bn1
        self.relu = model.relu
        self.maxpool = model.maxpool
        self.layer1 = model.layer1
        self.layer2 = model.layer2
        self.layer3 = model.layer3
        self.layer4 = model.layer4

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        return self.layer4(self.layer3(self.layer2(self.layer1(self.maxpool(self.relu(self.bn1(self.conv1(X))))))))

class SeqToMap(torch.nn.Module):
    def __init__(self, num_channels: int, height: int, out_features: int):
        super().__init__()
        self.project = torch.nn.Linear(in_features=num_channels * height, out_features=out_features)

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        B, C, H, W = X.size()
        X = X.permute(0, 3, 1, 2).reshape(B, W, C * H)
        return self.project(X)

class BiLSTM(torch.nn.Module):
    def __init__(self, input_size: int, hidden_size: int, num_layers: int, dropout: float):
        super().__init__()
        self.rnn = torch.nn.LSTM(
            input_size=input_size, hidden_size=hidden_size, num_layers=num_layers,
            batch_first=True, bidirectional=True, dropout=dropout
        )

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        X, _ = self.rnn(X)
        return X

class ResNet18_CRNN(torch.nn.Module):
    def __init__(self, input_height: int, compression_size: int, hidden_size_bilstm: int, num_layers_bilstm: int, num_class: int, dropout_bilstm: float):
        super().__init__()
        self.resnet18 = ResNet18()
        self.seqtomap = SeqToMap(num_channels=512, height=input_height // 16, out_features=compression_size)
        self.bilstm = BiLSTM(input_size=compression_size, hidden_size=hidden_size_bilstm, num_layers=num_layers_bilstm, dropout=dropout_bilstm)
        self.fc = torch.nn.Linear(in_features=hidden_size_bilstm * 2, out_features=num_class)

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.log_softmax(self.fc(self.bilstm(self.seqtomap(self.resnet18(X)))), dim=2)

# --- LOAD WEIGHTS AT STARTUP ---
# Update this string if your character mapping differs from this default standard
CHARACTERS = ['<BLANK>', ' ', '!', '"', '#', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']

def load_ocr_model():
    weights_path = hf_hub_download(repo_id="Shad0wKillar/resnet18-crnn-ocr", filename="model18_crnn_ocr.pth")
    
    model = ResNet18_CRNN(
        input_height=96,
        compression_size=256,
        hidden_size_bilstm=256,
        num_layers_bilstm=2,
        dropout_bilstm=0.5,
        num_class=80
    )
    
    state_dict = torch.load(weights_path, map_location=torch.device("cpu"), weights_only=True)
    # Strip "module." prefix if saved using DataParallel
    clean_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
    model.load_state_dict(clean_state_dict)
    model.eval()
    return model

ocr_model = load_ocr_model()

# --- PREPROCESSING & DECODING ---
def preprocess_image(image_bytes: bytes) -> torch.Tensor:
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    w, h = image.size
    new_w = int(w * (96 / h))
    
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize((96, new_w)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transforms(image).unsqueeze(0)

def ctc_decode(logits: torch.Tensor) -> str:
    # Get the predicted classes directly using argmax
    preds = torch.argmax(logits, dim=2)[0]
    
    tokens = []
    prev_token = None
    for token in preds:
        token_idx = token.item()
        if token_idx != 0 and token_idx != prev_token:
            if token_idx < len(CHARACTERS):
                tokens.append(CHARACTERS[token_idx])
        prev_token = token_idx
        
    return "".join(tokens)

# --- ROUTES & FRONTEND HTML ---
@app.get("/", response_class=HTMLResponse)
async def read_root():
    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <title>ResNet18-CRNN OCR Inference</title>
        <style>
            body { font-family: sans-serif; background: #0b0f19; color: #e5e7eb; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0; }
            .card { background: #1e293b; padding: 30px; border-radius: 12px; width: 400px; text-align: center; box-shadow: 0 4px 10px rgba(0,0,0,0.3); }
            input[type="file"] { display: none; }
            .upload-btn { display: inline-block; padding: 12px 24px; background: #3b82f6; border-radius: 8px; cursor: pointer; font-weight: bold; margin-bottom: 20px; transition: 0.2s; }
            .upload-btn:hover { background: #2563eb; }
            button { width: 100%; padding: 12px; background: #10b981; border: none; border-radius: 8px; color: white; font-weight: bold; cursor: pointer; font-size: 16px; transition: 0.2s; }
            button:hover { background: #059669; }
            #preview { max-width: 100%; height: 96px; object-fit: contain; display: none; margin: 0 auto 20px auto; border: 1px solid #374151; border-radius: 4px; }
            #result { margin-top: 25px; font-size: 24px; font-weight: bold; color: #fbbf24; word-break: break-all; min-height: 30px; }
        </style>
    </head>
    <body>
        <div class="card">
            <h2 style="margin-top:0;">ResNet18 OCR</h2>
            <p style="color: #9ca3af; margin-bottom: 25px;">Upload a line image</p>
            
            <img id="preview" />
            
            <label class="upload-btn">
                Choose Image
                <input type="file" id="imgInput" accept="image/*" onchange="showPreview(event)">
            </label>
            <button onclick="performOcr()">Analyze Text</button>
            <div id="result"></div>
        </div>
        <script>
            function showPreview(event) {
                const file = event.target.files[0];
                if (file) {
                    const reader = new FileReader();
                    reader.onload = function(e) {
                        const img = document.getElementById('preview');
                        img.src = e.target.result;
                        img.style.display = 'block';
                    }
                    reader.readAsDataURL(file);
                }
            }

            async function performOcr() {
                const fileInput = document.getElementById('imgInput');
                const resultDiv = document.getElementById('result');
                if (!fileInput.files[0]) return alert("Please select an image first.");
                
                resultDiv.innerText = "Analyzing...";
                const formData = new FormData();
                formData.append("file", fileInput.files[0]);
                
                try {
                    const res = await fetch("/predict", { method: "POST", body: formData });
                    const data = await res.json();
                    resultDiv.innerText = data.text || "Error processing image";
                } catch(e) {
                    resultDiv.innerText = "Server Error";
                }
            }
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    image_bytes = await file.read()
    img_tensor = preprocess_image(image_bytes)
    
    with torch.no_grad():
        logits = ocr_model(img_tensor)
        
    predicted_text = ctc_decode(logits)
    return {"text": predicted_text}