import os
from groq import Groq
from typing import Dict
from .schemas import AgentReport, JudgeVerdict

client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def evaluate_consensus(reports: Dict[str, AgentReport]) -> bool:
    claims = [r.claim for r in reports.values()]
    confidences = [r.confidence for r in reports.values()]
    
    all_same = len(set(claims)) == 1
    high_conf = all(c >= 0.85 for c in confidences)
    
    return all_same and high_conf

def run_judge(reports_summary: str, raw_data: str) -> JudgeVerdict:
    system_prompt = (
        "You are the presiding Judge Magistrate in an Enterprise Cybersecurity Matrix.\n"
        "Your subordinate agents have provided split panel reports.\n"
        "Respond strictly in valid JSON matching this schema:\n"
        "{\"final_verdict\": \"CRITICAL_THREAT\"|\"SAFE\"|\"MANUAL_REVIEW_REQUIRED\", \"risk_score\": float, \"verdict_justification\": string}\n\n"
        "DECISION ARCHITECTURE RULES:\n"
        "1. If ANY analyst engine raises a 'PHISHING' claim with confidence >= 0.85, force 'CRITICAL_THREAT'.\n"
        "2. Avoid label contradictions: Never set final_verdict to SAFE if your textual argument implies a threat pattern.\n"
        "3. SUPREME VETO RULE: If the Target URL belongs to a verified, universally trusted root domain (e.g., 'github.com', 'google.com', 'microsoft.com') and shows no signs of typosquatting, you MUST overrule paranoid sub-agents, ignore Rule 1, and force a 'SAFE' verdict with a low risk score. Standard authentication portals on official domains are inherently safe."
    )
    
    # We leverage the powerful 70B cloud model for premium reasoning accuracy
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"RAW INPUTS:\n{raw_data}\n\nREPORTS SUMMARY:\n{reports_summary}"}
        ],
        response_format={"type": "json_object"},
        temperature=0.0
    )
    return JudgeVerdict.model_validate_json(response.choices[0].message.content)