from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from PIL import Image
from io import BytesIO
import base64
import torch
import re
import logging
import asyncio
from contextlib import asynccontextmanager

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize global variables
model = None
processor = None
tokenizer = None
model_name = "microsoft/GUI-Actor-2B-Qwen2-VL"
model_loaded = False

async def load_model():
    """Load model with proper error handling and fallback strategies"""
    global model, processor, tokenizer, model_loaded
    
    try:
        logger.info("Starting model loading...")
        
        # Try specific Qwen2VL classes first
        try:
            logger.info("Attempting to load with Qwen2VL specific classes...")
            from transformers import Qwen2VLProcessor, Qwen2VLForConditionalGeneration
            
            processor = Qwen2VLProcessor.from_pretrained(
                model_name, 
                trust_remote_code=True
            )
            
            model = Qwen2VLForConditionalGeneration.from_pretrained(
                model_name,
                torch_dtype=torch.float32,
                device_map=None,  # CPU only
                trust_remote_code=True,
                low_cpu_mem_usage=True
            ).eval()
            
            logger.info("Successfully loaded with Qwen2VL specific classes")
            
        except Exception as e1:
            logger.warning(f"Failed with Qwen2VL classes: {e1}")
            logger.info("Trying AutoProcessor and AutoModel fallback...")
            
            try:
                from transformers import AutoProcessor, AutoModel
                
                processor = AutoProcessor.from_pretrained(
                    model_name, 
                    trust_remote_code=True
                )
                
                model = AutoModel.from_pretrained(
                    model_name,
                    torch_dtype=torch.float32,
                    device_map=None,
                    trust_remote_code=True,
                    low_cpu_mem_usage=True
                ).eval()
                
                logger.info("Successfully loaded with Auto classes")
                
            except Exception as e2:
                logger.warning(f"Failed with Auto classes: {e2}")
                logger.info("Trying generic transformers approach...")
                
                # Last fallback - try loading as generic model
                from transformers import AutoConfig, AutoTokenizer
                import transformers
                
                config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
                logger.info(f"Model config type: {type(config)}")
                
                # Try to find the right model class
                if hasattr(transformers, 'Qwen2VLForConditionalGeneration'):
                    ModelClass = getattr(transformers, 'Qwen2VLForConditionalGeneration')
                elif hasattr(transformers, 'AutoModelForVision2Seq'):
                    ModelClass = getattr(transformers, 'AutoModelForVision2Seq')
                else:
                    raise Exception("No suitable model class found")
                
                processor = AutoProcessor.from_pretrained(
                    model_name, 
                    trust_remote_code=True
                )
                
                model = ModelClass.from_pretrained(
                    model_name,
                    config=config,
                    torch_dtype=torch.float32,
                    device_map=None,
                    trust_remote_code=True,
                    low_cpu_mem_usage=True
                ).eval()
        
        # Verify processor and model are loaded
        if processor is None or model is None:
            raise Exception("Failed to load processor or model")
            
        tokenizer = processor.tokenizer
        logger.info("Model and processor loaded successfully!")
        model_loaded = True
        return True
        
    except Exception as e:
        logger.error(f"Error loading model: {e}")
        model_loaded = False
        return False

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    logger.info("Starting up GUI-Actor API...")
    await load_model()
    yield
    # Shutdown
    logger.info("Shutting down GUI-Actor API...")

# Initialize FastAPI app with lifespan
app = FastAPI(
    title="GUI-Actor API", 
    version="1.0.0",
    lifespan=lifespan
)

class Base64Request(BaseModel):
    image_base64: str
    instruction: str

def extract_coordinates(text):
    """
    Extract coordinates from model output text
    """
    # Pattern untuk mencari koordinat dalam berbagai format
    patterns = [
        r'click\s*\(\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)',  # click(x, y)
        r'\[\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\]',          # [x, y]
        r'(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)',                    # x, y
        r'point:\s*\(\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)', # point: (x, y)
    ]
    
    for pattern in patterns:
        matches = re.findall(pattern, text.lower())
        if matches:
            try:
                x, y = float(matches[0][0]), float(matches[0][1])
                # Normalize jika koordinat > 1 (asumsi pixel coordinates)
                if x > 1 or y > 1:
                    # Asumsi resolusi 1920x1080 untuk normalisasi
                    x = x / 1920 if x > 1 else x
                    y = y / 1080 if y > 1 else y
                return [(x, y)]
            except (ValueError, IndexError):
                continue
    
    # Default ke center jika tidak ditemukan
    return [(0.5, 0.5)]

def cpu_inference(conversation, model, tokenizer, processor):
    try:
        prompt = processor.apply_chat_template(
            conversation, 
            tokenize=False,
            add_generation_prompt=True
        )

        image = conversation[1]["content"][0]["image"]
        inputs = processor(
            text=[prompt],
            images=[image],
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        )

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                do_sample=True,
                temperature=0.3,
                top_p=0.8,
                pad_token_id=tokenizer.eos_token_id or tokenizer.pad_token_id or 0
            )

        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
        response = tokenizer.decode(generated_ids, skip_special_tokens=True)
        coordinates = extract_coordinates(response)

        return {
            "topk_points": coordinates,
            "response": response,
            "success": True
        }

    except Exception as e:
        logger.error(f"Inference error: {e}")
        return {
            "topk_points": [(0.5, 0.5)],
            "response": f"Error during inference: {str(e)}",
            "success": False
        }

@app.get("/")
async def root():
    return {
        "message": "GUI-Actor API is running", 
        "status": "healthy",
        "model_loaded": model_loaded,
        "model_name": model_name
    }

@app.post("/click/base64")
async def predict_click_base64(data: Base64Request):
    if not model_loaded:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded properly"
        )
    
    try:
        # Decode base64 to image
        try:
            # Handle data URL format
            if "," in data.image_base64:
                image_data = base64.b64decode(data.image_base64.split(",")[-1])
            else:
                image_data = base64.b64decode(data.image_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 image: {e}")
        
        try:
            pil_image = Image.open(BytesIO(image_data)).convert("RGB")
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid image format: {e}")

        conversation = [
            {
                "role": "system",
                "content": [
                    {
                        "type": "text",
                        "text": "You are a GUI agent. You are given a task and a screenshot of the screen. You need to perform a series of pyautogui actions to complete the task. Please provide the click coordinates.",
                    }
                ]
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": pil_image,
                    },
                    {
                        "type": "text",
                        "text": data.instruction,
                    },
                ],
            },
        ]

        # Run inference
        pred = cpu_inference(conversation, model, tokenizer, processor)
        px, py = pred["topk_points"][0]
        
        return JSONResponse(content={
            "x": round(px, 4), 
            "y": round(py, 4),
            "response": pred["response"],
            "success": pred["success"]
        })
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        raise HTTPException(
            status_code=500,
            detail=f"Internal server error: {str(e)}"
        )

@app.get("/health")
async def health_check():
    return {
        "status": "healthy" if model_loaded else "unhealthy",
        "model": model_name,
        "device": "cpu",
        "torch_dtype": "float32",
        "model_loaded": model_loaded
    }

@app.get("/debug")
async def debug_info():
    """Debug endpoint to check model loading status"""
    import transformers
    available_classes = [attr for attr in dir(transformers) if 'Qwen' in attr or 'VL' in attr]
    
    return {
        "model_loaded": model_loaded,
        "processor_type": type(processor).__name__ if processor else None,
        "model_type": type(model).__name__ if model else None,
        "available_qwen_classes": available_classes,
        "transformers_version": transformers.__version__
    }