import json
from sqlalchemy import create_engine, text
from src.core.config import Config
from datetime import datetime

class MemoryManager:
    def __init__(self):
        self.config = Config()
        try:
            self.engine = create_engine(self.config.DB_URL)
        except:
            self.engine = None

    def get_conn(self): return self.engine.connect()

    def save_attachment(self, user_id, workspace_id, filename, filetype, analysis):
        try:
            with self.get_conn() as conn:
                # Get or Create Session
                row = conn.execute(text("SELECT id FROM chat_sessions WHERE user_id = :uid ORDER BY last_active DESC LIMIT 1"), {"uid": str(user_id)}).fetchone()
                if row: 
                    sid = row[0]
                else:
                    res = conn.execute(text("INSERT INTO chat_sessions (user_id, workspace_id, title) VALUES (:uid, :wid, 'New Chat') RETURNING id"), {"uid": str(user_id), "wid": str(workspace_id)}).fetchone()
                    sid = res[0]

                conn.execute(text("INSERT INTO chat_attachments (session_id, file_name, file_type, analysis_summary) VALUES (:sid, :f, :t, :a)"), {"sid": sid, "f": filename, "t": filetype, "a": analysis})
                conn.execute(text("UPDATE chat_sessions SET last_active = CURRENT_TIMESTAMP WHERE id = :sid"), {"sid": sid})
                conn.commit()
        except Exception as e: print(f"DB Error: {e}")

    def add_message(self, user_id, workspace_id, role, content):
        try:
            with self.get_conn() as conn:
                # Simple lookup for demo
                sid_res = conn.execute(text("SELECT id FROM chat_sessions WHERE user_id = :uid ORDER BY last_active DESC LIMIT 1"), {"uid": str(user_id)}).fetchone()
                sid = sid_res[0] if sid_res else 1
                
                conn.execute(text("INSERT INTO chat_messages (session_id, role, content) VALUES (:sid, :role, :content)"), {"sid": sid, "role": role, "content": str(content)})
                conn.commit()
        except: pass

    def get_context_string(self, user_id, limit=6):
        try:
            with self.get_conn() as conn:
                # 1. Get Messages
                rows = conn.execute(text("SELECT m.role, m.content FROM chat_messages m JOIN chat_sessions s ON m.session_id = s.id WHERE s.user_id = :uid ORDER BY m.created_at DESC LIMIT :lim"), {"uid": str(user_id), "lim": limit}).fetchall()
                history = "\n".join([f"{r[0]}: {r[1]}" for r in reversed(rows)])
                
                # 2. Get Attachments (SMARTER LOGIC)
                # We fetch the *Latest* one separately to mark it as CURRENT
                att_rows = conn.execute(text("SELECT a.file_name, a.analysis_summary, a.created_at FROM chat_attachments a JOIN chat_sessions s ON a.session_id = s.id WHERE s.user_id = :uid ORDER BY a.created_at DESC LIMIT 3"), {"uid": str(user_id)}).fetchall()
                
                vision_context = ""
                if att_rows:
                    # The first one is the latest
                    latest = att_rows[0]
                    vision_context += f"\n[CURRENTLY LOOKING AT IMAGE]:\nFile: {latest[0]}\nDescription: {latest[1]}\n"
                    
                    # The rest are history
                    if len(att_rows) > 1:
                        vision_context += "\n[PREVIOUS IMAGES (Context only)]:\n"
                        for r in att_rows[1:]:
                            vision_context += f"- {r[0]}: {r[1]}\n"
                            
                return vision_context + "\n" + history
        except Exception as e: 
            return f"Error: {e}"

    def save_workflow(self, workspace_id, name, json_data):
        # (Same as before)
        with self.get_conn() as conn:
            conn.execute(text("INSERT INTO scenarios (workspace_id, name, description, steps, status, created_at) VALUES (:wid, :name, 'AI Generated', :steps, 'active', :time)"),
            {"wid": workspace_id, "name": name, "steps": json.dumps(json_data), "time": datetime.now().isoformat()})
            conn.commit()
            return 1
    
    # Required for server.py imports
    def get_user_workspaces(self, uid): return []