database.py

import os
from google.cloud import firestore
from google.genai import types

class Database:
    def __init__(self):
        try:
            self.db = firestore.Client()
        except Exception as e:
            print(f"Firestore initialization failed (running locally without GCP auth?): {e}")
            self.db = None

    def _serialize_history(self, history: list[types.Content]) -> list[dict]:
        serialized = []
        for content in history:
            parts = []
            if not content.parts:
                continue
            for part in content.parts:
                if getattr(part, "text", None):
                    parts.append({"text": part.text})
                # Skip heavy binary audio/image data for persistent text memory
            if parts:
                serialized.append({"role": content.role, "parts": parts})
        return serialized

    def _deserialize_history(self, serialized: list[dict]) -> list[types.Content]:
        history = []
        for item in serialized:
            parts = []
            for p in item.get("parts", []):
                # google-genai 0.3.0 uses types.Part(text="...") or from_text
                parts.append(types.Part.from_text(text=p["text"]))
            history.append(types.Content(role=item["role"], parts=parts))
        return history

    def get_chat_history(self, user_id: int) -> list[types.Content]:
        if not self.db:
            return []
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc = doc_ref.get()
            if doc.exists:
                serialized = doc.to_dict().get("history", [])
                return self._deserialize_history(serialized)
        except Exception as e:
            print(f"Error reading from Firestore: {e}")
        return []

    def save_chat_history(self, user_id: int, history: list[types.Content]):
        if not self.db:
            return
        try:
            serialized = self._serialize_history(history)
            doc_ref = self.db.collection("users").document(str(user_id))
            doc_ref.set({"history": serialized}, merge=True)
        except Exception as e:
            print(f"Error writing to Firestore: {e}")
            
    def get_all_user_ids(self) -> list[int]:
        if not self.db:
            return []
        try:
            users = self.db.collection("users").stream()
            return [int(user.id) for user in users]
        except Exception as e:
            print(f"Error fetching users from Firestore: {e}")
            return []

    def add_todo(self, user_id: int, task: str) -> str:
        """Adds a task to the user's todo list."""
        if not self.db:
            return "Database not initialized."
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc_ref.set({"todos": firestore.ArrayUnion([task])}, merge=True)
            return f"Task '{task}' added successfully."
        except Exception as e:
            return f"Error adding task: {e}"

    def get_todos(self, user_id: int) -> list[str]:
        """Gets all tasks in the user's todo list."""
        if not self.db:
            return []
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc = doc_ref.get()
            if doc.exists:
                return doc.to_dict().get("todos", [])
        except Exception as e:
            print(f"Error fetching todos: {e}")
        return []

    def clear_todos(self, user_id: int) -> str:
        """Clears all tasks from the user's todo list."""
        if not self.db:
            return "Database not initialized."
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc_ref.set({"todos": []}, merge=True)
            return "To-Do list cleared."
        except Exception as e:
            return f"Error clearing todos: {e}"

    def add_expense(self, user_id: int, amount: float, category: str) -> str:
        """Adds an expense record for the user."""
        if not self.db:
            return "Database not initialized."
        try:
            expense_ref = self.db.collection("users").document(str(user_id)).collection("expenses").document()
            expense_ref.set({
                "amount": amount,
                "category": category,
                "timestamp": firestore.SERVER_TIMESTAMP
            })
            return f"Expense of {amount} added to {category}."
        except Exception as e:
            return f"Error adding expense: {e}"

    def get_expenses_summary(self, user_id: int) -> str:
        """Gets a summary of all expenses for the user."""
        if not self.db:
            return "Database not initialized."
        try:
            expenses_ref = self.db.collection("users").document(str(user_id)).collection("expenses")
            expenses = expenses_ref.get()
            total = 0.0
            categories = {}
            for exp in expenses:
                data = exp.to_dict()
                amt = float(data.get("amount", 0))
                cat = data.get("category", "Other")
                total += amt
                categories[cat] = categories.get(cat, 0) + amt
            
            if total == 0:
                return "No expenses recorded yet."
            
            summary = f"Total Expenses: {total}.\nBreakdown:\n"
            for k, v in categories.items():
                summary += f"- {k}: {v}\n"
            return summary
        except Exception as e:
            return f"Error fetching expenses: {e}"

    def get_raw_expenses(self, user_id: int) -> dict:
        """Gets raw expenses by category for charting."""
        if not self.db:
            return {}
        try:
            expenses_ref = self.db.collection("users").document(str(user_id)).collection("expenses")
            expenses = expenses_ref.get()
            categories = {}
            for exp in expenses:
                data = exp.to_dict()
                amt = float(data.get("amount", 0))
                cat = data.get("category", "Other")
                categories[cat] = categories.get(cat, 0) + amt
            return categories
        except Exception as e:
            print(f"Error fetching raw expenses: {e}")
            return {}

    def add_habit(self, user_id: int, habit: str) -> str:
        """Adds a habit to the user's habit tracking list."""
        if not self.db:
            return "Database not initialized."
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc_ref.set({"habits": firestore.ArrayUnion([habit])}, merge=True)
            return f"Habit '{habit}' added successfully."
        except Exception as e:
            return f"Error adding habit: {e}"

    def get_habits(self, user_id: int) -> list[str]:
        """Gets all habits in the user's habit tracking list."""
        if not self.db:
            return []
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc = doc_ref.get()
            if doc.exists:
                return doc.to_dict().get("habits", [])
        except Exception as e:
            print(f"Error fetching habits: {e}")
        return []

    def log_workout(self, user_id: int, workout_details: str) -> str:
        """Logs a workout for the user."""
        if not self.db:
            return "Database not initialized."
        try:
            workout_ref = self.db.collection("users").document(str(user_id)).collection("workouts").document()
            workout_ref.set({
                "details": workout_details,
                "timestamp": firestore.SERVER_TIMESTAMP
            })
            return f"Workout logged: {workout_details}."
        except Exception as e:
            return f"Error logging workout: {e}"

    def remember_fact(self, user_id: int, key: str, value: str) -> str:
        """Remembers a fact about the user in a key-value store."""
        if not self.db:
            return "Database not initialized."
        try:
            doc_ref = self.db.collection("users").document(str(user_id)).collection("memory").document(key)
            doc_ref.set({"value": value}, merge=True)
            return f"Successfully remembered: {key} = {value}"
        except Exception as e:
            return f"Error remembering fact: {e}"

    def recall_fact(self, user_id: int, key: str) -> str:
        """Recalls a specific fact about the user by key."""
        if not self.db:
            return "Database not initialized."
        try:
            doc_ref = self.db.collection("users").document(str(user_id)).collection("memory").document(key)
            doc = doc_ref.get()
            if doc.exists:
                return f"Fact '{key}': {doc.to_dict().get('value')}"
            return f"No memory found for '{key}'."
        except Exception as e:
            return f"Error recalling fact: {e}"

    def get_all_facts(self, user_id: int) -> str:
        """Recalls all saved facts about the user."""
        if not self.db:
            return "Database not initialized."
        try:
            mem_ref = self.db.collection("users").document(str(user_id)).collection("memory")
            memories = mem_ref.get()
            facts = [f"{m.id}: {m.to_dict().get('value')}" for m in memories]
            if not facts:
                return "No facts remembered yet."
            return "\n".join(facts)
        except Exception as e:
            return f"Error recalling facts: {e}"

    # ===== XP & GAMIFICATION =====

    def _xp_to_level(self, xp: int) -> dict:
        """Maps XP to a level name."""
        levels = [
            (1000, 5, "👑 Legend Macha"),
            (600,  4, "💪 Mass Bro"),
            (300,  3, "🔥 Sema Fellow"),
            (100,  2, "⚔️ Decent Da"),
            (0,    1, "🥉 Rookie Bro"),
        ]
        for threshold, level, name in levels:
            if xp >= threshold:
                return {"level": level, "level_name": name}
        return {"level": 1, "level_name": "🥉 Rookie Bro"}

    def get_user_xp(self, user_id: int) -> dict:
        """Returns current XP and level info for user."""
        if not self.db:
            return {"xp": 0}
        try:
            doc = self.db.collection("users").document(str(user_id)).get()
            xp = doc.to_dict().get("xp", 0) if doc.exists else 0
            return {"xp": xp, **self._xp_to_level(xp)}
        except Exception as e:
            return {"xp": 0}

    def add_xp(self, user_id: int, points: int) -> str:
        """Awards XP to the user and returns a status string."""
        if not self.db:
            return "Database not initialized."
        try:
            doc_ref = self.db.collection("users").document(str(user_id))
            doc = doc_ref.get()
            current_xp = doc.to_dict().get("xp", 0) if doc.exists else 0
            new_xp = current_xp + points
            doc_ref.set({"xp": new_xp}, merge=True)
            info = self._xp_to_level(new_xp)
            return f"+{points} XP! Total: {new_xp} XP | Level {info['level']}: {info['level_name']}"
        except Exception as e:
            return f"Error adding XP: {e}"

    def get_daily_challenge(self, user_id: int) -> dict:
        """Returns the current daily challenge dict."""
        if not self.db:
            return {}
        try:
            doc = self.db.collection("users").document(str(user_id)).get()
            return (doc.to_dict() or {}).get("daily_challenge", {}) if doc.exists else {}
        except Exception:
            return {}

    def set_daily_challenge(self, user_id: int, challenge: str, date_str: str) -> str:
        """Sets a new daily challenge for the user."""
        if not self.db:
            return "Database not initialized."
        try:
            self.db.collection("users").document(str(user_id)).set(
                {"daily_challenge": {"text": challenge, "date": date_str, "accepted": False, "completed": False}},
                merge=True
            )
            return f"Challenge set: {challenge}"
        except Exception as e:
            return f"Error: {e}"

    def accept_daily_challenge(self, user_id: int) -> str:
        """Marks today's challenge as accepted."""
        if not self.db:
            return "Database not initialized."
        try:
            self.db.collection("users").document(str(user_id)).update({"daily_challenge.accepted": True})
            return "Challenge accepted! Go for it macha 💪"
        except Exception as e:
            return f"Error: {e}"

    def complete_daily_challenge(self, user_id: int) -> str:
        """Marks today's challenge as completed and awards 50 XP."""
        if not self.db:
            return "Database not initialized."
        try:
            ch = self.get_daily_challenge(user_id)
            if not ch:
                return "No active challenge found."
            if ch.get("completed"):
                return "Challenge already completed today!"
            self.db.collection("users").document(str(user_id)).update({"daily_challenge.completed": True})
            return self.add_xp(user_id, 50) + " 🎉 Challenge complete!"
        except Exception as e:
            return f"Error: {e}"

    # ===== PERSONAL KNOWLEDGE BASE =====

    def save_knowledge(self, user_id: int, topic: str, content: str) -> str:
        """Saves a knowledge entry under a topic key."""
        if not self.db:
            return "Database not initialized."
        try:
            key = topic.lower().replace(" ", "_")[:100]
            self.db.collection("users").document(str(user_id)).collection("knowledge").document(key).set(
                {"topic": topic, "content": content, "timestamp": firestore.SERVER_TIMESTAMP}, merge=True
            )
            return f"Knowledge saved: '{topic}'"
        except Exception as e:
            return f"Error: {e}"

    def search_knowledge(self, user_id: int, query: str) -> str:
        """Keyword-searches the user's knowledge base."""
        if not self.db:
            return "No knowledge base."
        try:
            docs = self.db.collection("users").document(str(user_id)).collection("knowledge").stream()
            q = query.lower()
            results = [
                f"**{d.to_dict()['topic']}**: {d.to_dict()['content'][:300]}"
                for d in docs
                if q in d.to_dict().get("topic", "").lower() or q in d.to_dict().get("content", "").lower()
            ]
            return "\n\n".join(results[:3]) if results else "No matching knowledge found."
        except Exception as e:
            return f"Error: {e}"

    def get_all_knowledge(self, user_id: int) -> str:
        """Returns all knowledge base topics."""
        if not self.db:
            return "No knowledge base."
        try:
            docs = list(self.db.collection("users").document(str(user_id)).collection("knowledge").stream())
            if not docs:
                return "Knowledge base is empty."
            return "\n".join(f"• {d.to_dict()['topic']}" for d in docs)
        except Exception as e:
            return f"Error: {e}"

    # ===== SMART NOTES =====

    def save_note(self, user_id: int, text: str, tags: list) -> str:
        """Saves a smart note with AI-assigned tags."""
        if not self.db:
            return "Database not initialized."
        try:
            self.db.collection("users").document(str(user_id)).collection("notes").document().set(
                {"text": text, "tags": tags, "timestamp": firestore.SERVER_TIMESTAMP}
            )
            return f"Note saved! Tags: {', '.join(tags) if tags else 'general'}"
        except Exception as e:
            return f"Error: {e}"

    def search_notes(self, user_id: int, query: str) -> str:
        """Searches notes by keyword or tag."""
        if not self.db:
            return "No notes found."
        try:
            docs = self.db.collection("users").document(str(user_id)).collection("notes")\
                .order_by("timestamp", direction=firestore.Query.DESCENDING).limit(50).stream()
            q = query.lower()
            results = []
            for d in docs:
                data = d.to_dict()
                if q in data.get("text", "").lower() or any(q in t.lower() for t in data.get("tags", [])):
                    results.append(f"📝 {data['text'][:200]}  [Tags: {', '.join(data.get('tags', []))}]")
            return "\n\n".join(results[:5]) if results else "No matching notes."
        except Exception as e:
            return f"Error: {e}"

    def get_recent_notes(self, user_id: int, limit: int = 5) -> str:
        """Gets the most recent N notes."""
        if not self.db:
            return "No notes."
        try:
            docs = self.db.collection("users").document(str(user_id)).collection("notes")\
                .order_by("timestamp", direction=firestore.Query.DESCENDING).limit(limit).stream()
            results = [
                f"📝 {d.to_dict()['text'][:200]}  [Tags: {', '.join(d.to_dict().get('tags', []))}]"
                for d in docs
            ]
            return "\n\n".join(results) if results else "No notes saved yet."
        except Exception as e:
            return f"Error: {e}"

    # ===== DATE-RANGE EXPENSE REPORTS =====

    def get_expenses_by_date_range(self, user_id: int, start_dt, end_dt) -> dict:
        """Returns expense totals grouped by category for a date range."""
        if not self.db:
            return {}
        try:
            ref = self.db.collection("users").document(str(user_id)).collection("expenses")
            docs = ref.where("timestamp", ">=", start_dt).where("timestamp", "<=", end_dt).stream()
            categories = {}
            total = 0.0
            for d in docs:
                data = d.to_dict()
                amt = float(data.get("amount", 0))
                cat = data.get("category", "Other")
                total += amt
                categories[cat] = categories.get(cat, 0) + amt
            return {"total": total, "categories": categories}
        except Exception as e:
            print(f"Date-range expense error: {e}")
            return {}

    # ===== SUBSCRIPTION TRACKER =====

    # Auto-assigned color palette for subscriptions (12 vivid colors)
    _SUB_COLORS = [
        "#E50914", "#1DB954", "#00A8E1", "#FF9900", "#6A0DAD",
        "#FF6B6B", "#00BFFF", "#FFD700", "#32CD32", "#FF1493",
        "#00CED1", "#FF8C00"
    ]

    def add_subscription(self, user_id: int, name: str, amount: float,
                         billing_day: int, cycle: str = "monthly",
                         currency: str = "INR", billing_month: int = None) -> str:
        """Adds or updates a recurring subscription for the user."""
        if not self.db:
            return "Database not initialized."
        try:
            key = name.lower().replace(" ", "_")[:60]
            # Assign a color from the palette based on existing count
            existing = list(self.db.collection("users").document(str(user_id))
                            .collection("subscriptions").stream())
            color = self._SUB_COLORS[len(existing) % len(self._SUB_COLORS)]
            data = {
                "name": name,
                "amount": float(amount),
                "currency": currency,
                "billing_day": int(billing_day),
                "cycle": cycle,
                "color": color,
                "active": True,
                "timestamp": firestore.SERVER_TIMESTAMP,
            }
            if billing_month:
                data["billing_month"] = int(billing_month)
            self.db.collection("users").document(str(user_id)) \
                .collection("subscriptions").document(key).set(data, merge=True)
            symbol = "₹" if currency == "INR" else currency
            return f"Subscription '{name}' saved! {symbol}{amount:.0f}/{cycle} on day {billing_day}."
        except Exception as e:
            return f"Error adding subscription: {e}"

    def get_subscriptions(self, user_id: int) -> list:
        """Returns all active subscriptions for the user as a list of dicts."""
        if not self.db:
            return []
        try:
            docs = self.db.collection("users").document(str(user_id)) \
                .collection("subscriptions") \
                .where("active", "==", True).stream()
            return [d.to_dict() for d in docs]
        except Exception as e:
            print(f"Error fetching subscriptions: {e}")
            return []

    def remove_subscription(self, user_id: int, name: str) -> str:
        """Marks a subscription as inactive (soft delete)."""
        if not self.db:
            return "Database not initialized."
        try:
            key = name.lower().replace(" ", "_")[:60]
            ref = self.db.collection("users").document(str(user_id)) \
                .collection("subscriptions").document(key)
            doc = ref.get()
            if not doc.exists:
                return f"No subscription named '{name}' found."
            ref.update({"active": False})
            return f"Subscription '{name}' removed successfully."
        except Exception as e:
            return f"Error removing subscription: {e}"

    def get_monthly_subscription_total(self, user_id: int) -> float:
        """Calculates total monthly spend across all active subscriptions."""
        subs = self.get_subscriptions(user_id)
        total = 0.0
        for s in subs:
            amt = float(s.get("amount", 0))
            cycle = s.get("cycle", "monthly")
            if cycle == "yearly":
                total += amt / 12
            elif cycle == "weekly":
                total += amt * 4.33
            else:
                total += amt
        return total


db = Database()