import os
import json
from dataclasses import dataclass
from typing import List, Dict, Optional
from groq import Groq

client = Groq(api_key=os.getenv("GROQ_API_KEY"))

def call_llm(system_prompt: str, user_prompt: str, json_mode: bool = False):
    model_name = "llama-3.3-70b-versatile"
    response_format = {"type": "json_object"} if json_mode else None
    completion = client.chat.completions.create(
        model=model_name,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        response_format=response_format,
        max_tokens=700,
        temperature=0.4,
    )
    return completion.choices[0].message.content

@dataclass
class Memory:
    user_preferences: List[str]
    emotional_patterns: List[str]
    facts: List[str]

MEMORY_SYSTEM = """
Extract:
1. user_preferences
2. emotional_patterns
3. facts
Return JSON only.
"""

def history_to_text(history: List[Dict]):
    return "\n".join([f"{m['role'].upper()}: {m['content']}" for m in history])

def extract_memory(history: List[Dict]):
    raw = call_llm(MEMORY_SYSTEM, history_to_text(history), json_mode=True)
    try:
        data = json.loads(raw)
    except:
        data = json.loads(raw[raw.find("{"): raw.rfind("}")+1])
    return Memory(
        user_preferences=data.get("user_preferences", []),
        emotional_patterns=data.get("emotional_patterns", []),
        facts=data.get("facts", []),
    )

PERSONALITY_STYLES = {
    "neutral": "neutral helpful tone",
    "calm_mentor": "calm mentor tone",
    "witty_friend": "witty friendly tone",
    "therapist_style": "gentle therapist tone"
}

class PersonalityEngine:
    def __init__(self, style, memory=None):
        self.style = style
        self.desc = PERSONALITY_STYLES[style]
        self.memory = memory

    def transform(self, user_message, base_reply):
        rules = "Tone only. No new content. Meaning unchanged."
        system_prompt = f"Rewrite in {self.desc}. {rules}"
        return call_llm(system_prompt, base_reply)