from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import httpx
import json
import subprocess
import time
import os

app = FastAPI(title="AI Coding Server", version="1.0")

# Enable CORS for all origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

OLLAMA_API = "http://localhost:11434"
CODING_MODEL = "mistral"


def start_ollama():
    """Start Ollama service in background"""
    try:
        httpx.get(f"{OLLAMA_API}/api/tags", timeout=2)
        print("✓ Ollama already running")
        return
    except:
        print("Starting Ollama...")
        subprocess.Popen(["ollama", "serve"], 
                        stdout=subprocess.DEVNULL, 
                        stderr=subprocess.DEVNULL)
        time.sleep(5)
        
        try:
            print(f"Pulling {CODING_MODEL} model...")
            subprocess.run(["ollama", "pull", CODING_MODEL], 
                          timeout=300,
                          stdout=subprocess.DEVNULL,
                          stderr=subprocess.DEVNULL)
            print(f"✓ {CODING_MODEL} ready")
        except Exception as e:
            print(f"⚠ Model pull timeout (will download on first request): {e}")


@app.on_event("startup")
async def startup_event():
    """Initialize on startup"""
    start_ollama()


# ========== API Endpoints ==========

@app.get("/api/health")
async def health_check():
    """Health check endpoint"""
    try:
        async with httpx.AsyncClient() as client:
            await client.get(f"{OLLAMA_API}/api/tags", timeout=2)
        return {
            "status": "healthy",
            "service": "AI Coding Server",
            "model": CODING_MODEL,
            "api": OLLAMA_API
        }
    except Exception as e:
        raise HTTPException(status_code=503, detail=f"Ollama not responding: {str(e)}")


@app.get("/api/models")
async def list_models():
    """List available models"""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{OLLAMA_API}/api/tags", timeout=5)
            return response.json()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/generate")
async def generate_code(request_data: dict):
    """Generate code using the AI model"""
    prompt = request_data.get("prompt", "")
    model = request_data.get("model", CODING_MODEL)
    max_tokens = request_data.get("max_tokens", 2000)
    
    if not prompt:
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
    
    try:
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": False,
            "temperature": 0.3,
            "num_predict": max_tokens,
        }
        
        # 300 second timeout for free tier CPU inference
        async with httpx.AsyncClient(timeout=300) as client:
            response = await client.post(
                f"{OLLAMA_API}/api/generate",
                json=payload
            )
            response.raise_for_status()
            result = response.json()
            
            return {
                "prompt": prompt,
                "model": model,
                "response": result.get("response", ""),
                "done": result.get("done", True),
                "total_duration": result.get("total_duration", 0),
            }
    except httpx.TimeoutException:
        raise HTTPException(status_code=504, detail="⏳ Still generating... Free tier is slow. Try again in 30 seconds!")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")


@app.get("/api/status")
async def status():
    """Get current system status"""
    try:
        async with httpx.AsyncClient() as client:
            models_response = await client.get(f"{OLLAMA_API}/api/tags", timeout=2)
            models = models_response.json().get("models", [])
        
        return {
            "status": "running",
            "models_available": len(models),
            "models": [m.get("name", "unknown") for m in models],
            "default_model": CODING_MODEL,
        }
    except Exception as e:
        return {
            "status": "initializing",
            "error": str(e),
            "message": "Server is starting up, please wait..."
        }


# ========== Web UI ==========

@app.get("/", response_class=HTMLResponse)
async def serve_ui():
    """Serve the web UI"""
    return """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>🤖 AI Coding Server</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }
            
            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif;
                background: linear-gradient(135deg, #1e1e2e 0%, #2a2a3e 100%);
                min-height: 100vh;
                padding: 20px;
                color: #e0e0e0;
            }
            
            .container {
                max-width: 1200px;
                margin: 0 auto;
            }
            
            .header {
                text-align: center;
                margin-bottom: 40px;
                animation: fadeIn 0.6s ease-in;
            }
            
            .header h1 {
                font-size: 3em;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                -webkit-background-clip: text;
                -webkit-text-fill-color: transparent;
                background-clip: text;
                margin-bottom: 10px;
                font-weight: 800;
            }
            
            .header p {
                color: #a0a0a0;
                font-size: 1.1em;
            }
            
            .main-content {
                display: grid;
                grid-template-columns: 2fr 1fr;
                gap: 30px;
                margin-bottom: 40px;
            }
            
            @media (max-width: 768px) {
                .main-content {
                    grid-template-columns: 1fr;
                }
            }
            
            .card {
                background: rgba(255, 255, 255, 0.05);
                border: 1px solid rgba(255, 255, 255, 0.1);
                border-radius: 15px;
                padding: 30px;
                backdrop-filter: blur(10px);
                transition: all 0.3s ease;
            }
            
            .card:hover {
                border-color: rgba(102, 126, 234, 0.5);
                transform: translateY(-5px);
            }
            
            .card h2 {
                margin-bottom: 20px;
                font-size: 1.5em;
                color: #667eea;
            }
            
            .input-group {
                margin-bottom: 20px;
            }
            
            label {
                display: block;
                margin-bottom: 8px;
                font-weight: 600;
                color: #c0c0c0;
                font-size: 0.95em;
            }
            
            textarea {
                width: 100%;
                padding: 15px;
                border: 1px solid rgba(255, 255, 255, 0.1);
                border-radius: 10px;
                background: rgba(255, 255, 255, 0.05);
                color: #e0e0e0;
                font-family: 'Monaco', 'Menlo', monospace;
                font-size: 0.95em;
                resize: vertical;
                min-height: 120px;
                transition: all 0.3s ease;
            }
            
            textarea:focus {
                outline: none;
                border-color: #667eea;
                background: rgba(255, 255, 255, 0.08);
            }
            
            .slider-group {
                display: grid;
                grid-template-columns: 1fr 1fr;
                gap: 15px;
            }
            
            .slider-item input[type="range"] {
                width: 100%;
                height: 6px;
                border-radius: 5px;
                background: rgba(255, 255, 255, 0.1);
                outline: none;
                -webkit-appearance: none;
            }
            
            .slider-item input[type="range"]::-webkit-slider-thumb {
                -webkit-appearance: none;
                appearance: none;
                width: 20px;
                height: 20px;
                border-radius: 50%;
                background: #667eea;
                cursor: pointer;
                transition: all 0.2s ease;
            }
            
            .slider-item input[type="range"]::-webkit-slider-thumb:hover {
                background: #764ba2;
                transform: scale(1.2);
            }
            
            .slider-item input[type="range"]::-moz-range-thumb {
                width: 20px;
                height: 20px;
                border-radius: 50%;
                background: #667eea;
                cursor: pointer;
                border: none;
            }
            
            .slider-value {
                display: inline-block;
                background: #667eea;
                padding: 2px 8px;
                border-radius: 5px;
                font-size: 0.85em;
                margin-left: 10px;
                color: white;
            }
            
            .button-group {
                display: flex;
                gap: 10px;
                margin-top: 30px;
            }
            
            button {
                flex: 1;
                padding: 15px;
                border: none;
                border-radius: 10px;
                font-size: 1em;
                font-weight: 600;
                cursor: pointer;
                transition: all 0.3s ease;
            }
            
            .btn-primary {
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
            }
            
            .btn-primary:hover:not(:disabled) {
                transform: translateY(-2px);
                box-shadow: 0 10px 25px rgba(102, 126, 234, 0.4);
            }
            
            .btn-primary:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }
            
            .btn-secondary {
                background: rgba(255, 255, 255, 0.1);
                color: #667eea;
                border: 1px solid #667eea;
            }
            
            .btn-secondary:hover:not(:disabled) {
                background: rgba(102, 126, 234, 0.2);
            }
            
            .sidebar-card {
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
            }
            
            .sidebar-card h3 {
                margin-bottom: 15px;
                font-size: 1.2em;
            }
            
            .sidebar-card ul {
                list-style: none;
                line-height: 1.8;
                font-size: 0.95em;
            }
            
            .sidebar-card li::before {
                content: "✓ ";
                margin-right: 8px;
                color: #fff;
                font-weight: bold;
            }
            
            .output-section {
                margin-top: 30px;
            }
            
            .output-section h3 {
                margin-bottom: 15px;
                color: #667eea;
                font-size: 1.2em;
            }
            
            .code-output {
                background: rgba(0, 0, 0, 0.3);
                border: 1px solid rgba(255, 255, 255, 0.1);
                border-radius: 10px;
                padding: 20px;
                overflow-x: auto;
                max-height: 400px;
                font-family: 'Monaco', 'Menlo', monospace;
                font-size: 0.9em;
                line-height: 1.6;
                color: #a8ff60;
                white-space: pre-wrap;
                word-wrap: break-word;
            }
            
            .code-output.empty {
                color: #a0a0a0;
                font-style: italic;
                min-height: 100px;
                display: flex;
                align-items: center;
                justify-content: center;
            }
            
            .status {
                display: inline-block;
                padding: 8px 15px;
                border-radius: 20px;
                font-size: 0.9em;
                font-weight: 600;
                margin-bottom: 15px;
            }
            
            .status.loading {
                background: #667eea;
                color: white;
                animation: pulse 1.5s infinite;
            }
            
            .status.success {
                background: #4ade80;
                color: white;
            }
            
            .status.error {
                background: #f87171;
                color: white;
            }
            
            @keyframes fadeIn {
                from {
                    opacity: 0;
                    transform: translateY(-20px);
                }
                to {
                    opacity: 1;
                    transform: translateY(0);
                }
            }
            
            @keyframes pulse {
                0%, 100% {
                    opacity: 1;
                }
                50% {
                    opacity: 0.7;
                }
            }
            
            .copy-btn {
                padding: 8px 12px !important;
                font-size: 0.85em !important;
                margin-top: 15px !important;
            }
            
            .footer {
                text-align: center;
                color: #808080;
                margin-top: 40px;
                padding-top: 20px;
                border-top: 1px solid rgba(255, 255, 255, 0.1);
            }
        </style>
    </head>
    <body>
        <div class="container">
            <div class="header">
                <h1>🚀 AI Coding Assistant</h1>
                <p>Powered by Mistral AI • Generate code instantly • 100% Local & Private</p>
            </div>
            
            <div class="main-content">
                <div class="card">
                    <h2>📝 Code Generator</h2>
                    
                    <div class="input-group">
                        <label for="prompt">Describe what you want to code:</label>
                        <textarea 
                            id="prompt" 
                            placeholder="e.g., Write a Python function to sort a list using bubble sort..."
                        ></textarea>
                    </div>
                    
                    <div class="slider-group">
                        <div class="slider-item">
                            <label for="temperature">🎚️ Temperature</label>
                            <input type="range" id="temperature" min="0" max="1" step="0.1" value="0.3">
                            <span class="slider-value" id="tempValue">0.3</span>
                        </div>
                        <div class="slider-item">
                            <label for="maxTokens">📏 Max Tokens</label>
                            <input type="range" id="maxTokens" min="100" max="2000" step="100" value="1000">
                            <span class="slider-value" id="tokenValue">1000</span>
                        </div>
                    </div>
                    
                    <div class="button-group">
                        <button class="btn-primary" id="generateBtn" onclick="generateCode()">✨ Generate Code</button>
                        <button class="btn-secondary" onclick="clearAll()">🔄 Clear</button>
                    </div>
                    
                    <div class="output-section">
                        <h3>Generated Code:</h3>
                        <div id="status" class="status" style="display: none;"></div>
                        <div class="code-output empty" id="output">💡 Your generated code will appear here...</div>
                        <button class="btn-secondary copy-btn" id="copyBtn" style="display: none;" onclick="copyCode()">📋 Copy to Clipboard</button>
                    </div>
                </div>
                
                <div class="card sidebar-card">
                    <h3>⚡ Quick Tips</h3>
                    <ul>
                        <li>Be specific in prompts</li>
                        <li>Include language preference</li>
                        <li>Mention patterns/libraries</li>
                        <li>Free tier = slower first request</li>
                        <li>100% local & private</li>
                        <li>Works offline</li>
                    </ul>
                    
                    <hr style="margin: 20px 0; border: none; border-top: 1px solid rgba(255,255,255,0.2);">
                    
                    <h3 style="margin-top: 20px;">📚 Example Prompts</h3>
                    <ul style="font-size: 0.9em;">
                        <li>Write Fibonacci in Python</li>
                        <li>FastAPI hello world</li>
                        <li>React todo component</li>
                        <li>SQL duplicate finder</li>
                    </ul>
                </div>
            </div>
            
            <div class="footer">
                <p>🔒 100% Private • 🚀 Instant Generation • 💰 Completely Free</p>
                <p style="font-size: 0.9em; margin-top: 10px;">Mistral 7B Model • FastAPI Backend • HuggingFace Spaces</p>
            </div>
        </div>
        
        <script>
            // Update slider values
            document.getElementById('temperature').addEventListener('input', (e) => {
                document.getElementById('tempValue').textContent = e.target.value;
            });
            
            document.getElementById('maxTokens').addEventListener('input', (e) => {
                document.getElementById('tokenValue').textContent = e.target.value;
            });
            
            async function generateCode() {
                const prompt = document.getElementById('prompt').value;
                
                if (!prompt.trim()) {
                    showStatus('❌ Please enter a prompt', 'error');
                    return;
                }
                
                const generateBtn = document.getElementById('generateBtn');
                const output = document.getElementById('output');
                const statusDiv = document.getElementById('status');
                
                generateBtn.disabled = true;
                output.textContent = '';
                output.classList.add('empty');
                statusDiv.style.display = 'inline-block';
                statusDiv.textContent = '⏳ Generating code... (this may take 30-60 seconds on first run)';
                statusDiv.className = 'status loading';
                document.getElementById('copyBtn').style.display = 'none';
                
                try {
                    const response = await fetch('/api/generate', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json',
                        },
                        body: JSON.stringify({
                            prompt: prompt,
                            model: 'mistral',
                            max_tokens: parseInt(document.getElementById('maxTokens').value)
                        })
                    });
                    
                    if (!response.ok) {
                        const error = await response.json();
                        throw new Error(error.detail || 'Generation failed');
                    }
                    
                    const result = await response.json();
                    output.textContent = result.response;
                    output.classList.remove('empty');
                    statusDiv.textContent = '✅ Generation complete!';
                    statusDiv.className = 'status success';
                    document.getElementById('copyBtn').style.display = 'block';
                    
                } catch (error) {
                    output.textContent = `Error: ${error.message}`;
                    output.classList.remove('empty');
                    statusDiv.textContent = '❌ ' + error.message;
                    statusDiv.className = 'status error';
                    console.error('Error:', error);
                } finally {
                    generateBtn.disabled = false;
                }
            }
            
            function copyCode() {
                const code = document.getElementById('output').textContent;
                navigator.clipboard.writeText(code).then(() => {
                    const copyBtn = document.getElementById('copyBtn');
                    const original = copyBtn.textContent;
                    copyBtn.textContent = '✅ Copied!';
                    setTimeout(() => {
                        copyBtn.textContent = original;
                    }, 2000);
                }).catch(err => {
                    console.error('Copy failed:', err);
                });
            }
            
            function clearAll() {
                document.getElementById('prompt').value = '';
                document.getElementById('output').textContent = '💡 Your generated code will appear here...';
                document.getElementById('output').classList.add('empty');
                document.getElementById('status').style.display = 'none';
                document.getElementById('copyBtn').style.display = 'none';
                document.getElementById('prompt').focus();
            }
            
            function showStatus(message, type) {
                const statusDiv = document.getElementById('status');
                statusDiv.textContent = message;
                statusDiv.className = `status ${type}`;
                statusDiv.style.display = 'inline-block';
            }
            
            // Allow Enter key to generate code
            document.getElementById('prompt').addEventListener('keydown', (e) => {
                if (e.ctrlKey && e.key === 'Enter') {
                    generateCode();
                }
            });
        </script>
    </body>
    </html>
    """


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=int(os.getenv("PORT", 7860)),
        log_level="info"
    )