# Devon 2.0 - AI Memory & Model Fusion Architecture

from fastapi import FastAPI, HTTPException
import redis
import psycopg2
import boto3
from azure.storage.blob import BlobServiceClient
import dropbox
import json
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# FastAPI App
app = FastAPI()

# Redis Configuration (Short-Term Memory)
redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)

# PostgreSQL Configuration (Mid-Term Legal Memory)
pg_conn = psycopg2.connect(
    dbname='devon_db',
    user='devon_user',
    password='securepassword',
    host='localhost',
    port='5432'
)
pg_cursor = pg_conn.cursor()

# AWS S3 Configuration (Strategic Knowledge Base)
s3_client = boto3.client('s3', aws_access_key_id=os.getenv('AWS_ACCESS_KEY'), aws_secret_access_key=os.getenv('AWS_SECRET_KEY'))
S3_BUCKET = "devon-legal-memory"

# Azure Blob Storage Configuration
azure_blob_service = BlobServiceClient.from_connection_string(os.getenv("AZURE_STORAGE_CONNECTION_STRING"))
AZURE_CONTAINER = "devon-legal-data"

# Dropbox Configuration (Long-Term Archival Memory)
db_client = dropbox.Dropbox(os.getenv("DROPBOX_ACCESS_TOKEN"))

# AI Model Paths
NEOX_MODEL_PATH = "/opt/models/gpt-neox-20B"
MISTRAL_MODEL_PATH = "/opt/models/Mistral-7B-v0.1"

# Load AI Models (Fusion)
try:
    tokenizer = AutoTokenizer.from_pretrained(NEOX_MODEL_PATH)
    neox_model = AutoModelForCausalLM.from_pretrained(NEOX_MODEL_PATH, torch_dtype=torch.float16, local_files_only=True)
    mistral_model = AutoModelForCausalLM.from_pretrained(MISTRAL_MODEL_PATH, torch_dtype=torch.float16, local_files_only=True)
    print("✅ Models Loaded Successfully - NeoX & Mistral Fusion")
except Exception as e:
    print(f"🚨 Error loading models: {e}")

# Route: Store in Redis (Short-Term Case Memory)
@app.post("/memory/short-term")
def store_short_term_memory(case_id: str, data: dict):
    redis_client.set(case_id, json.dumps(data))
    return {"status": "Stored in short-term memory"}

# Route: Store in PostgreSQL (Mid-Term Legal Memory)
@app.post("/memory/mid-term")
def store_mid_term_memory(case_id: str, data: dict):
    pg_cursor.execute("INSERT INTO case_memory (case_id, data) VALUES (%s, %s) ON CONFLICT (case_id) DO UPDATE SET data = %s", (case_id, json.dumps(data), json.dumps(data)))
    pg_conn.commit()
    return {"status": "Stored in mid-term memory"}

# Route: Store in AWS S3 (Strategic Legal Knowledge Base)
@app.post("/memory/strategic")
def store_strategic_memory(case_id: str, data: dict):
    s3_client.put_object(Bucket=S3_BUCKET, Key=f"cases/{case_id}.json", Body=json.dumps(data))
    return {"status": "Stored in strategic memory"}

# Route: Retrieve Case Data (Unified Fetch System)
@app.get("/memory/{case_id}")
def get_case_memory(case_id: str):
    case_data = redis_client.get(case_id)
    if case_data:
        return {"source": "short-term", "data": json.loads(case_data)}
    
    pg_cursor.execute("SELECT data FROM case_memory WHERE case_id = %s", (case_id,))
    result = pg_cursor.fetchone()
    if result:
        return {"source": "mid-term", "data": json.loads(result[0])}
    
    try:
        obj = s3_client.get_object(Bucket=S3_BUCKET, Key=f"cases/{case_id}.json")
        return {"source": "strategic", "data": json.loads(obj['Body'].read())}
    except:
        pass
    
    return {"error": "Case data not found"}

# Route: AI Legal Analysis (Fusion Processing)
@app.post("/legal/analyze")
def analyze_legal_text(data: dict):
    input_text = data.get("text", "")
    if not input_text:
        raise HTTPException(status_code=400, detail="No text provided")
    
    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
    input_ids = inputs["input_ids"].to("cuda" if torch.cuda.is_available() else "cpu")
    attention_mask = inputs["attention_mask"].to("cuda" if torch.cuda.is_available() else "cpu")

    neox_output = neox_model.generate(input_ids, attention_mask=attention_mask, max_length=512)
    mistral_output = mistral_model.generate(input_ids, attention_mask=attention_mask, max_length=512)

    response_text_neox = tokenizer.decode(neox_output[0], skip_special_tokens=True)
    response_text_mistral = tokenizer.decode(mistral_output[0], skip_special_tokens=True)
    
    fusion_response = f"NeoX Analysis: {response_text_neox}\nMistral Analysis: {response_text_mistral}"
    
    return {"input": input_text, "analysis": fusion_response}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)