import os
import sys

# --- CRITICAL SYSTEM-LEVEL FIXES: FORCE 2 THREAD LIMITS ---
# These MUST be set before importing torch or transformers to block thread explosions
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
os.environ["OPENBLAS_NUM_THREADS"] = "2"
os.environ["RAYON_NUM_THREADS"] = "2"  # Blocks Rust safetensors threading spikes
os.environ["HF_HUB_OFFLINE"] = "0"

import re
import time
import gc
import json
import io
import psutil
import torch
import traceback
from fastapi import FastAPI, File
from pydantic import BaseModel
from transformers import AutoProcessor, AutoModelForVision2Seq, AutoConfig
from PIL import Image, ImageEnhance, ImageFilter
from contextlib import asynccontextmanager

# Force print statements to write to the console instantly
sys.stdout.reconfigure(line_buffering=True)

# Lock PyTorch's internal execution threadpool
torch.set_num_threads(2)

NORM_SIZE = 500
MODEL_ID = "docling-project/ScreenVLM"

def log_memory(label: str):
    """Utility to print precise, real-time RAM usage of the Python process."""
    process = psutil.Process(os.getpid())
    mem_mb = process.memory_info().rss / (1024 * 1024)
    print(f"[MEMORY DIAGNOSTIC] {label} -> Current process RAM: {mem_mb:.2f} MB", flush=True)

log_memory("0. Server Process Started")

class ScreenVLMEngine:
    def __init__(self):
        log_memory("1. Initializing ScreenVLMEngine")
        
        print("2. Loading processor from Hugging Face Hub...", flush=True)
        self.processor = AutoProcessor.from_pretrained(MODEL_ID)
        
        # Inject chat template fallback if not defined in tokenizer_config
        self.processor.chat_template = (
            "{% for message in messages %}"
            "{{ '<|start_of_role|>' + message['role'] + '<|end_of_role|>' }}"
            "{% for content in message['content'] %}"
            "{% if content['type'] == 'image' %}{{ '<image>' }}"
            "{% elif content['type'] == 'text' %}{{ content['text'] }}"
            "{% endif %}{% endfor %}"
            "{{ '<end_of_utterance>' + '\n' }}"
            "{% endfor %}"
            "{% if add_generation_prompt %}{{ '<|start_of_role|>assistant<|end_of_role|>' }}{% endif %}"
        )
        if hasattr(self.processor, "tokenizer"):
            self.processor.tokenizer.chat_template = self.processor.chat_template

        log_memory("3. Processor loaded. Preparing config...")
        config = AutoConfig.from_pretrained(MODEL_ID)
        config.use_cache = True
        config.text_config.use_cache = True

        log_memory("4. Config prepared. Starting file-load of weights on CPU...")
        t_load = time.time()
        
        # Load the model directly in float32 using 'eager' attention
        self.model = AutoModelForVision2Seq.from_pretrained(
            MODEL_ID,
            config=config,                 
            torch_dtype=torch.float32,     
            attn_implementation="eager",   
            low_cpu_mem_usage=True          
        )
        
        # Resize the model token embeddings to accommodate the tokenizer's special tokens
        if hasattr(self.processor, "tokenizer"):
            vocab_size = len(self.processor.tokenizer)
            print(f"[DIAGNOSTIC] Resizing token embeddings to: {vocab_size}", flush=True)
            self.model.resize_token_embeddings(vocab_size)
            
        print(f"5. File-load complete in {time.time() - t_load:.3f}s.", flush=True)
        log_memory("6. Model is AWAKE and ready in RAM (No quantization).")

    def parse_screentag(self, text: str, width: int, height: int):
        pattern = re.compile(
            r"<(?P<tag>[a-zA-Z][a-zA-Z0-9_]*)>"
            r"\s*<loc_(?P<l>\d+)><loc_(?P<t>\d+)><loc_(?P<r>\d+)><loc_(?P<b>\d+)>"
            r"(?P<text>[^<]*)"
        )
        elements = []
        for m in pattern.finditer(text):
            l, t, r, b = [max(0, min(int(m.group(k)), NORM_SIZE)) for k in ("l", "t", "r", "b")]
            if r < l: l, r = r, l
            if b < t: t, b = b, t
                
            x1 = int(l / NORM_SIZE * width)
            y1 = int(t / NORM_SIZE * height)
            x2 = int(r / NORM_SIZE * width)
            y2 = int(b / NORM_SIZE * height)
            
            elements.append({
                "label": m.group("tag"),
                "bbox": [x1, y1, x2, y2],
                "text": m.group("text").strip() or None,
            })
        return elements

    def analyze(self, image: Image.Image):
        orig_width, orig_height = image.size
        
        # Resize image safely to match native processor maximum bounds
        max_edge = 2048
        if max(orig_width, orig_height) > max_edge:
            scale = max_edge / float(max(orig_width, orig_height))
            new_w = int(orig_width * scale)
            new_h = int(orig_height * scale)
            image_to_process = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
        else:
            image_to_process = image
        
        # [NEW ENHANCEMENT STEP]
        # 1. Sharpen the image to define thin boundaries between small icons
        image_to_process = image_to_process.filter(ImageFilter.SHARPEN)
        
        # 2. Boost the contrast (factor of 1.5) to make text and borders stand out
        contrast_enhancer = ImageEnhance.Contrast(image_to_process)
        image_to_process = contrast_enhancer.enhance(1.1)
        
        # 3. Boost sharpness (factor of 2.0) to make adjacent icons visually separate
        sharpness_enhancer = ImageEnhance.Sharpness(image_to_process)
        image_to_process = sharpness_enhancer.enhance(1.4)
        
        prompt = self.processor.apply_chat_template([
            {
                "role": "user", 
                "content": [
                    {"type": "image"}, 
                    {"type": "text", "text": "Generate the screen representation for this UI:"}
                ]
            }
        ], tokenize=False, add_generation_prompt=True)
        
        log_memory("Inference Step A: Before Preprocessing")
        inputs = self.processor(text=prompt, images=[image_to_process], return_tensors="pt")
        log_memory("Inference Step B: Tensors Allocated")

        print("[DIAGNOSTIC] --- INPUT TENSORS DUMP ---", flush=True)
        for k, v in inputs.items():
            if torch.is_tensor(v):
                print(f"  Tensor '{k}': shape={list(v.shape)}, dtype={v.dtype}", flush=True)
            else:
                print(f"  Non-Tensor '{k}': type={type(v)}", flush=True)
        print("[DIAGNOSTIC] --------------------------", flush=True)

        print("Executing model.generate()...", flush=True)
        t0 = time.time()
        with torch.inference_mode():
            generated_ids = self.model.generate(
                **inputs, 
                max_new_tokens=1024,
                use_cache=True,                  
                eos_token_id=100257,            
                pad_token_id=100257,
                #repetition_penalty=1.15,      # Penalizes sequential tag repetition
                #num_beams=2,                  # Uses Beam Search instead of Greedy Decoding (slower but more precise)
                #no_repeat_ngram_size=4        # Prevents redundant coordinate loops

            )
            
        print(f"Inference took: {time.time() - t0:.3f}s", flush=True)
        
        prompt_length = inputs.input_ids.shape[1]
        raw_output = self.processor.batch_decode(
            generated_ids[:, prompt_length:], 
            skip_special_tokens=False
        )[0].lstrip()
        
        parsed_elements = self.parse_screentag(raw_output, orig_width, orig_height)
        return parsed_elements, raw_output

# --- FastAPI Server Setup ---
engine = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global engine
    os.environ["HF_HUB_OFFLINE"] = "0"
    engine = ScreenVLMEngine()
    os.environ["HF_HUB_OFFLINE"] = "1"
    yield
    engine = None

app = FastAPI(lifespan=lifespan)

class ParseRequest(BaseModel):
    image_path: str

@app.post("/parse_screen")
def parse_screen(req: ParseRequest):
    return {"status": "error", "message": "Use /parse_image with raw image bytes."}

@app.post("/parse_image")
async def parse_image(file: bytes = File(...)):
    try:
        gc.collect()
        log_memory("FastAPI Route: Request Received")
        image = Image.open(io.BytesIO(file)).convert("RGB")
        elements, raw_output = engine.analyze(image)
        del image
        gc.collect()
        log_memory("FastAPI Route: Response Ready")
        return {"status": "success", "elements": elements, "raw_output": raw_output}
    except Exception as e:
        traceback.print_exc()
        print(f"[CRITICAL ROUTE ERROR] {str(e)}", flush=True)
        return {"status": "error", "message": str(e)}

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