import streamlit as st
from transformers import pipeline
import json
from pathlib import Path
import os

class FastWebGenerator:
    def __init__(self):
        # Load the smallest model possible for quick text completion
        self.generator = pipeline(
            'text-generation',
            model='distilgpt2',
            device_map='auto'
        )
        
        # Load templates
        self.templates = {
            "react": {
                "component": """
import React, { useState, useEffect } from 'react';

const {component_name} = () => {
    const [data, setData] = useState([]);

    useEffect(() => {
        // Fetch data
        fetchData();
    }, []);

    const fetchData = async () => {
        try {
            const response = await fetch('/api/data');
            const jsonData = await response.json();
            setData(jsonData);
        } catch (error) {
            console.error('Error:', error);
        }
    };

    return (
        <div className="container mx-auto p-4">
            <h1 className="text-2xl font-bold mb-4">{title}</h1>
            {content}
        </div>
    );
};

export default {component_name};
""",
                "api": """
async function fetchData() {
    const response = await fetch('/api/endpoint');
    const data = await response.json();
    return data;
}

async function postData(data) {
    const response = await fetch('/api/endpoint', {
        method: 'POST',
        headers: {
            'Content-Type': 'application/json',
        },
        body: JSON.stringify(data),
    });
    return response.json();
}
"""
            },
            "fastapi": {
                "main": """
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

app = FastAPI()

class ItemModel(BaseModel):
    id: Optional[int]
    name: str
    description: Optional[str]

@app.get("/")
async def root():
    return {"message": "API is running"}

@app.get("/items")
async def get_items():
    return {"items": items}

@app.post("/items")
async def create_item(item: ItemModel):
    items.append(item)
    return item

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
""",
                "crud": """
@app.get("/items/{item_id}")
async def get_item(item_id: int):
    if item_id < 0 or item_id >= len(items):
        raise HTTPException(status_code=404, detail="Item not found")
    return items[item_id]

@app.put("/items/{item_id}")
async def update_item(item_id: int, item: ItemModel):
    if item_id < 0 or item_id >= len(items):
        raise HTTPException(status_code=404, detail="Item not found")
    items[item_id] = item
    return item

@app.delete("/items/{item_id}")
async def delete_item(item_id: int):
    if item_id < 0 or item_id >= len(items):
        raise HTTPException(status_code=404, detail="Item not found")
    item = items.pop(item_id)
    return item
"""
            },
            "mongodb": """
from motor.motor_asyncio import AsyncIOMotorClient

client = AsyncIOMotorClient('mongodb://localhost:27017')
db = client.database_name

async def get_all():
    cursor = db.collection.find({})
    return await cursor.to_list(length=100)

async def get_one(id: str):
    return await db.collection.find_one({"_id": id})

async def create_one(data: dict):
    result = await db.collection.insert_one(data)
    return result.inserted_id

async def update_one(id: str, data: dict):
    result = await db.collection.update_one(
        {"_id": id},
        {"$set": data}
    )
    return result.modified_count

async def delete_one(id: str):
    result = await db.collection.delete_one({"_id": id})
    return result.deleted_count
"""
        }

    def customize_code(self, template, params):
        # Quick text customization using the model
        prompt = f"Customize this code: {template}\nWith these parameters: {params}\n"
        result = self.generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
        return result

    def generate_project(self, specs):
        frontend_code = self.templates["react"]["component"].replace(
            "{component_name}", specs["component_name"]
        ).replace(
            "{title}", specs["title"]
        ).replace(
            "{content}", specs["content"]
        )
        
        backend_code = self.templates["fastapi"]["main"]
        if specs.get("crud", False):
            backend_code += self.templates["fastapi"]["crud"]
            
        db_code = self.templates["mongodb"] if specs["database"] == "MongoDB" else ""
        
        return {
            "frontend": frontend_code,
            "backend": backend_code,
            "database": db_code
        }

def main():
    st.set_page_config(page_title="Fast Code Generator", layout="wide")
    
    if 'generator' not in st.session_state:
        with st.spinner("Loading generator..."):
            st.session_state.generator = FastWebGenerator()
    
    st.title("⚡ Fast Code Generator")
    
    with st.form("generate_form"):
        col1, col2 = st.columns(2)
        
        with col1:
            component_name = st.text_input("Component Name", "MyComponent")
            title = st.text_input("Title", "My App")
            content = st.text_area("Content", "<p>Your content here</p>")
        
        with col2:
            database = st.selectbox("Database", ["MongoDB", "PostgreSQL"])
            add_crud = st.checkbox("Add CRUD Operations", value=True)
            add_auth = st.checkbox("Add Authentication", value=False)
        
        if st.form_submit_button("Generate"):
            specs = {
                "component_name": component_name,
                "title": title,
                "content": content,
                "database": database,
                "crud": add_crud,
                "auth": add_auth
            }
            
            result = st.session_state.generator.generate_project(specs)
            
            # Display results
            tabs = st.tabs(["Frontend", "Backend", "Database"])
            
            with tabs[0]:
                st.code(result["frontend"], language="javascript")
            with tabs[1]:
                st.code(result["backend"], language="python")
            with tabs[2]:
                st.code(result["database"], language="python")
            
            # Add download button
            combined_code = f"""
            // Frontend Code
            {result['frontend']}
            
            // Backend Code
            {result['backend']}
            
            // Database Code
            {result['database']}
            """
            
            st.download_button(
                "Download Code",
                combined_code,
                file_name="generated_code.txt",
                mime="text/plain"
            )

if __name__ == "__main__":
    main()