import subprocess
import time
import os
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn
import requests

app = FastAPI(
    title="AJ-DeepSeek API",
    description="FREE AI API by AJ STUDIOZ - Tamil Nadu Heritage Assistant",
    version="1.0.0"
)

class PromptRequest(BaseModel):
    prompt: str

# Global variable to track model status
model_ready = False

def setup_model():
    """Setup AJ-DeepSeek model"""
    global model_ready
    try:
        print("🚀 Starting Ollama server...")
        subprocess.Popen(["ollama", "serve"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(15)  # Wait longer for server to start
        
        print("📥 Pulling base model...")
        subprocess.run(["ollama", "pull", "qwen2.5:7b-instruct"], check=True, timeout=600)
        
        print("🤖 Creating AJ model...")
        subprocess.run(["ollama", "create", "aj-deepseek:latest", "-f", "Modelfile"], check=True, timeout=300)
        
        model_ready = True
        print("✅ AJ-DeepSeek model ready!")
        return True
    except Exception as e:
        print(f"❌ Error setting up model: {e}")
        return False

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

@app.get("/")
async def root():
    """API status endpoint"""
    return {
        "message": "AJ-DeepSeek API by AJ STUDIOZ",
        "heritage": "Tamil Nadu, India", 
        "status": "ready" if model_ready else "initializing",
        "endpoint": "/api/generate",
        "example": {
            "method": "POST",
            "url": "/api/generate", 
            "body": {"prompt": "Hello AJ!"}
        }
    }

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {"status": "healthy", "model_ready": model_ready}

@app.post("/api/generate")
async def generate(request: PromptRequest):
    """Generate response from AJ-DeepSeek model"""
    if not model_ready:
        return JSONResponse(
            status_code=503,
            content={
                "error": "Model is still initializing. Please wait a few minutes and try again.",
                "status": "initializing"
            }
        )
    
    try:
        response = requests.post(
            "http://localhost:11434/api/generate",
            json={
                "model": "aj-deepseek:latest",
                "prompt": request.prompt,
                "stream": False
            },
            timeout=60
        )
        
        if response.status_code == 200:
            result = response.json()
            return {
                "response": result.get("response", ""),
                "model": "aj-deepseek:latest",
                "heritage": "Tamil Nadu, India",
                "created_by": "AJ STUDIOZ"
            }
        else:
            return JSONResponse(
                status_code=500,
                content={"error": "Model inference failed", "details": response.text}
            )
            
    except requests.exceptions.Timeout:
        return JSONResponse(
            status_code=504,
            content={"error": "Request timeout. Please try with a shorter prompt."}
        )
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": f"Internal server error: {str(e)}"}
        )

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)