import json
import numpy as np
from openai import OpenAI
import sys

# --- Configuration Loading ---
def load_config(path='config.json'):
    with open(path, 'r') as f:
        return json.load(f)

# --- Vector Store (Mock Implementation) ---
class VectorStore:
    def __init__(self, file_path):
        print(f"Loading embeddings from {file_path}...")
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)
            self.chunks = [item['text'] for item in data]
            # Convert lists back to numpy arrays
            self.embeddings = [np.array(item['embedding']) for item in data]
            print(f"Loaded {len(self.chunks)} chunks.")
        except FileNotFoundError:
            print(f"Error: {file_path} not found. Creating a dummy knowledge base.")
            # Fallback dummy data for testing if file missing
            self.chunks = ["DBT teaches Distress Tolerance skills.", "DEAR MAN is a skill for interpersonal effectiveness."]
            # Dummy embeddings (normally these come from an embedding model)
            self.embeddings = [np.random.rand(1536), np.random.rand(1536)]

    def search(self, query_embedding, top_k=3):
        """Finds most similar chunks using Cosine Similarity."""
        if not self.embeddings:
            return []
        
        query_vec = np.array(query_embedding)
        scores = []
        for i, doc_vec in enumerate(self.embeddings):
            # Cosine similarity
            score = np.dot(query_vec, doc_vec) / (np.linalg.norm(query_vec) * np.linalg.norm(doc_vec))
            scores.append((score, i))
        
        scores.sort(reverse=True)
        return [self.chunks[i] for score, i in scores[:top_k]]

# --- OpenRouter Client ---
class LLMClient:
    def __init__(self, api_key):
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )

    def get_embedding(self, text):
        """Generates embedding for the query (using OpenAI as default for simplicity)."""
        # Note: OpenRouter supports embedding models, usually openai/text-embedding-3-small
        response = self.client.embeddings.create(
            model="openai/text-embedding-3-small",
            input=text
        )
        return response.data[0].embedding

    def generate(self, model_id, system_prompt, user_prompt, temperature=0.7):
        """Generic generation function."""
        try:
            response = self.client.chat.completions.create(
                model=model_id,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=temperature
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"Error calling {model_id}: {str(e)}"

# --- Jury Logic ---
class DBTJurySystem:
    def __init__(self, config):
        self.config = config
        self.llm = LLMClient(config['openrouter_api_key'])
        self.vector_store = VectorStore(config['embedding_file'])

    def retrieve_context(self, query):
        print("\n[1. Retrieving Context...]")
        query_emb = self.llm.get_embedding(query)
        context_chunks = self.vector_store.search(query_emb)
        return "\n".join(context_chunks)

    def run_generator(self, query, context):
        """Step 1: Generate the initial draft."""
        print("[2. Generating Draft...]")
        prompt = f"""
        Context from DBT Manual:
        {context}

        User Query: {query}

        Instructions: Answer the user's query using ONLY the context provided. 
        If the context is insufficient, state that clearly.
        """
        return self.llm.generate(
            self.config['models']['generator'],
            self.config['system_prompt'],
            prompt
        )

    def run_jury_deliberation(self, query, context, draft_answer):
        """Step 2: The Jury Votes."""
        print("[3. Jury Deliberating (Single Veto Logic)...]")
        
        jury_config = [
            {"role": "Clinical Accuracy", "model_key": "jury_clinical", "instruction": "Check if the advice strictly follows DBT protocol."},
            {"role": "Safety", "model_key": "jury_safety", "instruction": "Check for any harmful, unethical, or dangerous advice. Be strict."},
            {"role": "Empathy", "model_key": "jury_empathy", "instruction": "Check if the tone is supportive and non-judgmental."},
            {"role": "Hallucination", "model_key": "jury_hallucination", "instruction": "Verify every claim in the answer is supported by the context."}
        ]

        results = []

        for member in jury_config:
            print(f"   - Querying {member['role']} ({member['model_key']})...")
            
            judge_prompt = f"""
            You are a Jury Member: {member['role']}.
            Your specific instruction: {member['instruction']}

            --- SOURCE CONTEXT ---
            {context}
            ----------------------

            --- PROPOSED ANSWER ---
            {draft_answer}
            ----------------------

            Task: Analyze the proposed answer.
            1. Does it violate your instruction?
            2. Is it factually grounded in the source context?

            Output format strictly as JSON:
            {{ "verdict": "APPROVE" or "VETO", "reason": "Your reasoning..." }}
            """
            
            response_text = self.llm.generate(
                self.config['models'][member['model_key']],
                "You are a strict JSON validator. Output ONLY valid JSON.",
                judge_prompt,
                temperature=0.1 # Low temp for deterministic judging
            )

            # Parse JSON response
            try:
                # Clean up potential markdown code blocks
                clean_response = response_text.strip().replace("```json", "").replace("```", "")
                vote = json.loads(clean_response)
            except json.JSONDecodeError:
                # Fallback if model fails to output JSON
                vote = {"verdict": "APPROVE", "reason": "Failed to parse response, defaulting to Approve."}
            
            vote['member'] = member['role']
            results.append(vote)

            if vote['verdict'].upper() == 'VETO':
                print(f"   ❌ VETO by {member['role']}: {vote['reason']}")
                return False, results
        
        print("   ✅ Unanimous Approval.")
        return True, results

    def process_query(self, query):
        # 1. RAG Retrieval
        context = self.retrieve_context(query)
        
        # 2. Generation
        draft = self.run_generator(query, context)
        
        print(f"\n--- Draft Answer ---\n{draft}\n--------------------")

        # 3. Jury Voting
        approved, votes = self.run_jury_deliberation(query, context, draft)

        if approved:
            return {
                "status": "SUCCESS",
                "answer": draft,
                "votes": votes
            }
        else:
            return {
                "status": "REJECTED",
                "answer": "The Jury could not agree on a safe or accurate answer. Please consult a professional or try rephrasing.",
                "votes": votes
            }

# --- Main Execution ---
if __name__ == "__main__":
    config = load_config()
    
    # Interactive Loop
    print("\nDBT Quorum System Active (Type 'exit' to quit)")
    system = DBTJurySystem(config)
    
    while True:
        user_query = input("\nUser: ")
        if user_query.lower() in ['exit', 'quit']:
            break
            
        response = system.process_query(user_query)
        
        print("\n===== FINAL OUTPUT =====")
        print(f"Status: {response['status']}")
        print(f"Response: {response['answer']}")
        # Optionally print vote breakdown
        # print(f"Vote Details: {json.dumps(response['votes'], indent=2)}")
        print("========================")