from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import requests
import yaml

app = FastAPI(title="AI CI/CD Assistant")

# ЗАМЕНИТЕ на IP вашего контейнера с Ollama
OLLAMA_API_URL = "http://10.10.10.136:11434/api/generate" 

class GenerationRequest(BaseModel):
    description: str
    platform: str = "GitHub Actions" # или GitLab CI
    model: str = "codellama"

def validate_config(config_text: str):
    """Простая валидация YAML синтаксиса"""
    try:
        yaml.safe_load(config_text)
        return True, "Success"
    except Exception as e:
        return False, str(e)

@app.post("/generate")
async def generate_pipeline(req: GenerationRequest):
    prompt = f"Write only a valid {req.platform} YAML configuration for: {req.description}. No explanations, just code."
    
    payload = {
        "model": req.model,
        "prompt": prompt,
        "stream": False
    }

    try:
        response = requests.post(OLLAMA_API_URL, json=payload, timeout=60)
        raw_text = response.json().get("response", "")
        
        # Очистка от Markdown-оберток (```yaml ... ```)
        clean_code = raw_text.replace("```yaml", "").replace("```", "").strip()
        
        is_valid, error_msg = validate_config(clean_code)
        
        return {
            "status": "ok" if is_valid else "error",
            "config": clean_code,
            "validation_error": error_msg if not is_valid else None
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Ollama error: {str(e)}")

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