# main.py
from __future__ import annotations

import asyncio as _asyncio
import json as _json
import uuid as _uuid
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
from uuid import UUID

import httpx
from fastapi import FastAPI, HTTPException, Body
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse

from .nexus_config import settings
from .chat import generate_chat_response, stream_chat_response
from .ollama_manager import initialize_ollama, initialize_ollama_async, get_ollama_manager
from .playbook_manager import PlaybookManager

def _render_memory_block(facts) -> str:
    """Render memory items as grouped ## Section / - bullet markdown."""
    from collections import defaultdict
    sections: dict = defaultdict(list)
    for item in facts:
        sections[item.section or "General"].append(item.text)
    parts = []
    for section, lines in sections.items():
        parts.append(f"## {section}\n" + "\n".join(f" - {line}" for line in lines))
    return "\n\n".join(parts)


async def _auto_select_model() -> str:
    try:
        s = store.get_settings()
        if s.get("model"):
            return s["model"]
        return await get_ollama_manager().select_best_model()
    except Exception:
        return getattr(settings, "default_model", None) or "mistral"

from .memory.store import store, MemoryItem
from .playbooks.store import playbook_store, PlaybookItem

MEMORY_SERVICE = "http://localhost:8001"

app = FastAPI(title="Synapse Backend", version="1.0")

# Alias for startup scripts
sio_app = app

# --- CORS ---
# allow_credentials=True is incompatible with allow_origins=["*"] — browsers
# reject such responses. Since this is a local-only service with no cookies/auth,
# wildcard origins without credentials is correct.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- GLOBALS ---
playbooks = PlaybookManager()
ollama = None


# -------------------------
# Startup: Initialize Ollama service
# -------------------------
@app.on_event("startup")
async def startup_event():
    global ollama
    try:
        ollama = await initialize_ollama_async()
        print("[Synapse] Ollama service is running.")
    except Exception as e:
        # Start in degraded mode — chat endpoints will return errors until Ollama
        # is brought up, but the backend itself stays alive so the control panel
        # and other endpoints remain reachable.
        print(f"[Synapse] WARNING: Ollama unavailable at startup: {e}")
        from .ollama_manager import OllamaManager
        ollama = OllamaManager()


# -------------------------
# Root
# -------------------------
@app.get("/")
async def root():
    try:
        status = ollama.get_status() if (ollama is not None and hasattr(ollama, "get_status")) else None
    except Exception:
        status = None
    return {"status": "online", "ollama": status}


# -------------------------
# Chat (non-streaming)
# -------------------------
@app.post("/chat")
async def chat_endpoint(payload: Dict[str, Any]):
    try:
        message = payload.get("message", "")
        app_settings = store.get_settings()
        model = payload.get("model") or await _auto_select_model()
        context = payload.get("context", {})
        history = payload.get("history", [])
        temperature = payload.get("temperature", app_settings.get("temperature"))

        if not message:
            raise HTTPException(status_code=400, detail="Missing 'message'")

        rendered_message = playbooks.render_prompt(message)
        system_prompt = playbooks.get_system_prompt() or app_settings.get("system_prompt", "")

        memory_facts = store.all()
        if memory_facts:
            facts_block = _render_memory_block(memory_facts)
            system_prompt = (system_prompt + "\n\n---\nWhat you know about Jon:\n\n" + facts_block) if system_prompt else facts_block

        metadata: Dict[str, Any] = {"model": model, "context": context, "system": system_prompt, "temperature": temperature}

        result = await generate_chat_response(
            user_message=rendered_message, metadata=metadata, history=history
        )
        return {"response": result.get("response", ""), "model": model}

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# -------------------------
# Chat (streaming)
# -------------------------
@app.post("/chat/stream")
async def chat_stream_endpoint(payload: Dict[str, Any]):
    try:
        message = payload.get("message", "")
        app_settings = store.get_settings()
        model = payload.get("model") or await _auto_select_model()
        context = payload.get("context", {})
        conversation_id = payload.get("conversation_id") or str(_uuid.uuid4())
        history = payload.get("history", [])
        temperature = payload.get("temperature", app_settings.get("temperature"))

        if not message:
            raise HTTPException(status_code=400, detail="Missing 'message'")

        rendered_message = playbooks.render_prompt(message)
        system_prompt = playbooks.get_system_prompt() or app_settings.get("system_prompt", "")

        # Append reference playbooks to the system prompt
        context_pbs = playbooks.get_context_playbooks()
        if context_pbs:
            refs = "\n\n".join(
                f"### {pb.title}\nGoal: {pb.goal}\n\n{pb.instructions}"
                for pb in context_pbs
            )
            separator = "\n\n---\nReference playbooks (read these as additional context):\n\n"
            system_prompt = (system_prompt + separator + refs) if system_prompt else refs

        # Inject persistent memory facts about Jon
        memory_facts = store.all()
        if memory_facts:
            facts_block = _render_memory_block(memory_facts)
            system_prompt = (system_prompt + "\n\n---\nWhat you know about Jon:\n\n" + facts_block) if system_prompt else facts_block

        # Search past conversations for relevant context and inject the top matches.
        # This gives the model memory of prior exchanges without requiring tool-calling support.
        past_context = store.search_conversations(message, limit=2)
        if past_context:
            snippets = []
            for conv in past_context:
                exchange = "\n".join(
                    f"  {m['role'].upper()}: {m['content']}"
                    for m in conv["matches"]
                )
                snippets.append(exchange)
            memory_block = "\n\n".join(snippets)
            separator = "\n\n---\nRelevant past exchanges (use as background context only):\n\n"
            system_prompt = (system_prompt + separator + memory_block) if system_prompt else memory_block

        metadata: Dict[str, Any] = {"model": model, "context": context, "system": system_prompt, "temperature": temperature}

        # Persist conversation and user message before streaming
        store.create_conversation(conversation_id)
        store.add_message(conversation_id, "user", rendered_message)

        async def event_stream() -> AsyncGenerator[str, None]:
            response_chunks: list[str] = []
            meta: dict = {}
            try:
                async for chunk in stream_chat_response(
                    user_message=rendered_message,
                    metadata=metadata,
                    history=history,
                ):
                    if chunk.startswith("__meta__"):
                        try:
                            meta = _json.loads(chunk[8:])
                        except Exception:
                            pass
                        yield f"event: meta\ndata: {chunk[8:]}\n\n"
                        continue
                    response_chunks.append(chunk)
                    yield f"data: {chunk}\n\n"
                # Persist the completed assistant response with model and token stats
                if response_chunks:
                    store.add_message(
                        conversation_id, "assistant", "".join(response_chunks),
                        model=meta.get("model") or model,
                        tokens=meta.get("tokens"),
                    )
            except Exception as e:
                yield f"event: error\ndata: {_json.dumps({'detail': str(e)})}\n\n"
                return

            # Ask the memory service curator to evaluate this exchange
            if response_chunks:
                try:
                    async with httpx.AsyncClient(timeout=25.0) as _mc:
                        r = await _mc.post(
                            f"{MEMORY_SERVICE}/memories/extract",
                            json={
                                "user_message": rendered_message,
                                "assistant_response": "".join(response_chunks),
                            },
                        )
                    if r.status_code == 200:
                        data = r.json()
                        if data.get("saved"):
                            mem_result = {"section": data["section"], "text": data["text"]}
                            yield f"event: memory\ndata: {_json.dumps(mem_result)}\n\n"
                except Exception:
                    pass

        return StreamingResponse(event_stream(), media_type="text/event-stream")

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -------------------------
# Playbook Execution
# -------------------------
@app.post("/playbook/run")
async def run_playbook(payload: Dict[str, Any]):
    name = payload.get("name")
    variables = payload.get("variables", {})

    if not name:
        raise HTTPException(status_code=400, detail="Missing playbook name")

    try:
        result = playbooks.render_prompt(name, variables=variables)
        return {"result": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# -------------------------
# Settings
# -------------------------
@app.get("/settings")
async def get_settings_endpoint():
    return store.get_settings()


@app.put("/settings")
async def put_settings_endpoint(payload: Dict[str, Any] = Body(...)):
    store.update_settings(payload)
    return store.get_settings()


# -------------------------
# Memory
# -------------------------
@app.get("/memory")
async def get_memory():
    return {"items": [{"id": m.id, "section": m.section, "text": m.text, "tags": m.tags} for m in store.all()]}


@app.post("/memory")
async def add_memory(payload: Dict[str, Any] = Body(...)):
    text = (payload.get("text") or "").strip()
    if not text:
        raise HTTPException(status_code=400, detail="Missing 'text'")
    from .memory.store import MemoryItem
    import uuid as _mem_uuid
    item = MemoryItem(
        id=str(_mem_uuid.uuid4()),
        section=(payload.get("section") or "General").strip(),
        text=text,
        tags=payload.get("tags", []),
    )
    store.add(item)
    return {"id": item.id, "section": item.section, "text": item.text, "tags": item.tags}


@app.patch("/memory/{item_id}")
async def update_memory(item_id: str, payload: Dict[str, Any] = Body(...)):
    existing = store.get(item_id)
    if not existing:
        raise HTTPException(status_code=404, detail="Not found")
    from .memory.store import MemoryItem
    updated = MemoryItem(
        id=item_id,
        section=(payload.get("section") or existing.section or "General").strip(),
        text=(payload.get("text") or existing.text).strip(),
        tags=payload.get("tags", existing.tags),
    )
    store.update(updated)
    return {"id": updated.id, "section": updated.section, "text": updated.text, "tags": updated.tags}


@app.delete("/memory/{item_id}")
async def delete_memory(item_id: str):
    if not store.get(item_id):
        raise HTTPException(status_code=404, detail="Not found")
    store.delete(item_id)
    return {"status": "deleted"}


# -------------------------
# Models
# -------------------------
@app.get("/models")
async def get_models():
    try:
        mgr = get_ollama_manager()
        models = await mgr.list_models()
        selected = await mgr.select_best_model()
        return {"models": models, "selected": selected}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/models/details")
async def get_model_details():
    """Return full model objects (name, size, modified_at) from Ollama."""
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            r = await client.get(f"{settings.ollama_host.rstrip('/')}/api/tags")
            r.raise_for_status()
            return r.json()
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Ollama unreachable: {e}")


@app.post("/models/pull")
async def pull_model(payload: Dict[str, Any] = Body(...)):
    """Proxy a streaming model pull from Ollama."""
    name = payload.get("name", "").strip()
    if not name:
        raise HTTPException(status_code=400, detail="Missing model name")

    async def _stream():
        try:
            async with httpx.AsyncClient(timeout=600.0) as client:
                async with client.stream(
                    "POST", f"{settings.ollama_host.rstrip('/')}/api/pull",
                    json={"name": name},
                ) as resp:
                    async for line in resp.aiter_lines():
                        if line:
                            yield line + "\n"
        except Exception as e:
            yield f'{{"error": "{e}"}}\n'
        finally:
            # Bust the model cache so new model is visible immediately
            get_ollama_manager().invalidate_model_cache()

    return StreamingResponse(_stream(), media_type="application/x-ndjson")


@app.delete("/models/{name:path}")
async def delete_model(name: str):
    """Proxy a model deletion to Ollama."""
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            r = await client.request(
                "DELETE", f"{settings.ollama_host.rstrip('/')}/api/delete",
                json={"name": name},
            )
            if not r.is_success:
                raise HTTPException(status_code=r.status_code, detail=r.text)
        get_ollama_manager().invalidate_model_cache()
        return {"status": "deleted", "name": name}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Ollama unreachable: {e}")


# -------------------------
# Ollama Status
# -------------------------
@app.get("/ollama/status")
async def ollama_status_endpoint():
    try:
        if ollama is not None and hasattr(ollama, "get_status"):
            status = ollama.get_status()
        else:
            status = None
        return {"status": status}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -------------------------
# Ollama Start
# -------------------------
@app.post("/ollama/start")
async def ollama_start_endpoint():
    try:
        global ollama
        if ollama is None:
            ollama = initialize_ollama()
        
        # Try to start if it has a start method
        if hasattr(ollama, "start"):
            ollama.start()
        
        is_running = ollama.is_running() if hasattr(ollama, "is_running") else True
        return {"status": "started" if is_running else "failed", "running": is_running}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -------------------------
# Ollama Stop
# -------------------------
@app.post("/ollama/stop")
async def ollama_stop_endpoint():
    try:
        global ollama
        if ollama is not None and hasattr(ollama, "stop"):
            ollama.stop()

        is_running = ollama.is_running() if (ollama is not None and hasattr(ollama, "is_running")) else False
        return {"status": "stopped" if not is_running else "still_running", "running": is_running}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -------------------------
# Playbooks List (existing)
# -------------------------
@app.get("/playbooks")
async def get_playbooks():
    try:
        playbook_list = playbook_store.all_playbooks()
        return {
            "playbooks": [
                {
                    "id": p.id,
                    "title": p.title,
                    "goal": p.goal,
                    "instructions": getattr(p, "instructions", ""),
                    "tags": getattr(p, "tags", []),
                }
                for p in playbook_list
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -------------------------
# Playbook retrieve / update / replace (new)
# -------------------------
def _find_playbook_by_id(playbook_id: str) -> Tuple[Optional[Any], Optional[str]]:
    """
    Try common store getters, then fall back to scanning store.all_playbooks().
    Returns (playbook_obj, key) where key is the identifier used by store if applicable.
    """
    try:
        pb = playbook_store.get_playbook(playbook_id)
        if pb:
            return pb, playbook_id
    except Exception:
        pass

    try:
        for p in playbook_store.all_playbooks():
            pid = str(getattr(p, "id", None) or "")
            if pid == str(playbook_id):
                return p, pid
    except Exception:
        pass

    return None, None


def _persist_playbook(playbook_dict: Dict[str, Any]) -> Dict[str, Any]:
    """
    Persist a playbook dict to the store by converting to PlaybookItem.
    """
    try:
        # Preserve existing order on update; use provided order (or tail) on create
        existing = playbook_store.get_playbook(str(playbook_dict["id"]))
        order = existing.order if existing else playbook_dict.get("order", len(playbook_store.all_playbooks()))

        playbook_item = PlaybookItem(
            id=str(playbook_dict["id"]),
            title=playbook_dict.get("title", ""),
            goal=playbook_dict.get("goal", ""),
            instructions=playbook_dict.get("instructions", ""),
            tags=playbook_dict.get("tags", []),
            order=order
        )

        playbook_store.add_playbook(playbook_item)
        
        # Return as dict
        return {
            "id": playbook_item.id,
            "title": playbook_item.title,
            "goal": playbook_item.goal,
            "instructions": playbook_item.instructions,
            "tags": playbook_item.tags,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to persist playbook: {str(e)}")
    
@app.post("/playbooks/reorder")
async def reorder_playbooks(payload: Dict[str, Any] = Body(...)):
    """
    Accepts {"ids": ["id1", "id2", "id3"]} in the desired order.
    """
    try:
        ids = payload.get("ids", [])
        if not ids:
            raise HTTPException(status_code=400, detail="Missing ids list")
        playbook_store.reorder_playbooks(ids)
        return {"status": "reordered"}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/playbooks")
async def create_playbook(payload: Dict[str, Any] = Body(...)):
    try:
        if not payload.get("title") or not payload.get("goal") or not payload.get("instructions"):
            raise HTTPException(status_code=400, detail="Missing required fields: title, goal, instructions")
        payload["id"] = str(_uuid.uuid4())
        payload.setdefault("tags", [])
        created = _persist_playbook(payload)
        return created
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/playbooks/{id}")
async def get_playbook(id: UUID):
    try:
        pb, _ = _find_playbook_by_id(str(id))
        if not pb:
            raise HTTPException(status_code=404, detail="Not Found")

        # Build dict in the exact order you want
        result = {
            "id": getattr(pb, "id", None) or (pb.get("id") if isinstance(pb, dict) else None),
            "title": getattr(pb, "title", None) or (pb.get("title") if isinstance(pb, dict) else None),
            "goal": getattr(pb, "goal", None) or (pb.get("goal") if isinstance(pb, dict) else None),
            "instructions": getattr(pb, "instructions", "") or (pb.get("instructions") if isinstance(pb, dict) else ""),
            "tags": getattr(pb, "tags", []) or (pb.get("tags") if isinstance(pb, dict) else []),
        }
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.patch("/playbooks/{id}")
async def patch_playbook(id: UUID, payload: Dict[str, Any] = Body(...)):
    """
    Partial update: accepts a JSON object with fields to update (e.g., {"instructions":"..."}).
    """
    try:
        pb, key = _find_playbook_by_id(str(id))
        if not pb:
            raise HTTPException(status_code=404, detail="Not Found")

        # Normalize existing to dict
        if isinstance(pb, dict):
            existing = dict(pb)
        else:
            existing = {k: getattr(pb, k) for k in ("id", "title", "goal", "instructions", "tags") if hasattr(pb, k)}

        merged = {**existing, **payload}

        # Try to persist via store API
        updated = _persist_playbook(merged)
        if not isinstance(updated, dict):
            return {k: getattr(updated, k) for k in ("id", "title", "goal", "instructions", "tags") if hasattr(updated, k)}
        return updated
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.put("/playbooks/{id}")
async def put_playbook(id: UUID, payload: Dict[str, Any] = Body(...)):
    """
    Full replace: replace the playbook with the provided payload (payload should include title, goal, instructions, tags).
    """
    try:
        pb, key = _find_playbook_by_id(str(id))
        if not pb:
            raise HTTPException(status_code=404, detail="Not Found")

        payload["id"] = str(id)

        replaced = _persist_playbook(payload)
        if not isinstance(replaced, dict):
            return {k: getattr(replaced, k) for k in ("id", "title", "goal", "instructions", "tags") if hasattr(replaced, k)}
        return replaced
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/playbooks/{id}")
async def delete_playbook_endpoint(id: UUID):
    try:
        pb, _ = _find_playbook_by_id(str(id))
        if not pb:
            raise HTTPException(status_code=404, detail="Not Found")
        playbook_store.delete_playbook(str(id))
        return {"status": "deleted"}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# -------------------------
# Unified Search
# -------------------------
@app.get("/search")
async def search_endpoint(q: Optional[str] = None):
    """Search conversations and playbooks by keyword."""
    if not q or not q.strip():
        raise HTTPException(status_code=400, detail="Missing query parameter 'q'")
    try:
        conversations = store.search_conversations(q, limit=5)
        playbook_hits = playbook_store.search_playbooks(q)
        return {
            "query": q,
            "conversations": conversations,
            "playbooks": [
                {"id": p.id, "title": p.title, "goal": p.goal, "tags": p.tags}
                for p in playbook_hits
            ],
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# -------------------------
# Conversations
# -------------------------
@app.get("/conversations")
async def get_conversations(q: Optional[str] = None):
    try:
        conversations = store.all_conversations()
        if q:
            q_lower = q.lower()
            conversations = [
                c for c in conversations
                if any(q_lower in m.content.lower() for m in c.messages)
                or q_lower in c.preview.lower()
            ]
        return {
            "conversations": [
                {
                    "id": c.id,
                    "timestamp": c.created_at,
                    "updated_at": c.updated_at,
                    "preview": c.preview,
                }
                for c in conversations
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/conversations/{conversation_id}")
async def get_conversation(conversation_id: str):
    try:
        conv = store.get_conversation(conversation_id)
        if not conv:
            raise HTTPException(status_code=404, detail="Conversation not found")
        return {
            "id": conv.id,
            "timestamp": conv.created_at,
            "messages": [
                {"role": m.role, "content": m.content, "timestamp": m.timestamp, "model": m.model, "tokens": m.tokens}
                for m in conv.messages
            ]
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/conversations/{conversation_id}")
async def delete_conversation(conversation_id: str):
    try:
        conv = store.get_conversation(conversation_id)
        if not conv:
            raise HTTPException(status_code=404, detail="Conversation not found")
        store.delete_conversation(conversation_id)
        return {"status": "deleted"}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -------------------------
# End of file
# -------------------------