"""
Diagram-to-Code
Upload a hand-drawn or digital diagram and get generated code.
Supports: Database schemas → SQL, Flowcharts → Python, UML → Classes
"""

import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import base64
from io import BytesIO
import os
import sys

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from shared.components import create_method_panel, create_premium_hero

# Initialize Hugging Face Inference Client. A token improves reliability on Spaces,
# but public inference may still work depending on provider availability.
client = InferenceClient(token=os.getenv("HF_TOKEN") or None)

PROMPTS = {
    "Database Schema → SQL": """Analyze this database schema diagram and generate SQL code.

Look for:
- Tables (rectangles/boxes with table names)
- Columns (listed inside tables)
- Data types (if shown)
- Primary keys (PK, underlined, or marked)
- Foreign keys (arrows, FK, or relationships)
- Relationships (one-to-many, many-to-many)

Generate:
1. CREATE TABLE statements for each table
2. Define columns with appropriate data types
3. Set PRIMARY KEY constraints
4. SET FOREIGN KEY constraints
5. Add indexes for foreign keys
6. Include comments explaining relationships

Use standard SQL (compatible with PostgreSQL/MySQL). Make the code clean and well-commented.""",

    "Flowchart → Python": """Analyze this flowchart diagram and generate Python code.

Look for:
- Start/End (ovals/rounded rectangles)
- Process boxes (rectangles)
- Decision diamonds (conditions)
- Arrows showing flow
- Input/Output (parallelograms)
- Loops

Generate:
1. Python function(s) implementing the logic
2. Use appropriate control structures (if/else, while, for)
3. Add descriptive variable names
4. Include docstrings
5. Add comments for major steps
6. Handle edge cases

Make the code clean, Pythonic, and well-structured.""",

    "UML Class Diagram → Classes": """Analyze this UML class diagram and generate Python class code.

Look for:
- Classes (boxes with class names)
- Attributes (variables in top section)
- Methods (functions in bottom section)
- Relationships (inheritance, composition, association)
- Access modifiers (+public, -private, #protected)
- Multiplicities

Generate:
1. Python class definitions
2. __init__ methods with attributes
3. Methods with appropriate signatures
4. Inheritance using class(BaseClass)
5. Type hints
6. Docstrings for classes and methods
7. Properties where appropriate

Follow Python best practices and PEP 8.""",

    "UI Mockup → HTML/CSS": """Analyze this UI mockup/wireframe and generate HTML & CSS code.

Look for:
- Layout structure (header, nav, main, sidebar, footer)
- Components (buttons, forms, cards, navigation)
- Text elements (headings, paragraphs)
- Images/placeholders
- Colors and styling
- Spacing and alignment

Generate:
1. Semantic HTML5 structure
2. Clean, modern CSS (or Tailwind classes)
3. Responsive design patterns
4. Proper class naming (BEM methodology)
5. Accessibility attributes (ARIA labels)
6. Comments explaining sections
7. Mobile-first approach

Make it robust, readable, and easy to extend with modern best practices.""",

    "Whiteboard/Napkin Sketch → Code": """This appears to be a hand-drawn sketch or whiteboard diagram. Analyze what's drawn and generate appropriate code.

First, identify what this is:
- Database schema? → Generate SQL
- App workflow/flowchart? → Generate Python logic
- UI mockup/wireframe? → Generate HTML/CSS
- Class diagram? → Generate Python classes
- API structure? → Generate API endpoints

Then generate clean, commented code based on the drawing. Be creative in interpreting rough sketches. Add helpful comments explaining what you inferred from the sketch."""
}


def fallback_scaffold(diagram_type, reason):
    """Return a polished deterministic scaffold when live VLM inference is unavailable."""
    templates = {
        "Database Schema → SQL": """```sql
-- Diagram interpretation scaffold
-- Replace the example entities below with the tables, columns, and relationships visible in your diagram.

CREATE TABLE users (
    id BIGSERIAL PRIMARY KEY,
    email TEXT NOT NULL UNIQUE,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE projects (
    id BIGSERIAL PRIMARY KEY,
    owner_id BIGINT NOT NULL REFERENCES users(id),
    name TEXT NOT NULL,
    status TEXT NOT NULL DEFAULT 'active',
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_projects_owner_id ON projects(owner_id);
```
""",
        "Flowchart → Python": """```python
def run_workflow(input_payload: dict) -> dict:
    \"\"\"Implementation scaffold for a flowchart-style process.\"\"\"
    state = {"status": "started", "steps": []}

    if not input_payload:
        return {"status": "rejected", "reason": "missing input"}

    state["steps"].append("validate_input")

    if input_payload.get("requires_review"):
        state["steps"].append("manual_review")
        state["status"] = "pending_review"
    else:
        state["steps"].append("automatic_processing")
        state["status"] = "completed"

    return state
```
""",
        "UML Class Diagram → Classes": """```python
from dataclasses import dataclass, field


@dataclass
class Entity:
    \"\"\"Base class inferred from a UML class diagram.\"\"\"
    id: int


@dataclass
class Project(Entity):
    name: str
    owner_id: int
    tasks: list[str] = field(default_factory=list)

    def add_task(self, task: str) -> None:
        self.tasks.append(task)
```
""",
        "UI Mockup → HTML/CSS": """```html
<main class="dashboard-shell">
  <section class="hero-card">
    <p class="eyebrow">Product Insight</p>
    <h1>Readable UI scaffold from a visual mockup</h1>
    <p>Replace these sections with the components visible in your uploaded wireframe.</p>
    <button>Primary action</button>
  </section>
</main>
```

```css
.dashboard-shell {
  min-height: 100vh;
  padding: 48px;
  background: linear-gradient(135deg, #f6efe4, #dfefff);
  font-family: ui-sans-serif, system-ui;
}

.hero-card {
  max-width: 760px;
  padding: 40px;
  border-radius: 28px;
  background: rgba(255, 255, 255, 0.82);
  box-shadow: 0 24px 80px rgba(20, 32, 54, 0.16);
}
```
""",
        "Whiteboard/Napkin Sketch → Code": """```python
def interpret_sketch(elements: list[dict]) -> dict:
    \"\"\"Starter adapter for turning rough whiteboard elements into implementation tasks.\"\"\"
    entities = [item for item in elements if item.get("kind") == "entity"]
    actions = [item for item in elements if item.get("kind") == "action"]
    relationships = [item for item in elements if item.get("kind") == "relationship"]

    return {
        "entities_to_model": entities,
        "functions_to_write": actions,
        "interfaces_to_define": relationships,
    }
```
""",
    }

    return f"""# {diagram_type}

## Implementation Scaffold

Live vision inference is temporarily unavailable in this runtime, so the Space is returning a professional starter scaffold for the selected diagram type.

{templates.get(diagram_type, templates["Whiteboard/Napkin Sketch → Code"])}

## How To Use This Result

1. Replace the placeholder names with the labels visible in your diagram.
2. Map each box to a table, class, function, component, or process step.
3. Map arrows to dependencies, foreign keys, inheritance, or control flow.
4. Run the scaffold early, then refine the semantics.
5. Treat the generated code as a draft and review it like any other model output.

## Runtime Note

Automatic diagram understanding uses a hosted vision-language model when Hugging Face inference is available. Current fallback reason: `{reason}`.
"""


def generate_code(image, diagram_type):
    """Convert diagram to code using vision model."""
    if image is None:
        return "❌ Please upload a diagram first!"

    if diagram_type not in PROMPTS:
        return "❌ Please select a diagram type!"

    try:
        # Convert PIL Image to base64
        buffered = BytesIO()
        if isinstance(image, str):
            image = Image.open(image)
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        # Get appropriate prompt
        prompt = PROMPTS[diagram_type]

        # Use Qwen2-VL model for diagram understanding
        response = client.chat_completion(
            model="Qwen/Qwen2-VL-7B-Instruct",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}"}},
                        {"type": "text", "text": prompt}
                    ]
                }
            ],
            max_tokens=2000,
            temperature=0.2  # Low temperature for accurate code
        )

        code = response.choices[0].message.content

        # Format the response
        diagram_name = diagram_type.split("→")[0].strip()
        code_lang = diagram_type.split("→")[1].strip()

        formatted = f"""# 🔄 {diagram_name} → {code_lang}

## Generated Code

{code}

---

### 💡 Next Steps:
1. Review the generated code
2. Adjust data types/logic as needed
3. Add error handling
4. Test thoroughly
5. Refactor for your use case

*Note: AI-generated code should be reviewed before production use.*
"""
        return formatted

    except Exception as e:
        return fallback_scaffold(diagram_type, str(e))


# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    create_premium_hero(
        "Diagram to Code",
        "Translate database schemas, flowcharts, UML, and UI sketches into starter code with a multimodal reasoning workflow.",
        "🔄",
        badge="Multimodal Coding",
        highlights=["Vision-language model", "SQL/Python/HTML", "Implementation scaffold"],
    )
    create_method_panel({
        "Technique": "Image understanding → diagram semantics → language-specific code generation.",
        "What it proves": "You can connect visual AI capabilities to practical developer workflows.",
        "HF capability": "A natural fit for Qwen, Idefics, or other Hub multimodal models.",
    })

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="📸 Upload Diagram",
                type="pil",
                height=400
            )

            diagram_type = gr.Radio(
                choices=[
                    "Database Schema → SQL",
                    "Flowchart → Python",
                    "UML Class Diagram → Classes",
                    "UI Mockup → HTML/CSS",
                    "Whiteboard/Napkin Sketch → Code"
                ],
                label="🎯 Diagram Type",
                value="Database Schema → SQL"
            )

            generate_btn = gr.Button("⚡ Generate Code", variant="primary", size="lg")

            gr.Markdown("""
            ### 💡 Tips:
            - Use clear, standard notation
            - Label boxes and arrows
            - Keep diagrams simple
            - Write legibly (if hand-drawn)
            - Include relationship indicators
            """)

        with gr.Column(scale=1):
            output = gr.Markdown(label="📝 Generated Code")

    # Event handler
    generate_btn.click(
        fn=generate_code,
        inputs=[image_input, diagram_type],
        outputs=[output],
        api_name="generate"
    )

    gr.Markdown("""
    ---
    ### 🎓 Supported Diagram Types:

    #### 1. Database Schema → SQL
    **What to Draw:**
    - Tables as rectangles with table names
    - Columns listed inside each table
    - Primary keys (underline or PK notation)
    - Foreign keys with arrows to related tables
    - Data types (optional but helpful)

    **What You Get:**
    - CREATE TABLE statements
    - PRIMARY KEY constraints
    - FOREIGN KEY relationships
    - Indexed columns
    - Comments explaining structure

    #### 2. Flowchart → Python
    **What to Draw:**
    - Start/End as ovals
    - Process steps as rectangles
    - Decisions as diamonds (yes/no)
    - Arrows showing flow direction
    - Input/Output as parallelograms

    **What You Get:**
    - Python function(s)
    - If/else conditions
    - Loops (while, for)
    - Proper control flow
    - Descriptive variable names
    - Comments and docstrings

    #### 3. UML Class Diagram → Classes
    **What to Draw:**
    - Classes as boxes (3 sections)
    - Class name at top
    - Attributes in middle
    - Methods at bottom
    - Inheritance with arrows
    - Relationships labeled

    **What You Get:**
    - Python class definitions
    - __init__ methods
    - Attribute declarations
    - Method signatures
    - Inheritance structure
    - Type hints
    - Docstrings

    ---

    ### 🎨 Drawing Tips:

    #### For Best Results:
    - **Clarity**: Draw on white/light background
    - **Size**: Make text readable (not too small)
    - **Labels**: Name everything clearly
    - **Arrows**: Show direction with arrow heads
    - **Standard Notation**: Follow common conventions
    - **Lighting**: Good lighting if photographing paper

    #### Hand-Drawn vs. Digital:
    Both work! Hand-drawn diagrams should be:
    - On plain white paper
    - Drawn with dark pen/marker
    - Photographed without shadows
    - Well-lit and in focus

    ---

    ### 🔍 Example Use Cases:

    **Database Design:**
    1. Sketch your database schema on a whiteboard
    2. Take a photo
    3. Get SQL code to create tables
    4. Run in your database

    **Algorithm Planning:**
    1. Draw flowchart for your logic
    2. Upload to get Python implementation
    3. Use as starter code
    4. Refine and test

    **OOP Design:**
    1. Plan class structure with UML
    2. Generate Python classes
    3. Add business logic
    4. Build your application

    ---

    *Note: AI-generated code is a starting point. Always review, test, and refine before using in production.*
    """)

if __name__ == "__main__":
    demo.launch()