import gradio as gr
import torch
import json
import re
from PIL import Image
import numpy as np
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
from qwen_vl_utils import process_vision_info
import warnings
warnings.filterwarnings("ignore")

# ============================================
# MODEL CONFIGURATION
# ============================================
MODEL_ID = "zackriya/diagram2graph"
MAX_PIXELS = 1280 * 28 * 28
MIN_PIXELS = 256 * 28 * 28

print("🔄 Loading Flowchart Extractor (diagram2graph)...")
print("⏳ This may take 1-2 minutes on first load...")

# Load model on GPU if available, otherwise CPU
device_map = "auto" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    device_map=device_map,
    torch_dtype=torch_dtype
)

processor = Qwen2_5_VLProcessor.from_pretrained(
    MODEL_ID,
    min_pixels=MIN_PIXELS,
    max_pixels=MAX_PIXELS
)

print(f"✅ Model loaded successfully on: {device_map.upper()}")
if torch.cuda.is_available():
    print(f"   GPU: {torch.cuda.get_device_name(0)}")

# System prompt for structured extraction
SYSTEM_MESSAGE = """You are a Vision Language Model specialized in extracting structured data from visual representations of process and flow diagrams.
Your task is to analyze the provided image of a diagram and extract the relevant information into a well-structured JSON format.
The diagram includes details such as nodes and edges. each of them have their own attributes.
Focus on identifying key data fields and ensuring the output adheres to the requested JSON structure.
Provide only the JSON output based on the extracted information. Avoid additional explanations or comments."""


# ============================================
# SHAPE CLASSIFICATION (Heuristic)
# ============================================
def classify_shape(node_text, node_type):
    """
    Classify shape based on node type and text content
    """
    text_lower = node_text.lower() if node_text else ""
    type_lower = node_type.lower() if node_type else ""
    
    # Decision nodes (diamonds)
    if type_lower == "decision" or "?" in text_lower or "is " in text_lower or "if" in text_lower:
        return "Diamond (Decision)"
    
    # Start/End nodes (ovals/circles)
    if type_lower in ["start", "end", "terminator"] or text_lower in ["start", "stop", "begin", "end"]:
        return "Circle/Oval (Terminator)"
    
    # Input/Output nodes (parallelograms)
    if "read" in text_lower or "print" in text_lower or "output" in text_lower or "input" in text_lower:
        return "Parallelogram (Input/Output)"
    
    # Process nodes (rectangles) - default
    return "Rectangle (Process)"


# ============================================
# POST-PROCESSING FUNCTIONS
# ============================================
def extract_json_from_text(text):
    """
    Extract JSON from model output text
    """
    # Find JSON-like structure in the text
    json_match = re.search(r'\{.*\}|\[.*\]', text, re.DOTALL)
    if json_match:
        return json_match.group()
    return text


def parse_flowchart_output(raw_output):
    """
    Parse and validate the flowchart extraction output
    """
    try:
        # First, extract JSON
        json_str = extract_json_from_text(raw_output)
        
        # Try to parse as JSON
        data = json.loads(json_str)
        
        # Ensure required fields exist
        if "nodes" not in data:
            data["nodes"] = []
        if "edges" not in data:
            data["edges"] = []
        
        return data, None
        
    except json.JSONDecodeError as e:
        # If JSON parsing fails, try to salvage
        return {
            "nodes": [],
            "edges": [],
            "raw_output": raw_output,
            "error": f"JSON parse error: {str(e)}"
        }, str(e)


# ============================================
# MAIN ANALYSIS FUNCTION
# ============================================
def analyze_flowchart(image):
    """
    Extract flowchart structure: nodes, edges, shapes, paths
    """
    if image is None:
        return ("Please upload an image first", "{}")
    
    try:
        # Convert to PIL if needed
        if isinstance(image, np.ndarray):
            image_pil = Image.fromarray(image).convert("RGB")
        else:
            image_pil = image
        
        # Prepare messages for the model
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": SYSTEM_MESSAGE}],
            },
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_pil},
                    {"type": "text", "text": "Extract data in JSON format. Only give the JSON."},
                ],
            },
        ]
        
        # Process with the model
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, _ = process_vision_info(messages)
        
        inputs = processor(
            text=[text],
            images=image_inputs,
            return_tensors="pt",
        )
        
        # Move to appropriate device
        if torch.cuda.is_available():
            inputs = inputs.to('cuda')
        
        # Generate
        with torch.no_grad():
            generated_ids = model.generate(**inputs, max_new_tokens=2048)
        
        # Trim input tokens from output
        generated_ids_trimmed = [
            out_ids[len(in_ids):]
            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        
        output_text = processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False
        )[0]
        
        # Parse the output
        data, parse_error = parse_flowchart_output(output_text)
        
        nodes = data.get("nodes", [])
        edges = data.get("edges", [])
        
        # Add shape classification to each node
        for i, node in enumerate(nodes):
            node_text = node.get("text", "")
            node_type = node.get("type", "")
            nodes[i]["shape"] = classify_shape(node_text, node_type)
        
        # Count shape distribution
        shape_counts = {}
        for node in nodes:
            shape = node.get("shape", "Unknown")
            shape_counts[shape] = shape_counts.get(shape, 0) + 1
        
        # Build output text
        output_md = "## Flowchart Analysis Complete!\n\n"
        output_md += "### Summary\n"
        output_md += f"- Total Nodes: {len(nodes)}\n"
        output_md += f"- Total Edges/Paths: {len(edges)}\n"
        output_md += f"- Model Device: {device_map.upper()}\n\n"
        
        output_md += "### Shape Distribution\n"
        for shape, count in shape_counts.items():
            output_md += f"- {shape}: {count}\n"
        
        if nodes:
            output_md += f"\n### Nodes ({len(nodes)} detected)\n\n"
            for i, node in enumerate(nodes, 1):
                node_text = node.get("text", "Unnamed")[:60]
                node_shape = node.get("shape", "Unknown")
                output_md += f"{i}. **{node_text}** -> {node_shape}\n"
        
        if edges:
            output_md += f"\n### Connections/Paths ({len(edges)} detected)\n\n"
            for i, edge in enumerate(edges, 1):
                from_id = edge.get("from", "?")
                to_id = edge.get("to", "?")
                label = edge.get("label", "")
                if label:
                    output_md += f"{i}. {from_id} -> {to_id} (label: {label})\n"
                else:
                    output_md += f"{i}. {from_id} -> {to_id}\n"
        
        if len(nodes) == 0 and len(edges) == 0:
            output_md += "\nNo nodes or edges detected. Try a clearer image.\n"
        
        # Prepare JSON output
        json_output = {
            "success": True,
            "total_nodes": len(nodes),
            "total_edges": len(edges),
            "shape_distribution": shape_counts,
            "nodes": nodes,
            "edges": edges
        }
        
        return output_md, json.dumps(json_output, indent=2)
        
    except Exception as e:
        import traceback
        error_msg = "Error During Analysis\n\n"
        error_msg += f"Error: {str(e)}\n\n"
        error_msg += "Troubleshooting:\n"
        error_msg += "1. Make sure your Space has GPU enabled\n"
        error_msg += "2. Try a different image\n"
        error_msg += "3. Refresh and try again\n"
        return error_msg, json.dumps({"error": str(e)}, indent=2)


# ============================================
# GRADIO INTERFACE
# ============================================
with gr.Blocks(title="Flowchart & Diagram Analyzer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Flowchart & Diagram Analyzer
    
    Extract structured data from flowchart and diagram images including nodes, shapes, connections, and paths.
    
    ### What This Tool Does
    - Detects all shapes in your flowchart (rectangles, diamonds, circles, etc.)
    - Counts total nodes - every shape/element in the diagram
    - Extracts paths - all connections between nodes
    - Reads text inside shapes
    - Outputs JSON for programmatic use
    
    ### How to Use
    1. Upload your flowchart image below
    2. Click "Analyze Flowchart"
    3. View results in the output panels
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                type="pil", 
                label="Upload Flowchart/Diagram Image",
                height=400
            )
            analyze_btn = gr.Button("Analyze Flowchart", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output_text = gr.Markdown(
                label="Analysis Results",
                value="Upload a flowchart and click 'Analyze Flowchart' to see results here."
            )
    
    with gr.Row():
        output_json = gr.Code(
            label="Raw JSON Output",
            language="json",
            lines=20,
            value="{}"
        )
    
    analyze_btn.click(
        fn=analyze_flowchart,
        inputs=input_image,
        outputs=[output_text, output_json]
    )
    
    gr.Markdown("""
    ---
    ### Understanding the Output
    
    | Field | Description |
    |-------|-------------|
    | Nodes | Each shape/box in your flowchart |
    | Edges | Arrows/connections showing the flow between nodes |
    | Shape | Classification of each node (Rectangle, Diamond, Circle, etc.) |
    | Paths | The complete flow routes through your diagram |
    
    ### Troubleshooting
    
    Issue: Model not loading
    - Make sure your Space has GPU enabled: Settings -> Hardware -> T4 Small
    - First load may take 1-2 minutes to download the model
    
    Issue: Poor detection results
    - Try a clearer image
    - Ensure your flowchart has good contrast
    - Simple, clean flowcharts work best
    
    Built with: diagram2graph by Zackriya Solutions | Gradio Interface
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)