#!/usr/bin/env python3
"""
CyberBrain26 SAFE-OI API v0.4 - AsyncInterpreter s approval tokem
"""

import os
import sys
import uuid
import json
import logging
import asyncio
import threading
import time
from typing import Dict, Optional, Any, List
from pathlib import Path
from datetime import datetime
from collections import deque

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn

# --- Logging ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("/tmp/cb26-safeoi-api.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("safe_oi_api")

# --- Konfigurace ---
LITELLM_PROXY = os.getenv("LITELLM_PROXY", "https://openrouter.ai/api/v1")
LITELLM_MASTER_KEY = os.getenv("OPENROUTER_API_KEY") or os.getenv("LITELLM_MASTER_KEY")
DEFAULT_MODEL = "nvidia/nemotron-3-ultra-550b-a55b:free"
HOST = "127.0.0.1"
PORT = 7682
OI_MAX_TOKENS = 1024

MODEL_ALIASES = {
    "nemo-free": "nvidia/nemotron-3-ultra-550b-a55b:free",
    "gpt-oss-free": "openai/gpt-oss-20b:free",
    "glm-free": "z-ai/glm-5.2:free",
    "north-code-free": "cohere/north-mini-code:free",
    "claude-haiku": "anthropic/claude-haiku-4.5",
    "gemini-flash": "google/gemini-2.5-flash",
    "llama-3.1": "meta-llama/llama-3.1-70b-instruct",
    "mistral-large": "mistralai/mistral-large-2512",
}

app = FastAPI(title="CyberBrain26 SAFE-OI API", version="0.4")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:8080", "http://127.0.0.1:8080"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

sessions: Dict[str, dict] = {}
pending_approvals: Dict[str, dict] = {}

class ChatRequest(BaseModel):
    thread_id: Optional[str] = None
    message: str
    model: str = DEFAULT_MODEL
    mode: str = "SAFE"
    depth: str = "normal"
    confirm_full: bool = False

class ApproveRequest(BaseModel):
    action: str = "approve"

def resolve_model(model: str) -> str:
    return MODEL_ALIASES.get(model, model)

def get_or_create_session(thread_id: Optional[str], model: str, mode: str) -> str:
    if not thread_id or thread_id not in sessions:
        thread_id = thread_id or str(uuid.uuid4())
        sessions[thread_id] = {
            "history": [],
            "mode": mode.upper(),
            "model": model,
            "interpreter": None,
            "last_auto_run": None,
            "pending_approval": None,
            "loop": None,
        }
        logger.info(f"Nová session: {thread_id} | mode={mode}")
    return thread_id

def create_async_interpreter(model: str, auto_run: bool):
    """Vytvoří AsyncInterpreter instanci"""
    try:
        from interpreter.core.async_core import AsyncInterpreter
        
        interp = AsyncInterpreter()
        interp.llm.model = f"openai/{model}"
        interp.llm.api_base = LITELLM_PROXY
        interp.llm.api_key = LITELLM_MASTER_KEY
        interp.llm.headers = {"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}
        interp.llm.max_tokens = OI_MAX_TOKENS
        interp.auto_run = auto_run
        interp.offline = False
        interp.safe_mode = "off"
        interp.print = False
        interp.debug = False
        
        logger.info(f"AsyncInterpreter vytvořen | model={model} | auto_run={auto_run}")
        return interp
    except Exception as e:
        logger.error(f"Nepodařilo se vytvořit AsyncInterpreter: {e}")
        raise HTTPException(status_code=500, detail=f"Chyba při inicializaci: {e}")

def classify_action(code: str, language: str) -> tuple:
    """Konzervativní klasifikace - vše kromě čistého čtení vyžaduje potvrzení"""
    lang = (language or "").strip().lower()
    
    # Bezpečné read-only příkazy
    read_only = {"python", "py"}
    read_only_commands = {"pwd", "ls", "dir", "find", "cat", "head", "tail", "grep", "echo", "stat"}
    
    if lang in read_only:
        # Python kód - zkontrolujeme zda neobsahuje zápis
        if any(x in code for x in ["open(", "write(", "remove(", "rmdir", "unlink", "rmtree"]):
            return "write", ["zápis nebo mazání souborů"]
        if any(x in code for x in ["subprocess", "os.system", "os.popen"]):
            return "shell", ["spouštění shell příkazů"]
        return "read", ["pouze čtení"]
    
    elif lang in ("shell", "bash", "sh"):
        for cmd in read_only_commands:
            if code.strip().startswith(cmd + " "):
                return "read", ["pouze čtení"]
        return "write", ["shell příkaz"]
    
    return "unknown", ["neznámá akce"]

@app.get("/health")
def health():
    return {
        "status": "ok",
        "service": "safe_oi_api",
        "version": "0.4",
        "sessions": len(sessions),
        "proxy": LITELLM_PROXY
    }

@app.post("/chat")
def chat(req: ChatRequest):
    mode = req.mode.upper()
    
    if mode == "FULL" and not req.confirm_full:
        raise HTTPException(
            status_code=403,
            detail="FULL-OI je možné aktivovat pouze s confirm_full=true"
        )
    
    resolved_model = resolve_model(req.model)
    thread_id = get_or_create_session(req.thread_id, resolved_model, mode)
    session = sessions[thread_id]
    session["mode"] = mode
    session["model"] = resolved_model
    
    # Vytvořit nebo získat AsyncInterpreter
    auto_run = (mode == "FULL")
    if session["interpreter"] is None or session.get("last_auto_run") != auto_run:
        session["interpreter"] = create_async_interpreter(resolved_model, auto_run=auto_run)
        session["last_auto_run"] = auto_run
    
    interp = session["interpreter"]
    
    # Připravit zprávu
    if req.depth == "fast":
        prefix = "[Odpověz stručně] "
    elif req.depth == "deep":
        prefix = "[Přemýšlej důkladně, postupuj krok za krokem] "
    else:
        prefix = ""
    
    user_message = prefix + req.message
    logger.info(f"[{thread_id}] {mode} | {resolved_model} | {user_message[:80]}...")
    
    try:
        # Odeslat zprávu do AsyncInterpreter
        async def run_chat():
            # Reset messages pro novou interakci
            interp.messages = interp.messages[-10:] if len(interp.messages) > 10 else interp.messages
            
            # Přidat user zprávu
            interp.messages.append({
                "role": "user",
                "type": "message",
                "content": user_message
            })
            
            # Spustit respond thread
            interp.stop_event.clear()
            interp.respond_thread = threading.Thread(
                target=interp.respond,
                args=(None,)  # run_code=None použije auto_run
            )
            interp.respond_thread.start()
            
            # Sbírat output chunky
            chunks = []
            code_chunks = []
            timeout = time.time() + 60  # 60s timeout
            
            while interp.respond_thread.is_alive() or not interp.output_queue.sync_q.empty():
                if time.time() > timeout:
                    interp.stop_event.set()
                    raise TimeoutError("Request timeout")
                
                try:
                    chunk = interp.output_queue.sync_q.get(timeout=0.5)
                    chunks.append(chunk)
                    
                    if chunk.get("type") == "code" and "content" in chunk:
                        code_chunks.append(chunk)
                    
                    if chunk.get("type") == "status" and chunk.get("content") == "complete":
                        break
                except:
                    if not interp.respond_thread.is_alive():
                        break
            
            return chunks, code_chunks
        
        chunks, code_chunks = asyncio.run(run_chat())
        
        # Zpracovat odpověď
        reply_parts = []
        for chunk in chunks:
            if chunk.get("type") == "message" and "content" in chunk:
                reply_parts.append(chunk["content"])
        
        reply = "".join(reply_parts) if reply_parts else "Žádná odpověď."
        
        # Pokud jsou code chunky a jsme v SAFE, vytvořit approval
        if code_chunks and mode == "SAFE":
            code_content = "".join(c.get("content", "") for c in code_chunks)
            language = code_chunks[0].get("format", "unknown") if code_chunks else "unknown"
            category, reasons = classify_action(code_content, language)
            
            approval_id = str(uuid.uuid4())
            pending_approvals[approval_id] = {
                "thread_id": thread_id,
                "code": code_content,
                "language": language,
                "category": category,
                "reasons": reasons,
                "status": "pending",
                "created_at": datetime.now().isoformat(),
            }
            
            session["pending_approval"] = approval_id
            
            reply += f"\n\n---\n**Návrh kódu čeká na potvrzení**\n\nTyp: {category}\nDůvod: {', '.join(reasons)}\n\nPoužij /approvals/{approval_id} pro schválení nebo zamítnutí."
        
        # Uložit do historie
        session["history"].append({"role": "user", "content": req.message})
        session["history"].append({"role": "assistant", "content": reply})
        
        return {
            "thread_id": thread_id,
            "reply": reply,
            "mode": mode,
            "model": req.model,
            "auto_run": auto_run,
            "has_approval": session["pending_approval"] is not None
        }
        
    except Exception as e:
        logger.exception(f"Chyba v /chat: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/approvals")
def list_approvals():
    return {
        "count": len(pending_approvals),
        "approvals": [
            {
                "approval_id": aid,
                "category": a.get("category"),
                "status": a.get("status"),
                "created_at": a.get("created_at")
            }
            for aid, a in pending_approvals.items()
        ]
    }

@app.get("/approvals/{approval_id}")
def get_approval(approval_id: str):
    if approval_id not in pending_approvals:
        raise HTTPException(status_code=404, detail="Approval not found")
    return pending_approvals[approval_id]

@app.post("/approvals/{approval_id}")
def approve_action(approval_id: str, req: ApproveRequest):
    if approval_id not in pending_approvals:
        raise HTTPException(status_code=404, detail="Approval not found")
    
    approval = pending_approvals[approval_id]
    thread_id = approval["thread_id"]
    
    if thread_id not in sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    session = sessions[thread_id]
    interp = session.get("interpreter")
    
    if interp is None:
        raise HTTPException(status_code=500, detail="Interpreter not found")
    
    if req.action == "approve":
        approval["status"] = "approved"
        session["pending_approval"] = None
        
        # Odeslat "go" příkaz
        async def send_go():
            await interp.input({
                "role": "user",
                "type": "command",
                "content": "go"
            })
        
        try:
            asyncio.run(send_go())
            logger.info(f"Schváleno: {approval_id}")
            return {"status": "approved", "approval_id": approval_id}
        except Exception as e:
            logger.error(f"Chyba při schvalování: {e}")
            raise HTTPException(status_code=500, detail=str(e))
    
    elif req.action == "deny":
        approval["status"] = "denied"
        session["pending_approval"] = None
        logger.info(f"Zamítnuto: {approval_id}")
        return {"status": "denied", "approval_id": approval_id}
    
    else:
        raise HTTPException(status_code=400, detail="Neznámá akce")

@app.get("/sessions")
def list_sessions():
    return {
        "count": len(sessions),
        "sessions": [
            {
                "thread_id": tid,
                "mode": s.get("mode"),
                "model": s.get("model"),
                "messages": len(s.get("history", [])),
                "has_pending": s.get("pending_approval") is not None
            }
            for tid, s in sessions.items()
        ]
    }

@app.delete("/sessions/{thread_id}")
def delete_session(thread_id: str):
    if thread_id in sessions:
        try:
            interp = sessions[thread_id].get("interpreter")
            if interp:
                interp.stop_event.set()
        except Exception:
            pass
        del sessions[thread_id]
        logger.info(f"Session smazána: {thread_id}")
        return {"status": "deleted"}
    raise HTTPException(status_code=404, detail="Session not found")

if __name__ == "__main__":
    logger.info(f"Starting SAFE-OI API v0.4 on http://{HOST}:{PORT}")
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")