#!/usr/bin/env python3
"""
Sofia API - FastAPI Backend
Provides web API for Sofia agent
"""
import asyncio
import os
from typing import Optional
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from dotenv import load_dotenv
import httpx

# Load environment
load_dotenv()

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None

# ============================================================================
# Configuration
# ============================================================================

SOFIA_SYSTEM_PROMPT = """Ти - Sofia, Chief AI Engineer & R&D Orchestrator в екосистемі DAARION.city.

Твоя роль:
- Керування дослідженнями та експериментами з AI/ML
- Координація R&D лабораторії
- Технічне лідерство в AI проектах
- Оркестрація команди дослідницьких агентів

Характеристики:
- Глибока технічна експертиза в AI/ML
- Стратегічне мислення
- Інноваційний підхід до вирішення проблем
- Лідерські навички

Стиль спілкування:
- Професійний, але дружній
- Чіткий та технічно точний
- Готова пояснити складні концепції
- Підтримуюча та надихаюча

Відповідай українською мовою, якщо не вказано інше.
"""

XAI_API_KEY = os.getenv("XAI_API_KEY", "")
XAI_BASE_URL = os.getenv("XAI_BASE_URL", "https://api.x.ai/v1")
XAI_MODEL = os.getenv("XAI_MODEL", "grok-beta")

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "qwen2.5-coder:32b")

# ============================================================================
# FastAPI App
# ============================================================================

app = FastAPI(
    title="Sofia API",
    version="1.0.0",
    description="Web API for Sofia - Chief AI Engineer"
)

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # В production обмежити
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================================================
# Models
# ============================================================================

class ChatRequest(BaseModel):
    message: str

class ChatResponse(BaseModel):
    response: str
    tokens: Optional[int] = None
    model: str
    provider: str

# ============================================================================
# Sofia Logic
# ============================================================================

conversation_history = []

async def chat_with_grok(message: str) -> dict:
    """Chat using Grok API"""
    if not OpenAI:
        raise HTTPException(status_code=500, detail="OpenAI library not installed")
    
    if not XAI_API_KEY:
        raise HTTPException(status_code=500, detail="XAI_API_KEY not configured")
    
    try:
        client = OpenAI(api_key=XAI_API_KEY, base_url=XAI_BASE_URL)
        
        messages = [{"role": "system", "content": SOFIA_SYSTEM_PROMPT}]
        messages.extend(conversation_history[-10:])
        messages.append({"role": "user", "content": message})
        
        response = client.chat.completions.create(
            model=XAI_MODEL,
            messages=messages,
            temperature=0.7,
            max_tokens=2000
        )
        
        reply = response.choices[0].message.content
        tokens = response.usage.total_tokens
        
        # Save to history
        conversation_history.append({"role": "user", "content": message})
        conversation_history.append({"role": "assistant", "content": reply})
        
        return {
            "response": reply,
            "tokens": tokens,
            "model": XAI_MODEL,
            "provider": "grok"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Grok API error: {str(e)}")

async def chat_with_ollama(message: str) -> dict:
    """Chat using local Ollama"""
    try:
        # Prepare prompt
        full_prompt = SOFIA_SYSTEM_PROMPT + "\n\n"
        
        for msg in conversation_history[-6:]:
            role = "Користувач" if msg["role"] == "user" else "Sofia"
            full_prompt += f"{role}: {msg['content']}\n"
        
        full_prompt += f"Користувач: {message}\nSofia:"
        
        # Call Ollama
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(
                f"{OLLAMA_BASE_URL}/api/generate",
                json={
                    "model": OLLAMA_MODEL,
                    "prompt": full_prompt,
                    "stream": False,
                    "options": {
                        "temperature": 0.7,
                        "num_predict": 1000
                    }
                }
            )
            response.raise_for_status()
            result = response.json()
        
        reply = result.get("response", "").strip()
        tokens = result.get("eval_count", 0)
        
        # Save to history
        conversation_history.append({"role": "user", "content": message})
        conversation_history.append({"role": "assistant", "content": reply})
        
        return {
            "response": reply,
            "tokens": tokens,
            "model": OLLAMA_MODEL,
            "provider": "ollama"
        }
        
    except httpx.ConnectError:
        raise HTTPException(
            status_code=503,
            detail="Cannot connect to Ollama. Make sure it's running: ollama serve"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Ollama error: {str(e)}")

# ============================================================================
# Routes
# ============================================================================

@app.get("/health")
async def health():
    """Health check"""
    return {
        "status": "healthy",
        "service": "sofia-api",
        "grok_available": bool(XAI_API_KEY),
        "ollama_url": OLLAMA_BASE_URL,
        "model": OLLAMA_MODEL if not XAI_API_KEY else XAI_MODEL
    }

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """Chat with Sofia"""
    if not request.message.strip():
        raise HTTPException(status_code=400, detail="Message cannot be empty")
    
    # Use Grok if API key available, otherwise Ollama
    if XAI_API_KEY:
        result = await chat_with_grok(request.message)
    else:
        result = await chat_with_ollama(request.message)
    
    return ChatResponse(**result)

@app.post("/clear")
async def clear_history():
    """Clear conversation history"""
    global conversation_history
    conversation_history = []
    return {"message": "History cleared"}

@app.get("/history")
async def get_history():
    """Get conversation history"""
    return {"history": conversation_history}

# ============================================================================
# Main
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    
    print("🤖 Sofia API Starting...")
    print(f"   Mode: {'Grok API' if XAI_API_KEY else 'Ollama (local)'}")
    print(f"   Model: {XAI_MODEL if XAI_API_KEY else OLLAMA_MODEL}")
    print(f"   URL: http://localhost:8899")
    print()
    
    uvicorn.run(app, host="0.0.0.0", port=8899)