"""MySQL Storage Layer"""
import uuid
import json
import logging
from datetime import datetime
from typing import Optional, List, Dict, Any

import pymysql
from pymysql.cursors import DictCursor

logger = logging.getLogger(__name__)


class MySQLStore:
    """MySQL-backed persistent storage for memories."""

    def __init__(self, host, port, user, password, database, charset="utf8mb4", unix_socket=None):
        self.conn_kwargs = dict(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            charset=charset,
            cursorclass=DictCursor,
            autocommit=True,
        )
        if unix_socket:
            self.conn_kwargs["unix_socket"] = unix_socket
        self._conn = None

    def _get_conn(self):
        if self._conn is None or not self._conn.open:
            self._conn = pymysql.connect(**self.conn_kwargs)
        try:
            self._conn.ping(reconnect=True)
        except Exception:
            self._conn = pymysql.connect(**self.conn_kwargs)
        return self._conn

    def _query(self, sql, args=None, fetch=False):
        conn = self._get_conn()
        with conn.cursor() as cur:
            cur.execute(sql, args)
            if fetch:
                return cur.fetchall()
        return None

    # ── Team management ──────────────────────────────────────────

    def create_team(self, team_id: str, name: str, description: str = "", config: dict = None):
        self._query(
            "INSERT INTO teams (id, name, description, config) VALUES (%s, %s, %s, %s)",
            (team_id, name, description, json.dumps(config) if config else None),
        )
        return self.get_team(team_id)

    def get_team(self, team_id: str) -> Optional[dict]:
        rows = self._query("SELECT * FROM teams WHERE id = %s", (team_id,), fetch=True)
        return rows[0] if rows else None

    def delete_team(self, team_id: str) -> bool:
        self._query("DELETE FROM teams WHERE id = %s", (team_id,))
        return True

    # ── Agent management ─────────────────────────────────────────

    def create_agent(self, agent_id: str, team_id: str, name: str, role: str = ""):
        self._query(
            "INSERT INTO agents (id, team_id, name, role) VALUES (%s, %s, %s, %s)",
            (agent_id, team_id, name, role),
        )
        return self.get_agent(agent_id)

    def get_agent(self, agent_id: str) -> Optional[dict]:
        rows = self._query("SELECT * FROM agents WHERE id = %s", (agent_id,), fetch=True)
        return rows[0] if rows else None

    def delete_agent(self, agent_id: str) -> bool:
        self._query("DELETE FROM agents WHERE id = %s", (agent_id,))
        return True

    def get_agents_by_team(self, team_id: str) -> list:
        return self._query(
            "SELECT * FROM agents WHERE team_id = %s ORDER BY created_at DESC",
            (team_id,), fetch=True,
        )

    # ── Personal memories ────────────────────────────────────────

    def add_personal_memory(self, agent_id: str, team_id: str, content: str,
                            embedding: bytes, importance: float = 0.5,
                            metadata: dict = None) -> dict:
        mid = str(uuid.uuid4())[:16]
        self._query(
            """INSERT INTO personal_memories
               (id, agent_id, team_id, content, embedding, importance, metadata)
               VALUES (%s, %s, %s, %s, %s, %s, %s)""",
            (mid, agent_id, team_id, content, embedding, importance,
             json.dumps(metadata) if metadata else None),
        )
        return self.get_personal_memory(mid)

    def get_personal_memory(self, memory_id: str) -> Optional[dict]:
        rows = self._query("SELECT * FROM personal_memories WHERE id = %s", (memory_id,), fetch=True)
        return rows[0] if rows else None

    def get_personal_memories_by_agent(self, agent_id: str, limit: int = 20) -> List[dict]:
        return self._query(
            "SELECT * FROM personal_memories WHERE agent_id = %s ORDER BY created_at DESC LIMIT %s",
            (agent_id, limit), fetch=True,
        )

    def get_personal_memories_with_embeddings(self, agent_id: str) -> List[dict]:
        """Get all personal memories with embeddings for vector search."""
        return self._query(
            "SELECT id, content, embedding, importance, metadata FROM personal_memories WHERE agent_id = %s",
            (agent_id,), fetch=True,
        )

    def update_personal_memory(self, memory_id: str, content: str = None,
                               importance: float = None, metadata: dict = None,
                               embedding: bytes = None) -> bool:
        sets, args = [], []
        if content is not None:
            sets.append("content = %s")
            args.append(content)
        if importance is not None:
            sets.append("importance = %s")
            args.append(importance)
        if metadata is not None:
            sets.append("metadata = %s")
            args.append(json.dumps(metadata))
        if embedding is not None:
            sets.append("embedding = %s")
            args.append(embedding)
        if not sets:
            return False
        args.append(memory_id)
        self._query(f"UPDATE personal_memories SET {', '.join(sets)} WHERE id = %s", args)
        return True

    def touch_personal_memory(self, memory_id: str):
        self._query(
            "UPDATE personal_memories SET access_count = access_count + 1, last_accessed = NOW() WHERE id = %s",
            (memory_id,),
        )

    def delete_personal_memory(self, memory_id: str) -> bool:
        self._query("DELETE FROM personal_memories WHERE id = %s", (memory_id,))
        return True

    # ── Team memories ────────────────────────────────────────────

    def add_team_memory(self, team_id: str, content: str, embedding: bytes,
                        importance: float = 0.5, category: str = "general",
                        metadata: dict = None) -> dict:
        mid = str(uuid.uuid4())[:16]
        self._query(
            """INSERT INTO team_memories
               (id, team_id, content, embedding, importance, category, metadata)
               VALUES (%s, %s, %s, %s, %s, %s, %s)""",
            (mid, team_id, content, embedding, importance, category,
             json.dumps(metadata) if metadata else None),
        )
        return self.get_team_memory(mid)

    def get_team_memory(self, memory_id: str) -> Optional[dict]:
        rows = self._query("SELECT * FROM team_memories WHERE id = %s", (memory_id,), fetch=True)
        return rows[0] if rows else None

    def get_team_memories_by_team(self, team_id: str, limit: int = 20) -> List[dict]:
        return self._query(
            "SELECT * FROM team_memories WHERE team_id = %s ORDER BY created_at DESC LIMIT %s",
            (team_id, limit), fetch=True,
        )

    def get_team_memories_with_embeddings(self, team_id: str) -> List[dict]:
        """Get all team memories with embeddings for vector search."""
        return self._query(
            "SELECT id, content, embedding, importance, metadata FROM team_memories WHERE team_id = %s",
            (team_id,), fetch=True,
        )

    def update_team_memory(self, memory_id: str, content: str = None,
                           importance: float = None, embedding: bytes = None) -> bool:
        sets, args = [], []
        if content is not None:
            sets.append("content = %s")
            args.append(content)
        if importance is not None:
            sets.append("importance = %s")
            args.append(importance)
        if embedding is not None:
            sets.append("embedding = %s")
            args.append(embedding)
        if not sets:
            return False
        args.append(memory_id)
        self._query(f"UPDATE team_memories SET {', '.join(sets)} WHERE id = %s", args)
        return True

    def touch_team_memory(self, memory_id: str):
        self._query(
            "UPDATE team_memories SET access_count = access_count + 1, last_accessed = NOW() WHERE id = %s",
            (memory_id,),
        )

    def delete_team_memory(self, memory_id: str) -> bool:
        self._query("DELETE FROM team_memories WHERE id = %s", (memory_id,))
        return True

    # ── Cleanup ──────────────────────────────────────────────────

    def cleanup_personal_memories(self, team_id: str = None, max_age_days: int = 90,
                                  min_importance: float = 0.2) -> int:
        sql = ("DELETE FROM personal_memories WHERE importance < %s "
               "AND COALESCE(last_accessed, created_at) < DATE_SUB(NOW(), INTERVAL %s DAY)")
        args = [min_importance, max_age_days]
        if team_id:
            sql += " AND team_id = %s"
            args.append(team_id)
        conn = self._get_conn()
        with conn.cursor() as cur:
            cur.execute(sql, args)
            return cur.rowcount

    def cleanup_team_memories(self, team_id: str = None, max_age_days: int = 90,
                              min_importance: float = 0.2) -> int:
        sql = ("DELETE FROM team_memories WHERE importance < %s "
               "AND COALESCE(last_accessed, created_at) < DATE_SUB(NOW(), INTERVAL %s DAY)")
        args = [min_importance, max_age_days]
        if team_id:
            sql += " AND team_id = %s"
            args.append(team_id)
        conn = self._get_conn()
        with conn.cursor() as cur:
            cur.execute(sql, args)
            return cur.rowcount

    # ── Stats ────────────────────────────────────────────────────

    def get_stats(self, team_id: str = None) -> dict:
        stats = {}
        if team_id:
            row = self._query(
                "SELECT COUNT(*) as cnt FROM personal_memories WHERE team_id = %s",
                (team_id,), fetch=True,
            )
            stats["personal_memories"] = row[0]["cnt"] if row else 0
            row = self._query(
                "SELECT COUNT(*) as cnt FROM team_memories WHERE team_id = %s",
                (team_id,), fetch=True,
            )
            stats["team_memories"] = row[0]["cnt"] if row else 0
            row = self._query(
                "SELECT COUNT(*) as cnt FROM agents WHERE team_id = %s",
                (team_id,), fetch=True,
            )
            stats["agents"] = row[0]["cnt"] if row else 0
        else:
            row = self._query("SELECT COUNT(*) as cnt FROM personal_memories", fetch=True)
            stats["personal_memories"] = row[0]["cnt"] if row else 0
            row = self._query("SELECT COUNT(*) as cnt FROM team_memories", fetch=True)
            stats["team_memories"] = row[0]["cnt"] if row else 0
            row = self._query("SELECT COUNT(*) as cnt FROM agents", fetch=True)
            stats["agents"] = row[0]["cnt"] if row else 0
            row = self._query("SELECT COUNT(*) as cnt FROM teams", fetch=True)
            stats["teams"] = row[0]["cnt"] if row else 0
        return stats

    # -- Pipeline config --

    def get_pipeline_config(self, team_id: str):
        rows = self._query(
            "SELECT * FROM pipeline_config WHERE team_id = %s",
            (team_id,), fetch=True,
        )
        return rows[0] if rows else None

    def create_pipeline_config(self, team_id: str, compress_every_n: int = 0,
                               compress_target_count: int = 5,
                               cleanup_idle_days: int = 0,
                               cleanup_min_importance: float = 0.2,
                               enabled: bool = True,
                               warmup_max_memories: int = 0,
                               warmup_compress_every_n: int = 1):
        self._query(
            "INSERT INTO pipeline_config (team_id, compress_every_n, compress_target_count, cleanup_idle_days, cleanup_min_importance, enabled, warmup_max_memories, warmup_compress_every_n) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)",
            (team_id, compress_every_n, compress_target_count, cleanup_idle_days, cleanup_min_importance, enabled, warmup_max_memories, warmup_compress_every_n),
        )

    def update_pipeline_config(self, team_id: str, **kwargs):
        sets, args = [], []
        for field in ["compress_every_n", "compress_target_count", "cleanup_idle_days", "cleanup_min_importance", "enabled", "warmup_max_memories", "warmup_compress_every_n"]:
            if field in kwargs and kwargs[field] is not None:
                sets.append(f"{field} = %s")
                args.append(kwargs[field])
        if not sets:
            return
        args.append(team_id)
        self._query(f"UPDATE pipeline_config SET {', '.join(sets)} WHERE team_id = %s", tuple(args))

    def get_all_pipeline_configs(self):
        return self._query("SELECT * FROM pipeline_config WHERE enabled = 1", fetch=True) or []

    # -- Scenario management --

    def add_scenario(self, team_id: str, agent_id: str, name: str,
                     summary: str, memory_ids: list) -> str:
        import uuid
        sid = str(uuid.uuid4())[:16]
        self._query(
            "INSERT INTO memory_scenarios (id, team_id, agent_id, name, summary, memory_ids) VALUES (%s, %s, %s, %s, %s, %s)",
            (sid, team_id, agent_id, name, summary, json.dumps(memory_ids)),
        )
        return sid

    def get_scenario(self, scenario_id: str):
        rows = self._query("SELECT * FROM memory_scenarios WHERE id = %s", (scenario_id,), fetch=True)
        return rows[0] if rows else None

    def get_scenarios_by_agent(self, agent_id: str):
        rows = self._query("SELECT * FROM memory_scenarios WHERE agent_id = %s ORDER BY created_at DESC", (agent_id,), fetch=True)
        for row in rows:
            if isinstance(row.get("memory_ids"), str):
                row["memory_ids"] = json.loads(row["memory_ids"])
        return rows or []

    def delete_scenario(self, scenario_id: str) -> bool:
        self._query("DELETE FROM memory_scenarios WHERE id = %s", (scenario_id,))
        return True

    # -- Persona management --

    def add_persona(self, team_id, agent_id, preferences, habits, expertise, communication_style, summary):
        import uuid
        pid = str(uuid.uuid4())[:16]
        self._query(
            'INSERT INTO user_personas (id, team_id, agent_id, preferences, habits, expertise, communication_style, summary) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)',
            (pid, team_id, agent_id, json.dumps(preferences), json.dumps(habits), json.dumps(expertise), communication_style, summary),
        )
        return pid

    def get_persona(self, persona_id):
        rows = self._query('SELECT * FROM user_personas WHERE id = %s', (persona_id,), fetch=True)
        if rows:
            row = rows[0]
            for field in ['preferences', 'habits', 'expertise']:
                if isinstance(row.get(field), str):
                    row[field] = json.loads(row[field])
            return row
        return None

    def get_persona_by_agent(self, agent_id):
        rows = self._query('SELECT * FROM user_personas WHERE agent_id = %s ORDER BY created_at DESC LIMIT 1', (agent_id,), fetch=True)
        if rows:
            row = rows[0]
            for field in ['preferences', 'habits', 'expertise']:
                if isinstance(row.get(field), str):
                    row[field] = json.loads(row[field])
            return row
        return None

    def delete_persona(self, persona_id):
        self._query('DELETE FROM user_personas WHERE id = %s', (persona_id,))
        return True