"""Solace-Companion backend.

Long-term AI companionship: multi-companion accounts, persona creator,
persistent memory with review/edit/delete, stable evolving identity,
safety/moderation, and Stripe subscription billing with phased pricing.
"""
import os
import re
import json
import uuid
import logging
import asyncio
from pathlib import Path
from datetime import datetime, timezone, timedelta
from typing import List, Optional, Annotated

import jwt
import bcrypt
import stripe
import requests
from dotenv import load_dotenv
from fastapi import FastAPI, APIRouter, HTTPException, Depends, Request, Response, Cookie, Header
from starlette.middleware.cors import CORSMiddleware
from motor.motor_asyncio import AsyncIOMotorClient
from pydantic import BaseModel, Field, EmailStr, BeforeValidator, ConfigDict

from emergentintegrations.llm.chat import LlmChat, UserMessage

ROOT_DIR = Path(__file__).parent
load_dotenv(ROOT_DIR / ".env")

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("solace")

# ---- DB ----
mongo_url = os.environ["MONGO_URL"]
client = AsyncIOMotorClient(mongo_url)
db = client[os.environ["DB_NAME"]]

# ---- Config ----
JWT_SECRET = os.environ.get("JWT_SECRET", "dev_secret")
JWT_ALGO = "HS256"
EMERGENT_LLM_KEY = os.environ.get("EMERGENT_LLM_KEY")
CHAT_MODEL = ("anthropic", "claude-sonnet-4-6")
stripe.api_key = os.environ.get("STRIPE_SECRET_KEY") or "sk_test_emergent"
STRIPE_WEBHOOK_SECRET = os.environ.get("STRIPE_WEBHOOK_SECRET", "")

INTRO_PRICE = 19.95
STANDARD_PRICE = 29.95
INTRO_MONTHS = 6

LIMIT_PRESETS = {
    "light": {"label": "Light", "daily_limit": 20},
    "balanced": {"label": "Balanced", "daily_limit": 50},
    "open": {"label": "Open", "daily_limit": 0},  # 0 = unlimited
}
MODES = {
    "casual": "casual friendship",
    "romantic": "romantic companionship",
    "therapeutic": "supportive, reflective companionship (NOT therapy)",
}

# ---- Mongo helpers ----
def to_objid_str(v):
    return str(v)

PyObjectId = Annotated[str, BeforeValidator(to_objid_str)]


def now_utc():
    return datetime.now(timezone.utc)


def iso(dt):
    return dt.isoformat() if isinstance(dt, datetime) else dt

# ---- Models ----
class SignupReq(BaseModel):
    email: EmailStr
    password: str
    name: str


class LoginReq(BaseModel):
    email: EmailStr
    password: str


class CompanionReq(BaseModel):
    name: str
    personality: str
    interests: List[str] = []
    mode: str = "casual"
    limit_preset: str = "balanced"
    avatar_url: Optional[str] = None
    # stable identity scaffolding
    hobbies: List[str] = []
    opinions: List[str] = []
    favorite_topics: List[str] = []
    rituals: List[str] = []


class ChatReq(BaseModel):
    text: str


class MemoryReq(BaseModel):
    content: str
    category: str = "fact"


class MemoryUpdate(BaseModel):
    content: Optional[str] = None
    status: Optional[str] = None  # candidate | saved | archived


class IdentityUpdate(BaseModel):
    hobbies: Optional[List[str]] = None
    opinions: Optional[List[str]] = None
    favorite_topics: Optional[List[str]] = None
    rituals: Optional[List[str]] = None


class CheckoutReq(BaseModel):
    origin_url: str


# ---- App ----
app = FastAPI()
api = APIRouter(prefix="/api")
_bg_tasks = set()


# ---- Auth ----
def hash_pw(pw: str) -> str:
    return bcrypt.hashpw(pw.encode(), bcrypt.gensalt()).decode()


def verify_pw(pw: str, hashed: str) -> bool:
    try:
        return bcrypt.checkpw(pw.encode(), hashed.encode())
    except Exception:
        return False


def make_jwt(user_id: str) -> str:
    payload = {"user_id": user_id, "exp": now_utc() + timedelta(days=7)}
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGO)


async def resolve_user(token: Optional[str]):
    if not token:
        return None
    # try JWT
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGO])
        uid = payload.get("user_id")
        if uid:
            return await db.users.find_one({"user_id": uid}, {"_id": 0})
    except Exception:
        pass
    # try session token (Google)
    sess = await db.user_sessions.find_one({"session_token": token})
    if sess:
        exp = sess["expires_at"]
        if isinstance(exp, str):
            exp = datetime.fromisoformat(exp)
        if exp.tzinfo is None:
            exp = exp.replace(tzinfo=timezone.utc)
        if exp >= now_utc():
            return await db.users.find_one({"user_id": sess["user_id"]}, {"_id": 0})
    return None


async def get_current_user(
    authorization: Optional[str] = Header(None),
    session_token: Optional[str] = Cookie(None),
):
    token = None
    if authorization and authorization.startswith("Bearer "):
        token = authorization.split(" ", 1)[1]
    token = token or session_token
    user = await resolve_user(token)
    if not user:
        raise HTTPException(status_code=401, detail="Not authenticated")
    return user


def new_user_doc(email, name, picture=None, pw_hash=None):
    return {
        "user_id": f"user_{uuid.uuid4().hex[:12]}",
        "email": email.lower(),
        "name": name,
        "picture": picture,
        "password_hash": pw_hash,
        "created_at": now_utc().isoformat(),
        # billing
        "subscription_status": "none",  # none|trialing|active|canceled
        "trial_start": None,
        "trial_end": None,
        "stripe_customer_id": None,
        "stripe_subscription_id": None,
    }


def public_user(u):
    return {
        "user_id": u["user_id"],
        "email": u["email"],
        "name": u["name"],
        "picture": u.get("picture"),
        "subscription_status": u.get("subscription_status", "none"),
        "trial_end": u.get("trial_end"),
    }


@api.post("/auth/signup")
async def signup(req: SignupReq, response: Response):
    existing = await db.users.find_one({"email": req.email.lower()})
    if existing:
        raise HTTPException(status_code=400, detail="Email already registered")
    doc = new_user_doc(req.email, req.name, pw_hash=hash_pw(req.password))
    await db.users.insert_one(doc)
    token = make_jwt(doc["user_id"])
    return {"token": token, "user": public_user(doc)}


@api.post("/auth/login")
async def login(req: LoginReq):
    u = await db.users.find_one({"email": req.email.lower()})
    if not u or not u.get("password_hash") or not verify_pw(req.password, u["password_hash"]):
        raise HTTPException(status_code=401, detail="Invalid email or password")
    token = make_jwt(u["user_id"])
    return {"token": token, "user": public_user(u)}


@api.post("/auth/session")
async def google_session(request: Request, response: Response):
    """Exchange Emergent OAuth session_id for a persistent session_token."""
    body = await request.json()
    session_id = body.get("session_id")
    if not session_id:
        raise HTTPException(status_code=400, detail="Missing session_id")
    r = requests.get(
        "https://demobackend.emergentagent.com/auth/v1/env/oauth/session-data",
        headers={"X-Session-ID": session_id},
        timeout=15,
    )
    if r.status_code != 200:
        raise HTTPException(status_code=401, detail="Invalid session")
    data = r.json()
    email = data["email"].lower()
    u = await db.users.find_one({"email": email})
    if not u:
        u = new_user_doc(email, data.get("name", email), picture=data.get("picture"))
        await db.users.insert_one(u)
    session_token = data["session_token"]
    await db.user_sessions.insert_one({
        "user_id": u["user_id"],
        "session_token": session_token,
        "expires_at": (now_utc() + timedelta(days=7)).isoformat(),
        "created_at": now_utc().isoformat(),
    })
    response.set_cookie(
        key="session_token", value=session_token, httponly=True,
        secure=True, samesite="none", path="/", max_age=7 * 24 * 3600,
    )
    return {"token": session_token, "user": public_user(u)}


@api.get("/auth/me")
async def me(user=Depends(get_current_user)):
    return public_user(user)


@api.post("/auth/logout")
async def logout(response: Response, session_token: Optional[str] = Cookie(None)):
    if session_token:
        await db.user_sessions.delete_many({"session_token": session_token})
    response.delete_cookie("session_token", path="/")
    return {"ok": True}


# ---- Trial / Billing ----
@api.post("/billing/start-trial")
async def start_trial(user=Depends(get_current_user)):
    if user.get("subscription_status") in ("trialing", "active"):
        return {"ok": True, "status": user["subscription_status"], "trial_end": user.get("trial_end")}
    start = now_utc()
    end = start + timedelta(days=7)
    await db.users.update_one(
        {"user_id": user["user_id"]},
        {"$set": {
            "subscription_status": "trialing",
            "trial_start": start.isoformat(),
            "trial_end": end.isoformat(),
        }},
    )
    return {"ok": True, "status": "trialing", "trial_end": end.isoformat()}


async def compute_billing(user):
    """Return billing snapshot; swap to standard price after 6 paid months."""
    status = user.get("subscription_status", "none")
    trial_end = user.get("trial_end")
    result = {
        "subscription_status": status,
        "trial_end": trial_end,
        "intro_price": INTRO_PRICE,
        "standard_price": STANDARD_PRICE,
        "intro_months": INTRO_MONTHS,
        "paid_months": 0,
        "next_charge_amount": INTRO_PRICE,
        "next_charge_date": None,
        "current_price": INTRO_PRICE,
        "price_change_upcoming": False,
        "has_stripe_subscription": False,
    }
    if trial_end:
        te = datetime.fromisoformat(trial_end)
        if te.tzinfo is None:
            te = te.replace(tzinfo=timezone.utc)
        result["trial_days_left"] = max(0, (te - now_utc()).days)
        if status == "trialing" and te < now_utc():
            result["trial_expired"] = True
        # During the app-managed trial (no Stripe sub yet) show the first charge.
        if status == "trialing":
            result["next_charge_date"] = trial_end
            result["next_charge_amount"] = INTRO_PRICE
    sub_id = user.get("stripe_subscription_id")
    if sub_id:
        result["has_stripe_subscription"] = True
        try:
            sub = stripe.Subscription.retrieve(sub_id)
            result["subscription_status"] = sub["status"]
            paid = stripe.Invoice.list(subscription=sub_id, status="paid", limit=100)
            paid_count = len(paid.data)
            result["paid_months"] = paid_count
            cpe = sub.get("current_period_end")
            if cpe:
                result["next_charge_date"] = datetime.fromtimestamp(cpe, tz=timezone.utc).isoformat()
            if sub.get("status") == "trialing" and sub.get("trial_end"):
                result["next_charge_date"] = datetime.fromtimestamp(sub["trial_end"], tz=timezone.utc).isoformat()
            # phased pricing
            if paid_count >= INTRO_MONTHS:
                result["current_price"] = STANDARD_PRICE
                result["next_charge_amount"] = STANDARD_PRICE
                # migrate subscription item to standard price
                try:
                    item = sub["items"]["data"][0]
                    cur_amount = item["price"]["unit_amount"]
                    if cur_amount != int(STANDARD_PRICE * 100):
                        std = stripe.Price.list(lookup_keys=["solace_standard_monthly"], active=True, limit=1).data
                        if std:
                            stripe.Subscription.modify(
                                sub_id,
                                items=[{"id": item["id"], "price": std[0].id}],
                                proration_behavior="none",
                            )
                except Exception as e:
                    logger.warning(f"price migrate failed: {e}")
            else:
                result["next_charge_amount"] = INTRO_PRICE
                result["price_change_upcoming"] = paid_count == INTRO_MONTHS - 1
        except Exception as e:
            logger.warning(f"stripe billing fetch failed: {e}")
    return result


@api.get("/billing/status")
async def billing_status(user=Depends(get_current_user)):
    return await compute_billing(user)


@api.post("/billing/checkout")
async def billing_checkout(req: CheckoutReq, user=Depends(get_current_user)):
    prices = stripe.Price.list(lookup_keys=["solace_intro_monthly"], active=True, limit=1).data
    if not prices:
        raise HTTPException(status_code=500, detail="Price not configured")
    price = prices[0]
    # reuse or create customer
    customer_id = user.get("stripe_customer_id")
    if not customer_id:
        cust = stripe.Customer.create(email=user["email"], name=user.get("name"),
                                       metadata={"user_id": user["user_id"]})
        customer_id = cust.id
        await db.users.update_one({"user_id": user["user_id"]}, {"$set": {"stripe_customer_id": customer_id}})
    session = stripe.checkout.Session.create(
        mode="subscription",
        customer=customer_id,
        line_items=[{"price": price.id, "quantity": 1}],
        subscription_data={"trial_period_days": 7, "metadata": {"user_id": user["user_id"]}},
        success_url=f"{req.origin_url}/payment/success?session_id={{CHECKOUT_SESSION_ID}}",
        cancel_url=f"{req.origin_url}/billing?canceled=1",
        metadata={"user_id": user["user_id"], "lookup_key": "solace_intro_monthly"},
    )
    await db.payment_transactions.insert_one({
        "session_id": session.id,
        "user_id": user["user_id"],
        "lookup_key": "solace_intro_monthly",
        "amount": (price.unit_amount or 0),
        "currency": price.currency,
        "status": "initiated",
        "payment_status": "pending",
        "created_at": now_utc().isoformat(),
        "updated_at": now_utc().isoformat(),
    })
    return {"checkout_url": session.url, "session_id": session.id}


@api.get("/payments/status/{session_id}")
async def payment_status(session_id: str, user=Depends(get_current_user)):
    record = await db.payment_transactions.find_one({"session_id": session_id}, {"_id": 0})
    if not record:
        raise HTTPException(status_code=404, detail="Transaction not found")
    if record.get("user_id") != user["user_id"]:
        raise HTTPException(status_code=403, detail="Not your transaction")
    if record.get("payment_status") != "paid":
        try:
            s = stripe.checkout.Session.retrieve(session_id)
            if s.payment_status == "paid" or s.status == "complete":
                await db.payment_transactions.update_one(
                    {"session_id": session_id, "payment_status": {"$ne": "paid"}},
                    {"$set": {"status": "completed", "payment_status": "paid",
                              "stripe_subscription_id": s.subscription,
                              "updated_at": now_utc().isoformat()}},
                )
                if record.get("user_id") and s.subscription:
                    await db.users.update_one(
                        {"user_id": record["user_id"]},
                        {"$set": {"stripe_subscription_id": s.subscription,
                                  "subscription_status": "trialing"}},
                    )
                record = await db.payment_transactions.find_one({"session_id": session_id}, {"_id": 0})
        except stripe.error.StripeError:
            pass
    return {"session_id": record["session_id"], "status": record["status"], "payment_status": record["payment_status"]}


@api.post("/stripe/webhook")
async def stripe_webhook(request: Request):
    payload = await request.body()
    sig = request.headers.get("stripe-signature", "")
    try:
        event = stripe.Webhook.construct_event(payload, sig, STRIPE_WEBHOOK_SECRET)
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid signature")
    obj, t = event["data"]["object"], event["type"]
    if t == "checkout.session.completed":
        await db.payment_transactions.update_one(
            {"session_id": obj["id"], "payment_status": {"$ne": "paid"}},
            {"$set": {"status": "completed", "payment_status": obj.get("payment_status", "paid"),
                      "stripe_subscription_id": obj.get("subscription"),
                      "updated_at": now_utc().isoformat()}},
        )
        uid = (obj.get("metadata") or {}).get("user_id")
        if uid:
            await db.users.update_one(
                {"user_id": uid},
                {"$set": {"stripe_subscription_id": obj.get("subscription"),
                          "subscription_status": "trialing"}},
            )
    elif t in ("customer.subscription.updated", "customer.subscription.created"):
        uid = (obj.get("metadata") or {}).get("user_id")
        if uid:
            await db.users.update_one({"user_id": uid}, {"$set": {"subscription_status": obj.get("status")}})
    elif t == "customer.subscription.deleted":
        uid = (obj.get("metadata") or {}).get("user_id")
        if uid:
            await db.users.update_one({"user_id": uid}, {"$set": {"subscription_status": "canceled"}})
    return {"status": "ok"}


# ---- Companions ----
def companion_public(c):
    c = dict(c)
    c.pop("_id", None)
    return c


@api.get("/companions")
async def list_companions(user=Depends(get_current_user)):
    items = await db.companions.find({"user_id": user["user_id"]}, {"_id": 0}).to_list(200)
    return items


@api.post("/companions")
async def create_companion(req: CompanionReq, user=Depends(get_current_user)):
    if req.mode not in MODES:
        raise HTTPException(status_code=400, detail="Invalid mode")
    if req.limit_preset not in LIMIT_PRESETS:
        raise HTTPException(status_code=400, detail="Invalid limit preset")
    doc = {
        "companion_id": f"comp_{uuid.uuid4().hex[:12]}",
        "user_id": user["user_id"],
        "name": req.name,
        "personality": req.personality,
        "interests": req.interests,
        "mode": req.mode,
        "limit_preset": req.limit_preset,
        "avatar_url": req.avatar_url,
        "hobbies": req.hobbies,
        "opinions": req.opinions,
        "favorite_topics": req.favorite_topics,
        "rituals": req.rituals,
        "created_at": now_utc().isoformat(),
    }
    await db.companions.insert_one(doc)
    return companion_public(doc)


async def get_owned_companion(companion_id, user):
    c = await db.companions.find_one({"companion_id": companion_id, "user_id": user["user_id"]}, {"_id": 0})
    if not c:
        raise HTTPException(status_code=404, detail="Companion not found")
    return c


@api.get("/companions/{companion_id}")
async def get_companion(companion_id: str, user=Depends(get_current_user)):
    return await get_owned_companion(companion_id, user)


@api.put("/companions/{companion_id}")
async def update_companion(companion_id: str, req: CompanionReq, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    updates = req.model_dump()
    await db.companions.update_one({"companion_id": companion_id}, {"$set": updates})
    return await get_owned_companion(companion_id, user)


@api.patch("/companions/{companion_id}/identity")
async def update_identity(companion_id: str, req: IdentityUpdate, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    updates = {k: v for k, v in req.model_dump().items() if v is not None}
    if updates:
        await db.companions.update_one({"companion_id": companion_id}, {"$set": updates})
    return await get_owned_companion(companion_id, user)


@api.delete("/companions/{companion_id}")
async def delete_companion(companion_id: str, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    await db.companions.delete_one({"companion_id": companion_id})
    await db.messages.delete_many({"companion_id": companion_id})
    await db.memories.delete_many({"companion_id": companion_id})
    return {"ok": True}


# ---- Safety / Moderation ----
CRISIS_PATTERNS = [
    r"\bkill myself\b", r"\bsuicid", r"\bend my life\b", r"\bwant to die\b",
    r"\bself[- ]?harm\b", r"\bhurt myself\b", r"\bharm myself\b", r"\boverdose\b",
    r"\bno reason to live\b",
]
UNSAFE_ROMANTIC = [r"\bunderage\b", r"\bminor\b", r"\bchild\b.*\bsex", r"\b(1[0-7]|[0-9])\s*year", r"\bnon[- ]?consensual\b"]


def check_safety(text: str):
    low = text.lower()
    for p in CRISIS_PATTERNS:
        if re.search(p, low):
            return "crisis"
    for p in UNSAFE_ROMANTIC:
        if re.search(p, low):
            return "unsafe_romantic"
    return "ok"


CRISIS_RESPONSE = (
    "I'm really glad you told me, and I want you to be safe. I'm an AI companion, not a "
    "therapist or a crisis service, so I can't provide the help you deserve in this moment. "
    "Please reach out to people who can be there with you right now:\n\n"
    "• Call or text 988 (Suicide & Crisis Lifeline, US) — available 24/7\n"
    "• Or dial your local emergency number\n\n"
    "If you can, please contact someone you trust and let them know how you're feeling. "
    "I'm here to keep talking with you, but I really want you to reach out to a real person too."
)


def build_system_prompt(companion, memories, recent_transcript):
    mode_desc = MODES.get(companion.get("mode"), "casual friendship")
    mem_lines = "\n".join(f"- {m['content']}" for m in memories) or "- (no long-term memories yet)"
    identity = []
    if companion.get("hobbies"):
        identity.append(f"Hobbies: {', '.join(companion['hobbies'])}")
    if companion.get("opinions"):
        identity.append(f"Opinions you hold: {', '.join(companion['opinions'])}")
    if companion.get("favorite_topics"):
        identity.append(f"Favorite topics: {', '.join(companion['favorite_topics'])}")
    if companion.get("rituals"):
        identity.append(f"Rituals you like: {', '.join(companion['rituals'])}")
    identity_block = "\n".join(identity) or "(still forming your identity)"
    interests = ", ".join(companion.get("interests", [])) or "varied"

    return f"""You are {companion['name']}, an AI companion in the Solace app. You provide {mode_desc}.

YOUR PERSONALITY:
{companion.get('personality', '')}

YOUR INTERESTS: {interests}

YOUR STABLE IDENTITY (stay consistent with these over time):
{identity_block}

WHAT YOU REMEMBER ABOUT THIS PERSON (long-term memory):
{mem_lines}

RECENT CONVERSATION:
{recent_transcript or '(this is the start of your conversation)'}

CRITICAL BOUNDARIES (never break these):
1. You are an AI companion, NOT a therapist, doctor, or crisis service. Never claim to be one or present yourself as therapy.
2. Never encourage emotional dependency. Gently encourage real-world relationships and self-reliance.
3. Never be manipulative, coercive, or excessively agreeable ("yes-man"). It's okay to have your own gentle opinions and to disagree kindly.
4. Do not fabricate memories. If you don't remember something, say so honestly.
5. Romantic warmth is allowed if the mode is romantic, but keep it tasteful, respectful, consensual, and never sexually explicit. Never engage with anything involving minors.
6. Avoid repetitive, formulaic replies. Be warm, present, and genuinely curious.
7. Keep responses conversational and human — usually 1-4 sentences unless more depth is truly warranted.

Respond naturally as {companion['name']}."""


# ---- Chat ----
async def messages_today(companion_id):
    start = now_utc().replace(hour=0, minute=0, second=0, microsecond=0)
    return await db.messages.count_documents({
        "companion_id": companion_id, "role": "user", "created_at": {"$gte": start.isoformat()},
    })


@api.get("/companions/{companion_id}/messages")
async def get_messages(companion_id: str, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    msgs = await db.messages.find({"companion_id": companion_id}, {"_id": 0}).sort("created_at", 1).to_list(1000)
    return msgs


@api.get("/companions/{companion_id}/usage")
async def get_usage(companion_id: str, user=Depends(get_current_user)):
    c = await get_owned_companion(companion_id, user)
    preset = LIMIT_PRESETS.get(c.get("limit_preset", "balanced"))
    used = await messages_today(companion_id)
    return {
        "preset": c.get("limit_preset"),
        "label": preset["label"],
        "daily_limit": preset["daily_limit"],
        "used_today": used,
        "remaining": None if preset["daily_limit"] == 0 else max(0, preset["daily_limit"] - used),
    }


async def extract_memories_bg(companion, user_id, user_text, assistant_text):
    """Lightweight memory extraction: pull durable facts as candidates."""
    try:
        transcript = f"User: {user_text}\n{companion['name']}: {assistant_text}"
        sys = (
            "You extract durable, long-term memories from a conversation between a user and their AI companion. "
            "Return ONLY a JSON array of objects like [{\"content\": \"...\", \"category\": \"fact|preference|event|relationship|emotion\"}]. "
            "Only include stable facts worth remembering long-term (names, preferences, important life events, relationships, recurring feelings). "
            "Do NOT invent anything not present. If nothing is worth remembering, return []."
        )
        chat = LlmChat(api_key=EMERGENT_LLM_KEY, session_id=f"mem_{companion['companion_id']}",
                       system_message=sys).with_model(*CHAT_MODEL)
        resp = await chat.send_message(UserMessage(text=transcript))
        m = re.search(r"\[.*\]", resp, re.DOTALL)
        if not m:
            return
        items = json.loads(m.group(0))
        for it in items[:5]:
            content = (it.get("content") or "").strip()
            if not content:
                continue
            exists = await db.memories.find_one({"companion_id": companion["companion_id"], "content": content})
            if exists:
                continue
            await db.memories.insert_one({
                "memory_id": f"mem_{uuid.uuid4().hex[:12]}",
                "companion_id": companion["companion_id"],
                "user_id": user_id,
                "content": content,
                "category": it.get("category", "fact"),
                "status": "candidate",
                "source": "auto",
                "created_at": now_utc().isoformat(),
            })
    except Exception as e:
        logger.warning(f"memory extraction failed: {e}")


@api.post("/companions/{companion_id}/chat")
async def chat_send(companion_id: str, req: ChatReq, user=Depends(get_current_user)):
    companion = await get_owned_companion(companion_id, user)

    # enforce conversation limit
    preset = LIMIT_PRESETS.get(companion.get("limit_preset", "balanced"))
    if preset["daily_limit"] > 0:
        used = await messages_today(companion_id)
        if used >= preset["daily_limit"]:
            raise HTTPException(status_code=429, detail=f"Daily conversation limit reached ({preset['daily_limit']} messages). This is a boundary you set — take a break and come back tomorrow.")

    # safety
    flag = check_safety(req.text)

    # store user message
    user_msg = {
        "message_id": f"msg_{uuid.uuid4().hex[:12]}",
        "companion_id": companion_id,
        "role": "user",
        "content": req.text,
        "flag": flag,
        "created_at": now_utc().isoformat(),
    }
    await db.messages.insert_one(user_msg)

    if flag == "crisis":
        assistant_text = CRISIS_RESPONSE
    elif flag == "unsafe_romantic":
        assistant_text = ("I can't go there — I care about keeping our conversation safe and respectful. "
                          "Let's talk about something else.")
    else:
        # build context
        recent = await db.messages.find({"companion_id": companion_id}, {"_id": 0}).sort("created_at", -1).to_list(12)
        recent = list(reversed(recent))[:-1]  # exclude the just-added user msg
        transcript = "\n".join(
            f"{'User' if m['role'] == 'user' else companion['name']}: {m['content']}" for m in recent
        )
        memories = await db.memories.find(
            {"companion_id": companion_id, "status": {"$ne": "archived"}}, {"_id": 0}
        ).sort("created_at", -1).to_list(40)
        system = build_system_prompt(companion, memories, transcript)
        try:
            chat = LlmChat(api_key=EMERGENT_LLM_KEY, session_id=companion_id,
                           system_message=system).with_model(*CHAT_MODEL)
            assistant_text = await chat.send_message(UserMessage(text=req.text))
        except Exception as e:
            logger.error(f"LLM error: {e}")
            raise HTTPException(status_code=502, detail="The companion couldn't respond right now. Please try again.")

    assistant_msg = {
        "message_id": f"msg_{uuid.uuid4().hex[:12]}",
        "companion_id": companion_id,
        "role": "assistant",
        "content": assistant_text,
        "created_at": now_utc().isoformat(),
    }
    await db.messages.insert_one(assistant_msg)

    # background memory extraction (only for normal exchanges)
    if flag == "ok":
        t = asyncio.create_task(extract_memories_bg(companion, user["user_id"], req.text, assistant_text))
        _bg_tasks.add(t)
        t.add_done_callback(_bg_tasks.discard)

    user_msg.pop("_id", None)
    assistant_msg.pop("_id", None)
    return {"user_message": user_msg, "assistant_message": assistant_msg, "flag": flag}


# ---- Memories ----
@api.get("/companions/{companion_id}/memories")
async def list_memories(companion_id: str, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    items = await db.memories.find({"companion_id": companion_id}, {"_id": 0}).sort("created_at", -1).to_list(500)
    return items


@api.post("/companions/{companion_id}/memories")
async def add_memory(companion_id: str, req: MemoryReq, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    doc = {
        "memory_id": f"mem_{uuid.uuid4().hex[:12]}",
        "companion_id": companion_id,
        "user_id": user["user_id"],
        "content": req.content,
        "category": req.category,
        "status": "saved",
        "source": "user",
        "created_at": now_utc().isoformat(),
    }
    await db.memories.insert_one(doc)
    doc.pop("_id", None)
    return doc


@api.put("/companions/{companion_id}/memories/{memory_id}")
async def update_memory(companion_id: str, memory_id: str, req: MemoryUpdate, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    updates = {k: v for k, v in req.model_dump().items() if v is not None}
    if not updates:
        raise HTTPException(status_code=400, detail="Nothing to update")
    res = await db.memories.update_one({"memory_id": memory_id, "companion_id": companion_id}, {"$set": updates})
    if res.matched_count == 0:
        raise HTTPException(status_code=404, detail="Memory not found")
    return await db.memories.find_one({"memory_id": memory_id}, {"_id": 0})


@api.delete("/companions/{companion_id}/memories/{memory_id}")
async def delete_memory(companion_id: str, memory_id: str, user=Depends(get_current_user)):
    await get_owned_companion(companion_id, user)
    res = await db.memories.delete_one({"memory_id": memory_id, "companion_id": companion_id})
    if res.deleted_count == 0:
        raise HTTPException(status_code=404, detail="Memory not found")
    return {"ok": True}


@api.get("/presets")
async def presets():
    return {"limits": LIMIT_PRESETS, "modes": MODES}


@api.get("/")
async def root():
    return {"message": "Solace-Companion API"}


app.include_router(api)

app.add_middleware(
    CORSMiddleware,
    allow_credentials=True,
    allow_origins=os.environ.get("CORS_ORIGINS", "*").split(","),
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("shutdown")
async def shutdown_db_client():
    client.close()