import cv2
import numpy as np
import torch
import pytesseract
from PIL import Image, ImageGrab
from transformers import CLIPProcessor, CLIPModel
import time

# --- ASETUKSET ---
VIEWPORT_TOP = 120 # Ohitetaan selaimen ylC$palkki (CDP/Tabs)
ICON_LABELS = [
    "a paperclip icon", "a microphone icon", "a smiley face icon", 
    "a send message button", "a user profile picture", 
    "a magnifying glass search icon", "random background noise"
]

def load_clip():
    print("[*] 1. Ladataan semanttinen Vision-malli (CLIP)...")
    model_id = "openai/clip-vit-base-patch32"
    return CLIPModel.from_pretrained(model_id), CLIPProcessor.from_pretrained(m
odel_id)

def detect_frames(gray_img):
    print("[*] 2. Analysoidaan kC$yttC6liittymC$n pC$C$rakenteet (Layout)...")
    screen_h, screen_w = gray_img.shape
    viewport_gray = gray_img[VIEWPORT_TOP:, :]
    
    edges = cv2.Canny(viewport_gray, 50, 150)
    col_sums = np.sum(edges, axis=0)
    
    search_start = int(screen_w * 0.20)
    search_end = int(screen_w * 0.50)
    middle_cols = col_sums[search_start:search_end]
    
    divider_x = search_start + np.argmax(middle_cols)
    if col_sums[divider_x] < (screen_h * 0.1 * 255): 
        divider_x = int(screen_w * 0.30)
        
    return [
        {"id": "Frame_1_Sidebar", "x": 0, "y": VIEWPORT_TOP, "w": divider_x,
"h": screen_h - VIEWPORT_TOP},
        {"id": "Frame_2_Main", "x": divider_x, "y": VIEWPORT_TOP, "w":
screen_w - divider_x, "h": screen_h - VIEWPORT_TOP}
    ]

def extract_nodes(gray_img, frame):
    fx, fy, fw, fh = frame["x"], frame["y"], frame["w"], frame["h"]
    roi_gray = gray_img[fy:fy+fh, fx:fx+fw]
    
    edges = cv2.Canny(roi_gray, 50, 150)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (10, 5))
    dilated = cv2.dilate(edges, kernel, iterations=1)
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
    
    nodes = []
    for i, c in enumerate(contours):
        x, y, w, h = cv2.boundingRect(c)
        if w < 12 or h < 12 or w > fw * 0.9: continue
            
        crop = roi_gray[y:y+h, x:x+w]
        text = pytesseract.image_to_string(crop, config='--psm 7').strip()
        text = "".join([ch for ch in text if ch.isalnum() or ch in "
.,?!-()/:"])
        
        abs_x, abs_y = fx + x, fy + y
        
        # Jos teksti on riittC$vC$n pitkC$, se on tekstiC$. Muuten ikoni.
        if len(text) > 2:
            nodes.append({"id": f"{frame['id']}_n{i}", "type": "TextNode",
"content": text, "box": (abs_x, abs_y, w, h)})
        else:
            nodes.append({"id": f"{frame['id']}_n{i}", "type": "IconNode",
"content": "Unknown", "box": (abs_x, abs_y, w, h)})
            
    # JC$rjestetC$C$n ylhC$C$ltC$ alas
    nodes.sort(key=lambda n: (n['box'][1] // 20, n['box'][0]))
    return nodes

def enrich_icons_with_clip(pil_img, nodes, clip_model, clip_processor):
    print(f"[*] 3. Rikastetaan ikonit tekoC$lyllC$ (Semanttinen analyysi)...")
    icon_nodes = [n for n in nodes if n['type'] == 'IconNode']
    if not icon_nodes: return nodes
    
    crops = []
    for node in icon_nodes:
        x, y, w, h = node['box']
        pad = 4 # Pieni marginaali auttaa CLIPiC$ nC$kemC$C$n ikonin rajat
        box = (max(0, x-pad), max(0, y-pad), min(pil_img.width, x+w+pad),
min(pil_img.height, y+h+pad))
        crops.append(pil_img.crop(box).convert("RGB"))
        
    # Batch-ajo: Kaikki ikonit kerralla CLIPin lC$pi
    inputs = clip_processor(text=ICON_LABELS, images=crops, return_tensors="pt"
, padding=True)
    with torch.no_grad():
        outputs = clip_model(**inputs)
        
    logits = outputs.logits_per_image.softmax(dim=1).numpy()
    
    # YhdistetC$C$n tulokset takaisin solmuihin
    for i, probs in enumerate(logits):
        best_idx = np.argmax(probs)
        best_score = probs[best_idx] * 100
        best_label = ICON_LABELS[best_idx]
        
        if best_score > 50.0 and "noise" not in best_label:
            icon_nodes[i]['content'] = f"{best_label} ({best_score:.0f}%)"
        else:
            icon_nodes[i]['content'] = "Generic Icon"
            
    return nodes

def generate_markdown(frames, all_nodes):
    print("[*] 4. Rakennetaan AST-MD semanttinen kartta...")
    md = "# AST-MD: Rikastettu KC$yttC6liittymC$kartta

"
    
    for frame in frames:
        md += f"## {frame['id']}
"
        md += f"> Rajat: X:{frame['x']}, Y:{frame['y']}, W:{frame['w']},
H:{frame['h']}

"
        
        frame_nodes = [n for n in all_nodes if n['id'].startswith(frame['id'])]
        for n in frame_nodes:
            x, y, w, h = n['box']
            if n['type'] == 'TextNode':
                md += f"  - `[{n['id']}]` **Text**: \"{n['content']}\" (X:{x},
Y:{y})
"
            else:
                md += f"  - `[{n['id']}]` **Icon**: *{n['content']}* (X:{x},
Y:{y})
"
        md += "
"
    return md

def main():
    start_time = time.time()
    clip_model, clip_processor = load_clip()
    
    pil_img = ImageGrab.grab()
    cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    gray_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)