#!/usr/bin/env python3
import argparse
import json
import logging
import os
import sqlite3
import sys
from datetime import datetime, timezone

logging.basicConfig(level=logging.INFO, format='%(name)s: [%(levelname)s] %(message)s')
logger = logging.getLogger("memory_doctor")

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(BASE_DIR, "bin"))
from m3_sdk import add_database_arg, resolve_db_path


def fix_missing_timestamps(conn):
    """Ensures all items have at least a created_at timestamp."""
    logger.info("Checking for missing timestamps...")
    now = datetime.now(timezone.utc).isoformat() + "Z"

    # Fix items with NULL created_at
    res = conn.execute(
        "UPDATE memory_items SET created_at = ? WHERE created_at IS NULL",
        (now,)
    )
    if res.rowcount:
        logger.info(f"Fixed {res.rowcount} items with missing created_at.")

def validate_relationships(conn):
    """Prunes relationships pointing to non-existent items."""
    logger.info("Validating relationship integrity...")
    res = conn.execute("""
        DELETE FROM memory_relationships
        WHERE from_id NOT IN (SELECT id FROM memory_items)
           OR to_id NOT IN (SELECT id FROM memory_items)
    """)
    if res.rowcount:
        logger.info(f"Pruned {res.rowcount} orphaned relationships.")

def fix_metadata_json(conn):
    """Ensures metadata_json is valid JSON."""
    logger.info("Validating metadata JSON strings...")
    cursor = conn.execute("SELECT id, metadata_json FROM memory_items WHERE metadata_json IS NOT NULL")
    rows = cursor.fetchall()
    fixed = 0
    for rid, meta in rows:
        try:
            if meta:
                json.loads(meta)
        except json.JSONDecodeError:
            logger.warning(f"Repairing invalid JSON for item {rid}")
            conn.execute("UPDATE memory_items SET metadata_json = '{}' WHERE id = ?", (rid,))
            fixed += 1
    if fixed:
        logger.info(f"Repaired {fixed} items with invalid metadata JSON.")

def check_sovereign_embedder():
    """Checks the status of the integrated sovereign embedder."""
    logger.info("Checking Sovereign Embedder status...")
    import asyncio

    # Import locally to avoid circular dependencies if any
    import memory_core

    status = asyncio.run(memory_core.embedder_status_impl())

    if status["binary_found"]:
        logger.info("✓ Sovereign binary found.")
    else:
        logger.info("- Sovereign binary not installed (optional).")

    if status["status"] == "online":
        logger.info(f"✓ Sovereign server is ONLINE on port {status['port']}.")
        if status["models"]:
            model_ids = [m["id"] for m in status["models"]]
            logger.info(f"  Loaded models: {', '.join(model_ids)}")
    elif status["binary_found"]:
        stat_str = status["status"].upper()
        logger.warning(f"⚠ Sovereign binary exists but server is {stat_str} on port {status['port']}.")
        logger.info("  Tip: Run `mcp-memory embedder start` to boot it.")


def main():
    parser = argparse.ArgumentParser(description="Memory health check and repair.")
    add_database_arg(parser)
    args = parser.parse_args()

    db_path = resolve_db_path(args.database)
    if not os.path.exists(db_path):
        logger.error(f"Database not found at {db_path}")
        sys.exit(1)

    conn = sqlite3.connect(db_path)
    try:
        check_sovereign_embedder()
        fix_missing_timestamps(conn)
        validate_relationships(conn)
        fix_metadata_json(conn)
        conn.commit()
        logger.info("Memory health check and repair completed.")
    except Exception as e:
        logger.error(f"Doctor failed: {e}")
        conn.rollback()
    finally:
        conn.close()

if __name__ == "__main__":
    main()
