import json
from agents import GroqClient, SerpClient
from prompts import build_master_prompt


GENERATOR_MODEL = "llama-3.1-70b-versatile"
VERIFIER_MODEL = "gemma2-27b-it"
FORMATTER_MODEL = "mixtral-8x7b-32768"


class MultiAgentOrchestrator:
    def __init__(self, groq: GroqClient, serp: SerpClient):
        self.groq = groq
        self.serp = serp

    def run_pipeline(self, subject, stream, partA, partB, partC, syllabus_text, ref_qp_text):

        # 1) Real-time search
        serp_results = self.serp.search(f"{subject} latest 2024 2025")
        snippets = "\n".join([r.get("snippet", "") for r in serp_results[:3]])

        # 2) Build master prompt
        master_prompt = build_master_prompt(
            subject, stream, partA, partB, partC,
            syllabus_text, ref_qp_text, snippets
        )

        # 3) AGENT 1 — GENERATOR
        gen_output = self.groq.ask(
            [{"role": "system", "content": "Exam generator agent."},
             {"role": "user", "content": master_prompt}],
            model=GENERATOR_MODEL
        )

        # Extract JSON from generator
        try:
            json_start = gen_output.rfind("{")
            qp_json = json.loads(gen_output[json_start:])
        except:
            qp_json = {"raw": gen_output}

        # 4) AGENT 2 — VERIFIER
        ver_prompt = f"Check this QP JSON for errors and propose corrections:\n{json.dumps(qp_json)}"
        verifier_output = self.groq.ask(
            [{"role": "system", "content": "Verifier agent."},
             {"role": "user", "content": ver_prompt}],
            model=VERIFIER_MODEL
        )

        try:
            corrections = json.loads(verifier_output)
        except:
            corrections = {"notes": verifier_output}

        # 5) AGENT 3 — FORMATTER
        fmt_prompt = f"""
Apply corrections to QP JSON and produce:
1. final_qp
2. answers
3. obe
All inside ONE valid JSON object.

Original QP:
{json.dumps(qp_json)}

Corrections:
{json.dumps(corrections)}
        """

        final_output = self.groq.ask(
            [{"role": "system", "content": "Formatter agent."},
             {"role": "user", "content": fmt_prompt}],
            model=FORMATTER_MODEL
        )

        try:
            final_json = json.loads(final_output)
        except:
            final_json = {"raw": final_output}

        return {
            "generator_raw": gen_output,
            "qp_json": qp_json,
            "verifier": corrections,
            "final": final_json
        }