"""
code-memory: OpenAI-compatible proxy with GraphRAG + persistent memory.
Endpoints:
  POST /v1/chat/completions  — drop-in Ollama replacement with memory
  POST /v1/index             — index a code repository
  GET  /v1/memories          — list memories
  DELETE /v1/memories/{id}   — delete a memory
  GET  /health               — health check
"""
import json, os, asyncio
from contextlib import asynccontextmanager
from typing import AsyncIterator

import asyncpg
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from dotenv import load_dotenv

load_dotenv()

import db as _db
import memory as mem
import graph as gph
import ollama_client as oll
from apscheduler.schedulers.asyncio import AsyncIOScheduler

# ── lifespan ──────────────────────────────────────────────────────────────────

scheduler = AsyncIOScheduler()

@asynccontextmanager
async def lifespan(app: FastAPI):
    await _db.init_db()
    # CMA consolidation runs every 6 hours
    scheduler.add_job(mem.consolidate_memories, "interval", hours=6,
                      kwargs={"user_id": "default"})
    scheduler.start()
    yield
    scheduler.shutdown()

app = FastAPI(title="code-memory", lifespan=lifespan)

# ── models ────────────────────────────────────────────────────────────────────

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: str = Field(default=None)
    messages: list[Message]
    stream: bool = False
    session_id: str = "default"
    user_id: str = "default"
    project_id: str | None = None

class IndexRequest(BaseModel):
    project_id: str
    root_dir: str
    extensions: list[str] = [".py"]

# ── models list endpoint ──────────────────────────────────────────────────────

@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{
            "id": "code-memory",
            "object": "model",
            "owned_by": "code-memory",
            "created": 1700000000,
        }]
    }

# ── context assembly ──────────────────────────────────────────────────────────

async def build_context(request: ChatRequest) -> list[dict]:
    user_msg = next((m.content for m in reversed(request.messages)
                     if m.role == "user"), "")

    # 1. Relevant archival memories
    memories = await mem.search_memories(user_msg, user_id=request.user_id, limit=6)
    mem_block = ""
    if memories:
        mem_block = "## Relevant memories\n" + "\n".join(
            f"- [{m['memory_type']}] {m['content']}" for m in memories
            if m.get("score", 0) > 0.4
        )

    # 2. GraphRAG code context
    code_block = ""
    if request.project_id:
        ctx = await gph.graphrag_context(user_msg, request.project_id)
        if ctx:
            code_block = f"## Code architecture context\n{ctx}"

    # 3. Conversation summary (if exists)
    summary = await mem.get_summary(request.session_id)
    summary_block = f"## Earlier conversation summary\n{summary}" if summary else ""

    # 4. Assemble system prompt
    base_system = """Ты — AI-ассистент по программированию с постоянной памятью.

У тебя есть три уровня памяти:
1. Долгосрочная память — факты о пользователе и проектах (показаны ниже в "Relevant memories")
2. История диалога — последние 20 сообщений текущей сессии
3. Контекст кода — граф архитектуры проекта (если проиндексирован)

Используй память чтобы давать точные ответы с учётом контекста пользователя. \
Если в памяти есть релевантный факт — опирайся на него, не переспрашивай то что уже известно.

Пользователь может управлять твоей памятью прямо в чате:
- "запомни что X" — ты сохранишь X с высоким приоритетом
- "забудь про X" — ты удалишь похожие воспоминания
- "что ты помнишь?" — покажи список активных воспоминаний через GET /v1/memories"""
    context_parts = [p for p in [base_system, summary_block, mem_block, code_block] if p]
    system_content = "\n\n".join(context_parts)

    # 5. Build messages: system + recent history + current
    recent = await mem.get_recent_turns(request.session_id, limit=20)
    messages = [{"role": "system", "content": system_content}]
    messages += [{"role": r["role"], "content": r["content"]} for r in recent]
    # add current user messages (skip if already in history)
    for m in request.messages:
        if not any(h["role"] == m.role and h["content"] == m.content for h in recent[-2:]):
            messages.append({"role": m.role, "content": m.content})

    return messages

# ── chat endpoint ─────────────────────────────────────────────────────────────

FORGET_PREFIXES = ("забудь про ", "забудь ", "forget ")
REMEMBER_PREFIXES = ("запомни что ", "запомни ", "remember ")

async def handle_memory_command(user_content: str, user_id: str) -> str | None:
    low = user_content.lower().strip()
    for prefix in FORGET_PREFIXES:
        if low.startswith(prefix):
            topic = user_content[len(prefix):].strip()
            pool = await _db.get_pool()
            vec = await oll.embed(topic)
            from memory import _vec_str
            async with pool.acquire() as conn:
                deleted = await conn.fetchval("""
                    DELETE FROM memories
                    WHERE user_id=$1
                      AND 1-(embedding<=>$2::vector) > 0.60
                    RETURNING content
                """, user_id, _vec_str(vec))
            if deleted:
                return f"Забыл: «{deleted}»"
            return f"Не нашёл ничего похожего на «{topic}»"
    for prefix in REMEMBER_PREFIXES:
        if low.startswith(prefix):
            fact = user_content[len(prefix):].strip()
            await mem.store_memory(fact, user_id=user_id, memory_type="manual", importance=0.9)
            return f"Запомнил: «{fact}»"
    return None

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest, bg: BackgroundTasks):
    user_content = next((m.content for m in reversed(req.messages)
                         if m.role == "user"), "")

    # intercept memory commands before building context
    cmd_response = await handle_memory_command(user_content, req.user_id)
    if cmd_response:
        return {
            "id": "cmpl-cmd",
            "object": "chat.completion",
            "model": req.model or os.getenv("CHAT_MODEL"),
            "choices": [{"index": 0,
                         "message": {"role": "assistant", "content": cmd_response},
                         "finish_reason": "stop"}]
        }

    messages = await build_context(req)
    # use real Ollama model, not our proxy name
    ollama_model = None if req.model == "code-memory" else req.model

    if req.stream:
        collected: list[str] = []

        async def gen() -> AsyncIterator[bytes]:
            try:
                async for chunk in oll.chat_stream(messages, model=ollama_model):
                    collected.append(chunk)
                    data = json.dumps({"choices": [{"delta": {"content": chunk},
                                                   "finish_reason": None}]})
                    yield f"data: {data}\n\n".encode()
                yield b"data: [DONE]\n\n"
            except Exception as e:
                yield f"data: {json.dumps({'error': str(e)})}\n\n".encode()

        async def _save_stream():
            response_text = "".join(collected)
            await mem.save_turn(req.session_id, "user",      user_content)
            await mem.save_turn(req.session_id, "assistant", response_text)
            await mem.extract_and_store(user_content,  user_id=req.user_id)
            await mem.extract_and_store(response_text, user_id=req.user_id)
            await mem.summarize_old_turns(req.session_id)

        bg.add_task(_save_stream)
        return StreamingResponse(gen(), media_type="text/event-stream")

    response_text = await oll.chat(messages, model=ollama_model)

    # background: save turns + extract facts
    async def _save():
        await mem.save_turn(req.session_id, "user",      user_content)
        await mem.save_turn(req.session_id, "assistant", response_text)
        await mem.extract_and_store(user_content,  user_id=req.user_id)
        await mem.extract_and_store(response_text, user_id=req.user_id)
        await mem.summarize_old_turns(req.session_id)

    bg.add_task(_save)

    return {
        "id": "cmpl-1",
        "object": "chat.completion",
        "model": req.model or os.getenv("CHAT_MODEL"),
        "choices": [{"index": 0,
                     "message": {"role": "assistant", "content": response_text},
                     "finish_reason": "stop"}]
    }

# ── index endpoint ────────────────────────────────────────────────────────────

@app.post("/v1/index")
async def index_code(req: IndexRequest, bg: BackgroundTasks):
    if not os.path.isdir(req.root_dir):
        raise HTTPException(404, f"Directory not found: {req.root_dir}")
    bg.add_task(gph.index_project, req.project_id, req.root_dir, req.extensions)
    return {"status": "indexing started", "project_id": req.project_id}

# ── memory endpoints ──────────────────────────────────────────────────────────

@app.get("/v1/memories")
async def list_memories(user_id: str = "default", limit: int = 50):
    pool = await _db.get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch("""
            SELECT id, content, importance, memory_type, access_count,
                   last_accessed, created_at
            FROM memories WHERE user_id=$1
              AND (metadata->>'archived') IS NULL
            ORDER BY importance DESC, last_accessed DESC LIMIT $2
        """, user_id, limit)
        return [dict(r) for r in rows]

@app.post("/v1/memories")
async def create_memory(body: dict, user_id: str = "default"):
    content = body.get("content", "").strip()
    if not content:
        raise HTTPException(400, "content required")
    importance = float(body.get("importance", 0.9))
    memory_type = body.get("memory_type", "manual")
    await mem.store_memory(content, user_id=user_id,
                           memory_type=memory_type, importance=importance)
    return {"status": "stored", "content": content}

@app.patch("/v1/memories/{memory_id}")
async def update_memory(memory_id: str, body: dict):
    import uuid as _uuid
    pool = await _db.get_pool()
    uid = _uuid.UUID(memory_id)
    async with pool.acquire() as conn:
        row = await conn.fetchrow("SELECT * FROM memories WHERE id=$1", uid)
        if not row:
            raise HTTPException(404, "memory not found")
        new_content = body.get("content", row["content"])
        new_importance = float(body.get("importance", row["importance"]))
        if new_content != row["content"]:
            vec = await oll.embed(new_content)
            from memory import _vec_str
            await conn.execute("""
                UPDATE memories SET content=$1, embedding=$2::vector,
                                    importance=$3, last_accessed=NOW()
                WHERE id=$4
            """, new_content, _vec_str(vec), new_importance, uid)
        else:
            await conn.execute("""
                UPDATE memories SET importance=$1, last_accessed=NOW()
                WHERE id=$2
            """, new_importance, uid)
    return {"updated": memory_id}

@app.delete("/v1/memories/{memory_id}")
async def delete_memory(memory_id: str):
    pool = await _db.get_pool()
    async with pool.acquire() as conn:
        await conn.execute("DELETE FROM memories WHERE id=$1",
                           __import__("uuid").UUID(memory_id))
    return {"deleted": memory_id}

@app.get("/health")
async def health():
    pool = await _db.get_pool()
    async with pool.acquire() as conn:
        await conn.fetchval("SELECT 1")
    return {"status": "ok", "db": "connected"}