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

app = FastAPI()

OLLAMA_URL = "http://10.10.10.136:11434/api/generate"

class Prompt(BaseModel):
    task_description: str
    model: str = "codellama"

def validate_yaml(content: str):
    try:
        yaml.safe_load(content)
        return True, "Valid YAML"
    except yaml.YAMLError as exc:
        return False, str(exc)

@app.post("/generate_config")
async def generate_config(data: Prompt):
    payload = {
        "model": data.model,
        "prompt": f"Generate only a valid CI/CD YAML configuration for the following task: {data.task_description}. Return ONLY the code block.",
        "stream": False
    }
    
    try:
        response = requests.post(OLLAMA_URL, json=payload)
        response_data = response.json()
        generated_text = response_data.get("response", "")
        
        # Очистка от markdown-тегов, если модель их добавила
        config_content = generated_text.replace("```yaml", "").replace("```", "").strip()
        
        is_valid, message = validate_yaml(config_content)
        
        return {
            "config": config_content,
            "is_valid": is_valid,
            "validation_message": message
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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