src/streamlit_app.py

import streamlit as st
import re
from io import BytesIO
import base64
import graphviz

# Page configuration
st.set_page_config(
    page_title="DBML Visualizer",
    page_icon="🗄️",
    layout="wide"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: bold;
        margin-bottom: 0.5rem;
        color: #1f77b4;
    }
    .sub-header {
        font-size: 1.2rem;
        color: #666;
        margin-bottom: 2rem;
    }
    .stTextArea textarea {
        font-family: 'Courier New', monospace;
    }
</style>
""", unsafe_allow_html=True)


class DBMLParser:
    """Parse DBML code into structured data"""
    
    def __init__(self, dbml_code):
        self.dbml_code = dbml_code
        self.tables = {}
        self.relationships = []
        
    def parse(self):
        """Parse DBML code"""
        self._parse_tables()
        self._parse_relationships()
        return self.tables, self.relationships
    
    def _parse_tables(self):
        """Extract tables and columns from DBML"""
        table_pattern = r'Table\s+(\w+)\s*{([^}]*)}'
        
        for match in re.finditer(table_pattern, self.dbml_code, re.DOTALL):
            table_name = match.group(1)
            table_body = match.group(2)
            
            columns = []
            for line in table_body.strip().split('\n'):
                line = line.strip()
                if not line or line.startswith('//'):
                    continue
                
                # Parse column definition
                col_match = re.match(r'(\w+)\s+(\w+)(\s+\[[^\]]*\])?', line)
                if col_match:
                    col_name = col_match.group(1)
                    col_type = col_match.group(2)
                    constraints = col_match.group(3) or ''
                    
                    is_pk = 'pk' in constraints.lower() or 'primary key' in constraints.lower()
                    is_unique = 'unique' in constraints.lower()
                    is_null = 'null' in constraints.lower() and 'not null' not in constraints.lower()
                    is_not_null = 'not null' in constraints.lower()
                    
                    columns.append({
                        'name': col_name,
                        'type': col_type,
                        'is_pk': is_pk,
                        'is_unique': is_unique,
                        'is_null': is_null,
                        'is_not_null': is_not_null
                    })
            
            self.tables[table_name] = columns
    
    def _parse_relationships(self):
        """Extract relationships from DBML"""
        # Pattern: Ref: table1.column1 > table2.column2
        ref_pattern = r'Ref:\s*(\w+)\.(\w+)\s*([<>-]+)\s*(\w+)\.(\w+)'
        
        for match in re.finditer(ref_pattern, self.dbml_code):
            from_table = match.group(1)
            from_column = match.group(2)
            rel_type = match.group(3)
            to_table = match.group(4)
            to_column = match.group(5)
            
            # Determine relationship type
            if '>' in rel_type:
                rel_name = 'one-to-many'
            elif '<' in rel_type:
                rel_name = 'many-to-one'
                from_table, to_table = to_table, from_table
                from_column, to_column = to_column, from_column
            else:
                rel_name = 'one-to-one'
            
            self.relationships.append({
                'from_table': from_table,
                'from_column': from_column,
                'to_table': to_table,
                'to_column': to_column,
                'type': rel_name
            })


class GraphvizGenerator:
    """Generate Graphviz diagram from parsed DBML"""
    
    # Color scheme
    COLORS = {
        'table_bg': '#E8F4F8',
        'table_border': '#2E86AB',
        'pk_bg': '#FFE5B4',
        'pk_text': '#8B4513',
        'header_bg': '#2E86AB',
        'header_text': '#FFFFFF',
        'text': '#333333',
        'arrow': '#555555'
    }
    
    def __init__(self, tables, relationships):
        self.tables = tables
        self.relationships = relationships
    
    def generate(self):
        """Generate Graphviz diagram"""
        dot = graphviz.Digraph(comment='Database Schema')
        dot.attr(rankdir='LR', bgcolor='white', splines='ortho', nodesep='1', ranksep='1.5')
        dot.attr('node', shape='plaintext')
        dot.attr('edge', color=self.COLORS['arrow'], penwidth='2')
        
        # Add tables
        for table_name, columns in self.tables.items():
            html_label = self._create_table_html(table_name, columns)
            dot.node(table_name, label=f'<{html_label}>')
        
        # Add relationships with proper arrows
        for rel in self.relationships:
            if rel['type'] == 'one-to-many':
                # One to many: arrow from parent to child
                dot.edge(
                    rel['to_table'], 
                    rel['from_table'],
                    label=f" {rel['to_column']} → {rel['from_column']} ",
                    arrowhead='crow',
                    fontsize='10',
                    fontcolor='#666666'
                )
            elif rel['type'] == 'many-to-one':
                # Many to one: arrow from child to parent
                dot.edge(
                    rel['from_table'],
                    rel['to_table'],
                    label=f" {rel['from_column']} → {rel['to_column']} ",
                    arrowhead='normal',
                    fontsize='10',
                    fontcolor='#666666'
                )
            else:
                # One to one
                dot.edge(
                    rel['from_table'],
                    rel['to_table'],
                    label=f" {rel['from_column']} ↔ {rel['to_column']} ",
                    arrowhead='none',
                    dir='both',
                    fontsize='10',
                    fontcolor='#666666'
                )
        
        return dot
    
    def _create_table_html(self, table_name, columns):
        """Create HTML table for Graphviz node"""
        html = f'''<
        <TABLE BORDER="2" CELLBORDER="1" CELLSPACING="0" CELLPADDING="8" BGCOLOR="{self.COLORS['table_bg']}" COLOR="{self.COLORS['table_border']}">
            <TR>
                <TD COLSPAN="2" BGCOLOR="{self.COLORS['header_bg']}" ALIGN="CENTER">
                    <FONT COLOR="{self.COLORS['header_text']}" POINT-SIZE="14"><B>{table_name}</B></FONT>
                </TD>
            </TR>
        '''
        
        for col in columns:
            # Determine background color
            bg_color = self.COLORS['pk_bg'] if col['is_pk'] else self.COLORS['table_bg']
            
            # Build column name with constraints
            col_name = col['name']
            if col['is_pk']:
                col_name = f"🔑 {col_name}"
            
            # Build type with constraints
            type_str = col['type']
            constraints = []
            if col['is_pk']:
                constraints.append('PK')
            if col['is_unique']:
                constraints.append('UQ')
            if col['is_not_null']:
                constraints.append('NOT NULL')
            
            if constraints:
                type_str += f" [{', '.join(constraints)}]"
            
            html += f'''
            <TR>
                <TD ALIGN="LEFT" BGCOLOR="{bg_color}">
                    <FONT COLOR="{self.COLORS['text']}" POINT-SIZE="11"><B>{col_name}</B></FONT>
                </TD>
                <TD ALIGN="LEFT" BGCOLOR="{bg_color}">
                    <FONT COLOR="{self.COLORS['text']}" POINT-SIZE="10">{type_str}</FONT>
                </TD>
            </TR>
            '''
        
        html += '''
        </TABLE>
        >'''
        
        return html


class PostgreSQLGenerator:
    """Generate PostgreSQL DDL from parsed DBML"""
    
    def __init__(self, tables, relationships):
        self.tables = tables
        self.relationships = relationships
    
    def generate(self):
        """Generate PostgreSQL CREATE TABLE statements"""
        sql_statements = []
        
        # Create tables
        for table_name, columns in self.tables.items():
            lines = [f"CREATE TABLE {table_name} ("]
            
            col_defs = []
            pk_columns = []
            
            for col in columns:
                col_def = f"    {col['name']} {self._map_type(col['type'])}"
                
                if col['is_pk']:
                    pk_columns.append(col['name'])
                
                if col['is_not_null'] and not col['is_pk']:
                    col_def += " NOT NULL"
                elif col['is_null']:
                    col_def += " NULL"
                
                if col['is_unique'] and not col['is_pk']:
                    col_def += " UNIQUE"
                
                col_defs.append(col_def)
            
            # Add primary key constraint
            if pk_columns:
                col_defs.append(f"    PRIMARY KEY ({', '.join(pk_columns)})")
            
            lines.append(',\n'.join(col_defs))
            lines.append(");")
            
            sql_statements.append('\n'.join(lines))
        
        # Add foreign key constraints
        fk_statements = []
        for rel in self.relationships:
            fk_name = f"fk_{rel['from_table']}_{rel['to_table']}"
            fk_stmt = f"""
ALTER TABLE {rel['from_table']}
ADD CONSTRAINT {fk_name}
FOREIGN KEY ({rel['from_column']})
REFERENCES {rel['to_table']}({rel['to_column']});"""
            fk_statements.append(fk_stmt)
        
        all_sql = '\n\n'.join(sql_statements)
        if fk_statements:
            all_sql += '\n\n-- Foreign Key Constraints\n' + '\n'.join(fk_statements)
        
        return all_sql
    
    def _map_type(self, dbml_type):
        """Map DBML types to PostgreSQL types"""
        type_mapping = {
            'int': 'INTEGER',
            'integer': 'INTEGER',
            'bigint': 'BIGINT',
            'smallint': 'SMALLINT',
            'varchar': 'VARCHAR(255)',
            'text': 'TEXT',
            'bool': 'BOOLEAN',
            'boolean': 'BOOLEAN',
            'date': 'DATE',
            'datetime': 'TIMESTAMP',
            'timestamp': 'TIMESTAMP',
            'decimal': 'DECIMAL',
            'float': 'REAL',
            'double': 'DOUBLE PRECISION',
            'uuid': 'UUID',
            'json': 'JSONB'
        }
        
        return type_mapping.get(dbml_type.lower(), dbml_type.upper())


def main():
    st.markdown('<div class="main-header">🗄️ DBML Visualizer & SQL Generator</div>', unsafe_allow_html=True)
    st.markdown('<div class="sub-header">Visualize database schemas and generate PostgreSQL code</div>', unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.header("📖 About")
        st.write("""
        This tool allows you to:
        - 📊 Visualize DBML schemas
        - 💾 Download diagram as PNG/SVG
        - 🔧 Generate PostgreSQL DDL
        - 📋 Copy SQL to clipboard
        """)
        
        st.header("💡 Example DBML")
        if st.button("Load Example"):
            st.session_state.example_loaded = True
        
        st.header("⚙️ Diagram Settings")
        image_format = st.selectbox("Download Format", ["PNG", "SVG"], index=0)
        st.session_state.image_format = image_format.lower()
    
    # Example DBML
    example_dbml = """Table users {
  id int [pk, increment]
  username varchar [unique, not null]
  email varchar [unique, not null]
  created_at timestamp
}

Table posts {
  id int [pk, increment]
  title varchar [not null]
  content text
  user_id int [not null]
  created_at timestamp
}

Table comments {
  id int [pk, increment]
  post_id int [not null]
  user_id int [not null]
  content text [not null]
  created_at timestamp
}

Ref: posts.user_id > users.id
Ref: comments.post_id > posts.id
Ref: comments.user_id > users.id"""
    
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("📝 DBML Input")
        
        default_value = example_dbml if st.session_state.get('example_loaded', False) else ""
        dbml_input = st.text_area(
            "Enter your DBML code:",
            value=default_value,
            height=400,
            placeholder="Table users {\n  id int [pk]\n  name varchar\n}\n\nTable posts {\n  id int [pk]\n  user_id int\n}\n\nRef: posts.user_id > users.id"
        )
        
        if st.button("🎨 Generate Visualization & SQL", type="primary", use_container_width=True):
            if dbml_input.strip():
                try:
                    with st.spinner("Parsing DBML and generating visualization..."):
                        # Parse DBML
                        parser = DBMLParser(dbml_input)
                        tables, relationships = parser.parse()
                        
                        if not tables:
                            st.error("❌ No tables found in DBML code. Please check your syntax.")
                        else:
                            # Generate Graphviz diagram
                            graph_gen = GraphvizGenerator(tables, relationships)
                            dot = graph_gen.generate()
                            
                            # Generate PostgreSQL
                            sql_gen = PostgreSQLGenerator(tables, relationships)
                            sql_code = sql_gen.generate()
                            
                            # Store in session state
                            st.session_state.dot = dot
                            st.session_state.sql_code = sql_code
                            st.session_state.tables = tables
                            st.session_state.relationships = relationships
                            
                            st.success(f"✅ Successfully parsed {len(tables)} tables and {len(relationships)} relationships!")
                
                except Exception as e:
                    st.error(f"❌ Error parsing DBML: {str(e)}")
                    import traceback
                    st.code(traceback.format_exc())
            else:
                st.warning("⚠️ Please enter some DBML code first.")
    
    with col2:
        if 'dot' in st.session_state:
            st.subheader("📊 Database Diagram")
            
            # Display the diagram
            st.graphviz_chart(st.session_state.dot)
            
            # Download buttons
            col_btn1, col_btn2 = st.columns(2)
            
            with col_btn1:
                # PNG download
                try:
                    png_data = st.session_state.dot.pipe(format='png')
                    st.download_button(
                        label="📥 Download PNG",
                        data=png_data,
                        file_name="database_schema.png",
                        mime="image/png",
                        use_container_width=True
                    )
                except:
                    st.warning("PNG export requires Graphviz installation")
            
            with col_btn2:
                # SVG download
                try:
                    svg_data = st.session_state.dot.pipe(format='svg')
                    st.download_button(
                        label="📥 Download SVG",
                        data=svg_data,
                        file_name="database_schema.svg",
                        mime="image/svg+xml",
                        use_container_width=True
                    )
                except:
                    st.warning("SVG export requires Graphviz installation")
            
            # Info about tables
            st.info(f"📋 **{len(st.session_state.tables)}** tables, **{len(st.session_state.relationships)}** relationships")
            
        else:
            st.info("👈 Enter DBML code and click 'Generate' to see the visualization")
    
    # SQL Output Section
    if 'sql_code' in st.session_state:
        st.markdown("---")
        st.subheader("🔧 Generated PostgreSQL Code")
        
        col_sql1, col_sql2 = st.columns([4, 1])
        
        with col_sql1:
            st.code(st.session_state.sql_code, language="sql", line_numbers=True)
        
        with col_sql2:
            st.download_button(
                label="📥 Download SQL",
                data=st.session_state.sql_code,
                file_name="schema.sql",
                mime="text/sql",
                use_container_width=True
            )
            
            if st.button("📋 Copy SQL", use_container_width=True):
                st.toast("Copy the SQL code from the code block above!", icon="📋")


if __name__ == "__main__":
    main()