#!/usr/bin/env python3
"""
Backend Code Generation API Service
===================================

Production-ready API service for serving the trained backend code generation model.
Provides RESTful endpoints for generating complete backend applications.
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Any
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import zipfile
import tempfile
import os
import uuid
from datetime import datetime
import asyncio
import logging
from pathlib import Path

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Pydantic models for API
class CodeGenerationRequest(BaseModel):
    description: str = Field(..., description="Description of the backend application to generate")
    framework: str = Field(..., description="Target framework (express, fastapi, django, flask)")
    language: str = Field(..., description="Programming language (javascript, python)")
    requirements: List[str] = Field(default=[], description="List of specific requirements")
    project_name: Optional[str] = Field(default=None, description="Custom project name")
    
    class Config:
        schema_extra = {
            "example": {
                "description": "E-commerce API with user authentication and product management",
                "framework": "fastapi",
                "language": "python",
                "requirements": [
                    "User registration and login",
                    "JWT authentication",
                    "Product CRUD operations",
                    "Shopping cart functionality",
                    "Order management"
                ],
                "project_name": "ecommerce-api"
            }
        }

class GenerationResponse(BaseModel):
    task_id: str
    status: str
    message: str
    estimated_time: int

class GenerationStatus(BaseModel):
    task_id: str
    status: str  # pending, processing, completed, failed
    progress: int  # 0-100
    message: str
    generated_files: Optional[Dict[str, str]] = None
    download_url: Optional[str] = None
    error: Optional[str] = None

class GeneratedProject(BaseModel):
    project_name: str
    framework: str
    language: str
    files: Dict[str, str]
    structure: Dict[str, Any]
    setup_instructions: List[str]
    features: List[str]

# Global model instance
class ModelManager:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.loaded = False
    
    async def load_model(self, model_path: str = "./trained_model"):
        """Load the trained model asynchronously"""
        try:
            logger.info(f"Loading model from {model_path} on {self.device}")
            
            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                device_map="auto" if self.device == "cuda" else None
            )
            
            if self.device == "cpu":
                self.model = self.model.to(self.device)
            
            self.loaded = True
            logger.info("Model loaded successfully!")
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise
    
    def generate_code(self, prompt: str, max_tokens: int = 1024) -> str:
        """Generate code using the trained model"""
        if not self.loaded:
            raise RuntimeError("Model not loaded")
        
        inputs = self.tokenizer.encode(prompt, return_tensors='pt')
        inputs = inputs.to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs,
                max_length=min(max_tokens, 1024),
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
                pad_token_id=self.tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        
        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return generated_text[len(self.tokenizer.decode(inputs[0], skip_special_tokens=True)):]

# Global instances
model_manager = ModelManager()
generation_tasks = {}  # Store generation tasks

# FastAPI app
app = FastAPI(
    title="Backend Code Generation API",
    description="AI-powered backend application generator",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
async def startup_event():
    """Load model on startup"""
    model_path = os.getenv("MODEL_PATH", "./trained_model")
    await model_manager.load_model(model_path)

@app.get("/")
async def root():
    """API root endpoint"""
    return {
        "service": "Backend Code Generation API",
        "version": "1.0.0",
        "status": "running",
        "model_loaded": model_manager.loaded,
        "endpoints": {
            "generate": "/api/v1/generate",
            "status": "/api/v1/status/{task_id}",
            "download": "/api/v1/download/{task_id}",
            "health": "/health"
        }
    }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "OK",
        "timestamp": datetime.utcnow().isoformat(),
        "model_loaded": model_manager.loaded,
        "device": model_manager.device if model_manager.loaded else None
    }

@app.post("/api/v1/generate", response_model=GenerationResponse)
async def generate_backend(
    request: CodeGenerationRequest,
    background_tasks: BackgroundTasks
):
    """Generate a complete backend application"""
    
    if not model_manager.loaded:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    # Create unique task ID
    task_id = str(uuid.uuid4())
    
    # Initialize task status
    generation_tasks[task_id] = GenerationStatus(
        task_id=task_id,
        status="pending",
        progress=0,
        message="Task queued for processing"
    )
    
    # Start background generation
    background_tasks.add_task(
        generate_project_background,
        task_id,
        request
    )
    
    return GenerationResponse(
        task_id=task_id,
        status="accepted",
        message="Code generation started",
        estimated_time=60  # seconds
    )

@app.get("/api/v1/status/{task_id}", response_model=GenerationStatus)
async def get_generation_status(task_id: str):
    """Get the status of a generation task"""
    
    if task_id not in generation_tasks:
        raise HTTPException(status_code=404, detail="Task not found")
    
    return generation_tasks[task_id]

@app.get("/api/v1/download/{task_id}")
async def download_generated_project(task_id: str):
    """Download the generated project as a ZIP file"""
    
    if task_id not in generation_tasks:
        raise HTTPException(status_code=404, detail="Task not found")
    
    task = generation_tasks[task_id]
    
    if task.status != "completed":
        raise HTTPException(status_code=400, detail="Generation not completed")
    
    if not task.download_url:
        raise HTTPException(status_code=404, detail="Download file not available")
    
    if not os.path.exists(task.download_url):
        raise HTTPException(status_code=404, detail="Download file not found")
    
    return FileResponse(
        path=task.download_url,
        filename=f"generated_project_{task_id}.zip",
        media_type="application/zip"
    )

@app.delete("/api/v1/cleanup/{task_id}")
async def cleanup_task(task_id: str):
    """Clean up task files and data"""
    
    if task_id not in generation_tasks:
        raise HTTPException(status_code=404, detail="Task not found")
    
    task = generation_tasks[task_id]
    
    # Remove download file if exists
    if task.download_url and os.path.exists(task.download_url):
        os.remove(task.download_url)
    
    # Remove task from memory
    del generation_tasks[task_id]
    
    return {"message": "Task cleaned up successfully"}

async def generate_project_background(task_id: str, request: CodeGenerationRequest):
    """Background task for generating the complete project"""
    
    task = generation_tasks[task_id]
    
    try:
        # Update status
        task.status = "processing"
        task.progress = 10
        task.message = "Analyzing requirements..."
        
        # Create the generation prompt
        prompt = create_generation_prompt(request)
        
        # Update progress
        task.progress = 30
        task.message = "Generating application structure..."
        
        # Generate code using the model
        generated_code = model_manager.generate_code(prompt, max_tokens=1024)
        
        # Update progress
        task.progress = 60
        task.message = "Processing generated code..."
        
        # Parse and structure the generated code
        project_files = parse_generated_code(generated_code, request)
        
        # Update progress
        task.progress = 80
        task.message = "Creating project files..."
        
        # Create downloadable ZIP file
        zip_path = create_project_zip(task_id, project_files, request)
        
        # Complete the task
        task.status = "completed"
        task.progress = 100
        task.message = "Project generated successfully"
        task.generated_files = {name: "Generated" for name in project_files.keys()}
        task.download_url = zip_path
        
    except Exception as e:
        logger.error(f"Generation failed for task {task_id}: {e}")
        task.status = "failed"
        task.error = str(e)
        task.message = "Generation failed"

def create_generation_prompt(request: CodeGenerationRequest) -> str:
    """Create the prompt for the model"""
    
    prompt_parts = [
        f"Description: {request.description}",
        f"Framework: {request.framework}",
        f"Language: {request.language}",
    ]
    
    if request.requirements:
        prompt_parts.append("Requirements:")
        for req in request.requirements:
            prompt_parts.append(f"- {req}")
    
    if request.project_name:
        prompt_parts.append(f"Project Name: {request.project_name}")
    
    prompt_parts.append("Generate the complete backend application with all necessary files:")
    
    return "\n".join(prompt_parts)

def parse_generated_code(generated_code: str, request: CodeGenerationRequest) -> Dict[str, str]:
    """Parse the generated code into individual files"""
    
    files = {}
    
    # Simple parsing logic - in production, this should be more sophisticated
    lines = generated_code.split('\n')
    current_file = None
    current_content = []
    
    for line in lines:
        if line.startswith('--- ') and line.endswith(' ---'):
            # Save previous file
            if current_file:
                files[current_file] = '\n'.join(current_content)
            
            # Start new file
            current_file = line.replace('--- ', '').replace(' ---', '').strip()
            current_content = []
            
        elif current_file and not line.startswith('--- End ---'):
            current_content.append(line)
    
    # Save last file
    if current_file and current_content:
        files[current_file] = '\n'.join(current_content)
    
    # If parsing failed, create basic structure based on framework
    if not files:
        files = create_fallback_structure(request)
    
    return files

def create_fallback_structure(request: CodeGenerationRequest) -> Dict[str, str]:
    """Create a basic project structure if parsing fails"""
    
    if request.framework.lower() == 'fastapi':
        return {
            'main.py': f'''from fastapi import FastAPI

app = FastAPI(title="{request.description}")

@app.get("/")
async def root():
    return {{"message": "Hello from {request.description}"}}

@app.get("/health")
async def health():
    return {{"status": "OK"}}
''',
            'requirements.txt': '''fastapi==0.104.1
uvicorn[standard]==0.24.0'''
        }
    
    elif request.framework.lower() == 'express':
        return {
            'app.js': f'''const express = require('express');
const app = express();

app.get('/', (req, res) => {{
    res.json({{ message: 'Hello from {request.description}' }});
}});

app.get('/health', (req, res) => {{
    res.json({{ status: 'OK' }});
}});

const PORT = process.env.PORT || 3000;
app.listen(PORT, () => {{
    console.log(`Server running on port ${{PORT}}`);
}});
''',
            'package.json': json.dumps({
                "name": request.project_name or "generated-backend",
                "version": "1.0.0",
                "main": "app.js",
                "dependencies": {
                    "express": "^4.18.2"
                }
            }, indent=2)
        }
    
    else:
        return {
            'README.md': f'# {request.description}\n\nGenerated backend application using {request.framework}'
        }

def create_project_zip(task_id: str, files: Dict[str, str], request: CodeGenerationRequest) -> str:
    """Create a ZIP file containing all project files"""
    
    # Create temporary directory for the ZIP file
    temp_dir = tempfile.gettempdir()
    zip_path = os.path.join(temp_dir, f"project_{task_id}.zip")
    
    project_name = request.project_name or f"generated_{request.framework}_app"
    
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for filename, content in files.items():
            # Add each file to the ZIP
            arcname = f"{project_name}/{filename}"
            zipf.writestr(arcname, content)
        
        # Add a README with setup instructions
        setup_instructions = get_setup_instructions(request.framework)
        zipf.writestr(f"{project_name}/SETUP.md", setup_instructions)
    
    return zip_path

def get_setup_instructions(framework: str) -> str:
    """Get setup instructions for the framework"""
    
    instructions = {
        'fastapi': '''# Setup Instructions

1. Install dependencies:
   ```bash
   pip install -r requirements.txt
   ```

2. Run the application:
   ```bash
   uvicorn main:app --reload
   ```

3. Access the API:
   - API: http://localhost:8000
   - Docs: http://localhost:8000/docs
''',
        'express': '''# Setup Instructions

1. Install dependencies:
   ```bash
   npm install
   ```

2. Run the application:
   ```bash
   node app.js
   ```

3. Access the API:
   - API: http://localhost:3000
''',
        'django': '''# Setup Instructions

1. Install dependencies:
   ```bash
   pip install -r requirements.txt
   ```

2. Run migrations:
   ```bash
   python manage.py migrate
   ```

3. Run the application:
   ```bash
   python manage.py runserver
   ```

4. Access the API:
   - API: http://localhost:8000
   - Admin: http://localhost:8000/admin
''',
        'flask': '''# Setup Instructions

1. Install dependencies:
   ```bash
   pip install -r requirements.txt
   ```

2. Run the application:
   ```bash
   python run.py
   ```

3. Access the API:
   - API: http://localhost:5000
'''
    }
    
    return instructions.get(framework, '# Setup Instructions\n\nRefer to the framework documentation for setup instructions.')

# Additional utility endpoints
@app.get("/api/v1/frameworks")
async def list_supported_frameworks():
    """List supported frameworks and languages"""
    return {
        "frameworks": [
            {
                "name": "fastapi",
                "language": "python",
                "description": "Modern, fast, web framework for building APIs"
            },
            {
                "name": "express",
                "language": "javascript",
                "description": "Fast, unopinionated web framework for Node.js"
            },
            {
                "name": "django",
                "language": "python",
                "description": "High-level Python web framework"
            },
            {
                "name": "flask",
                "language": "python",
                "description": "Lightweight WSGI web application framework"
            }
        ]
    }

@app.get("/api/v1/examples")
async def get_example_requests():
    """Get example generation requests"""
    return {
        "examples": [
            {
                "name": "E-commerce API",
                "request": {
                    "description": "Complete e-commerce backend with user management and product catalog",
                    "framework": "fastapi",
                    "language": "python",
                    "requirements": [
                        "User registration and authentication",
                        "Product CRUD operations",
                        "Shopping cart functionality",
                        "Order management",
                        "Payment processing integration"
                    ]
                }
            },
            {
                "name": "Task Management System",
                "request": {
                    "description": "Task management system with team collaboration",
                    "framework": "express",
                    "language": "javascript",
                    "requirements": [
                        "User authentication with JWT",
                        "Task CRUD operations",
                        "Team and project management",
                        "Real-time notifications",
                        "File attachments"
                    ]
                }
            },
            {
                "name": "Blog Platform",
                "request": {
                    "description": "Blog platform with content management",
                    "framework": "django",
                    "language": "python",
                    "requirements": [
                        "Article management",
                        "User comments and ratings",
                        "Category and tag system",
                        "SEO optimization",
                        "Media file handling"
                    ]
                }
            }
        ]
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "api_service:app",
        host="0.0.0.0",
        port=8000,
        reload=True
    )