from datetime import datetime

from common.config_loader import loadVisionConfig
from common.runtime_log import appendRuntimeLog
from screen.screen_actions import captureOcrRegion

_vision_pipeline = None


def log(level, message):
    now = datetime.now().strftime("%H:%M:%S")
    print(f"[{level}] {now} {message}", flush=True)


def getVisionPipeline():
    global _vision_pipeline

    if _vision_pipeline is not None:
        return _vision_pipeline

    config = loadVisionConfig()
    model_name = config["model_name"]

    if not model_name:
        raise RuntimeError("vision.model_name is not set")

    try:
        from transformers import AutoProcessor
        from transformers import Gemma4ForConditionalGeneration
        import torch
    except Exception as e:
        raise RuntimeError(f"vision dependencies are missing detail={e}")

    log("INFO", f"Vision model initializing model={model_name}")
    processor = AutoProcessor.from_pretrained(model_name)
    model = Gemma4ForConditionalGeneration.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
    ).eval()
    _vision_pipeline = (processor, model)
    return _vision_pipeline


def saveVisionText(image_path, text):
    return appendRuntimeLog(
        "VISION TEXT",
        f"image={image_path}\n{text}",
    )


def runVisionFromImage(image_path):
    config = loadVisionConfig()
    processor, model = getVisionPipeline()

    prompt = config["prompt"]
    messages = [
        {
            "role": "system",
            "content": [
                {"type": "text", "text": "You are a helpful vision-language assistant."},
            ],
        },
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": prompt},
            ],
        }
    ]

    try:
        import torch
        from PIL import Image
    except Exception as e:
        log("ERROR", f"Vision runtime imports failed detail={e}")
        return None

    try:
        image = Image.open(image_path).convert("RGB")
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        )
        inputs = inputs.to(model.device, dtype=model.dtype)
        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=config["max_new_tokens"],
                do_sample=config["temperature"] > 0,
            )
        input_len = inputs["input_ids"].shape[-1]
        generated = output[0][input_len:]
        text = processor.decode(generated, skip_special_tokens=True).strip()
    except Exception as e:
        log("ERROR", f"Vision inference failed detail={e}")
        return None

    text_path = saveVisionText(image_path, text)
    log("INFO", f"Vision text appended={text_path}")
    print(text, flush=True)
    return {
        "image_path": image_path,
        "text_path": text_path,
        "text": text,
    }


def runVisionFromScreen():
    log("ACTION", "runVisionFromScreen called")
    image_path = captureOcrRegion()
    return runVisionFromImage(image_path)