import concurrent.futures
import sys
import uuid
from contextlib import contextmanager
from dotenv import load_dotenv

from crewai import Agent, Crew, LLM, Process, Task
from crewai_tools import SerperDevTool

load_dotenv()

# ------ Session tracking (no-op fallback when OTel pkgs not installed) ------

try:
    from openinference.instrumentation import using_session as _using_session
except ImportError:
    @contextmanager
    def _using_session(session_id):  # noqa: F811
        yield

try:
    from openinference.semconv.trace import SpanAttributes
except ImportError:
    class _SpanAttributes:
        SESSION_ID = "session_id"
    SpanAttributes = _SpanAttributes()  # type: ignore

# ------ Phoenix / OpenTelemetry tracing ------

try:
    from phoenix.otel import register
    from huggingface_hub.utils import build_hf_headers
    from openinference.instrumentation.openai import OpenAIInstrumentor

    from opentelemetry import trace
    from openinference.instrumentation.crewai import CrewAIInstrumentor

    OpenAIInstrumentor().uninstrument()
    tp = trace.get_tracer_provider()
    if tp and hasattr(tp, "shutdown"):
        tp.shutdown()
    trace._TRACER_PROVIDER = None

    tracer_provider = register(
        project_name="crewai",
        endpoint="https://RCaz-phoenix-arize-observability.hf.space/v1/traces",
        headers=build_hf_headers(),
    )
    CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
except ImportError:
    pass


###### The agentic app

# ------ LLM endpoint constants ------
_VLLM_BASE_URL = "https://rcaz33--example-vllm-inference-serve.modal.run/v1"
_VLLM_MODEL = "openai/google/gemma-4-26B-A4B-it"  # LiteLLM prefix — for crewai
_VLLM_SERVED_MODEL = (
    "google/gemma-4-26B-A4B-it"  # actual vLLM served name — for direct API calls
)

# Define our LLM using the Modal-deployed Gemma 4 26B model via vLLM (OpenAI-compatible API)
llm = LLM(
    model=_VLLM_MODEL,
    base_url=_VLLM_BASE_URL,
    api_key="sk-dummy-key-not-needed",
    max_tokens=4096,
)

search_tool = SerperDevTool()

# ------ Modal service handles (lazy) ------
_flux_url = None
_vox_url = None
_transcribe_url = None


def _get_flux_url() -> str:
    global _flux_url
    if _flux_url is None:
        import modal as _modal

        _flux_url = _modal.Cls.from_name(
            "flux-image-generator", "FluxGenerator"
        )().generate.get_web_url()
    return _flux_url


def _get_vox_url() -> str:
    global _vox_url
    if _vox_url is None:
        import modal as _modal

        _vox_url = _modal.Cls.from_name(
            "voxcpm-generator", "VoxCPMGenerator"
        )().synthesize.get_web_url()
    return _vox_url


def _get_transcribe_url() -> str:
    global _transcribe_url
    if _transcribe_url is None:
        import modal as _modal

        _transcribe_url = _modal.Cls.from_name(
            "cohere-transcriber", "CohereTranscriber"
        )().transcribe.get_web_url()
    return _transcribe_url


CREW_TIMEOUT = 300  # seconds per crew


def _run_with_timeout(fn, *, timeout=CREW_TIMEOUT):
    """Run a callable with a hard timeout. Returns result or fallback str."""
    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
        fut = pool.submit(fn)
        try:
            return fut.result(timeout=timeout)
        except concurrent.futures.TimeoutError:
            print(f"\n⚠ Crew timed out after {timeout}s — using fallback\n")
            fut.cancel()
            return "The search timed out. Using general knowledge as fallback."


# ------ Research Pipeline ------


def run_pipeline(statement: str, session_id: str | None = None) -> dict:
    """
    Run the 3-crew research pipeline on a statement.

    Args:
        statement: The statement to research.
        session_id: Optional Phoenix session ID for trace grouping.

    Returns a dict with keys:
      prompt, result_corroborate, result_opposite
    """
    if session_id is None:
        session_id = str(uuid.uuid4())

    print(f"\n{'=' * 60}")
    print(f'Researching statement: "{statement}"')
    print(f"Session: {session_id}")
    print(f"{'=' * 60}\n")

    with _using_session(session_id):
        # -------------------------------------------------------
        # CREW 1 — Corroborative research
        # -------------------------------------------------------
        researcher_corroborate = Agent(
            role="Corroboration Researcher",
            goal=f"Find facts and evidence that support or corroborate the statement: '{statement}'",
            backstory="You are a researcher skilled at finding supporting evidence for a given claim. You provide detailed findings with sources.",
            verbose=True,
            allow_delegation=False,
            tools=[search_tool],
            llm=llm,
            max_iter=1,
        )

        task_corroborate = Task(
            description=f"Search the web for facts, data, and expert opinions that support the statement: '{statement}'. Provide a detailed list of corroborative evidence with sources.",
            expected_output="Detailed bullet-point list of corroborative facts with sources",
            agent=researcher_corroborate,
        )

        crew_corroborate = Crew(
            agents=[researcher_corroborate],
            tasks=[task_corroborate],
            verbose=True,
            process=Process.sequential,
        )

        print("\n>>> Crew 1: Corroborative research <<<\n")
        result_corroborate = _run_with_timeout(crew_corroborate.kickoff)

        # -------------------------------------------------------
        # CREW 2 — Opposite / complementary research
        # -------------------------------------------------------
        researcher_opposite = Agent(
            role="Opposition Researcher",
            goal=f"Find facts and evidence that challenge, contradict, or offer a complementary perspective to the statement: '{statement}'",
            backstory="You are a researcher skilled at finding counterarguments, alternative viewpoints, and complementary angles to a given claim.",
            verbose=True,
            allow_delegation=False,
            tools=[search_tool],
            llm=llm,
            max_iter=1,
        )

        task_opposite = Task(
            description=f"Search the web for facts, data, and expert opinions that contradict, challenge, or offer a different perspective on the statement: '{statement}'. Provide a detailed list of opposing or complementary evidence with sources.",
            expected_output="Detailed bullet-point list of opposing/complementary facts with sources",
            agent=researcher_opposite,
        )

        crew_opposite = Crew(
            agents=[researcher_opposite],
            tasks=[task_opposite],
            verbose=True,
            process=Process.sequential,
        )

        print("\n>>> Crew 2: Opposite / complementary research <<<\n")
        result_opposite = _run_with_timeout(crew_opposite.kickoff)

        # -------------------------------------------------------
        # CREW 3 — Synthesize into image prompt
        # -------------------------------------------------------
        synthesizer = Agent(
            role="Creative Director & Visual Designer",
            goal="Synthesize two research perspectives into a compelling visual concept",
            backstory="You are a world-class creative director who translates complex, contrasting ideas into powerful visual concepts.",
            verbose=True,
            allow_delegation=False,
            llm=llm,
            max_iter=1,
        )

        task_synthesize = Task(
            description=f"""A user made this statement: "{statement}"

Two research crews investigated it. Here are their findings:

=== CORROBORATIVE EVIDENCE ===
{result_corroborate}

=== OPPOSING / COMPLEMENTARY EVIDENCE ===
{result_opposite}

Your job:
Create a highly detailed visual scene description that captures the dialogue, contrast, or tension between these two perspectives. Describe the composition, colors, lighting, mood, subjects, setting, and visual metaphor in vivid detail — as if instructing an artist or image generation model.

Output exactly in this format:

IMAGE PROMPT: <detailed visual description, 3-5 sentences>""",
            expected_output="A detailed image prompt",
            agent=synthesizer,
        )

        crew_synthesize = Crew(
            agents=[synthesizer],
            tasks=[task_synthesize],
            verbose=True,
            process=Process.sequential,
        )

        print("\n>>> Crew 3: Generating image prompt <<<\n")
        result = crew_synthesize.kickoff()

        print(f"\n{'=' * 60}")
        print("IMAGE PROMPT")
        print(f"{'=' * 60}\n")
        print(result)

        # Extract IMAGE PROMPT
        result_text = str(result)
        prompt = None
        for line in result_text.split("\n"):
            if line.startswith("IMAGE PROMPT:"):
                prompt = line[len("IMAGE PROMPT:") :].strip()

        return {
            "prompt": prompt,
            "result_corroborate": str(result_corroborate),
            "result_opposite": str(result_opposite),
        }


# ------ Image Generation (Flux on Modal) ------


def generate_image(prompt: str) -> bytes:
    print(f"{'=' * 60}")
    print("Generating image with Flux on Modal...")
    print(f"{'=' * 60}\n")
    print(f"Prompt: {prompt}\n")

    import httpx as _httpx

    url = _get_flux_url()
    resp = _httpx.post(
        url,
        json={"prompt": prompt, "steps": 30},
        timeout=600,
        follow_redirects=True,
    )
    resp.raise_for_status()
    return resp.content


# ------ Caption Generation (direct LLM call with image prompt context) ------


def generate_caption(corroborate: str, opposite: str, image_prompt: str, session_id: str | None = None) -> dict:
    """
    Send the research + image prompt to Gemma 4 and get a 30-second
    spoken caption + voice style back.
    """
    if session_id is None:
        session_id = str(uuid.uuid4())

    import httpx as _httpx

    print(f"\n{'=' * 60}")
    print("Generating caption via Gemma 4...")
    print(f"{'=' * 60}\n")

    payload = {
        "model": _VLLM_SERVED_MODEL,
        "max_tokens": 1024,
        "messages": [
            {
                "role": "system",
                "content": "You are a creative narrator. Given two research perspectives and the description of an image they inspired, write a 30-second spoken narration (60-90 words) and choose a matching voice style. The caption should tie the research and the visual together.",
            },
            {
                "role": "user",
                "content": f"""Two research crews investigated this statement. Here are their findings:

=== CORROBORATIVE EVIDENCE ===
{corroborate}

=== OPPOSING / COMPLEMENTARY EVIDENCE ===
{opposite}

The image below was generated from this prompt (based on the findings above):
"{image_prompt}"

Your job:
1. Write a 30-second spoken narration (60-90 words) that weaves the two research perspectives together and describes what the image shows.
2. Choose a voice style for the narration (e.g. "gentle melancholic girl", "laid-back surfer dude", "authoritative news anchor", "warm thoughtful professor", etc.).

Output exactly in this format:

CAPTION: <30-second spoken narration, 60-90 words>
VOICE_STYLE: <voice style description>""",
            },
        ],
    }

    with _using_session(session_id):
        resp = _httpx.post(
            f"{_VLLM_BASE_URL}/chat/completions",
            json=payload,
            headers={"Authorization": "Bearer sk-dummy-key-not-needed"},
            timeout=300,
        )
        resp.raise_for_status()
        body = resp.json()
        text = body["choices"][0]["message"]["content"]

    print(f"Gemma response:\n{text}\n")

    caption = None
    voice_style = None
    for line in text.split("\n"):
        if line.startswith("CAPTION:"):
            caption = line[len("CAPTION:") :].strip()
        elif line.startswith("VOICE_STYLE:"):
            voice_style = line[len("VOICE_STYLE:") :].strip().lower()

    return {"caption": caption, "voice_style": voice_style}


# ------ Voice Generation (VoxCPM on Modal) ------


def generate_voice(voice_style: str, voice_script: str, session_id: str | None = None) -> bytes:
    """Send script to VoxCPM on Modal T4, return WAV bytes."""
    if session_id is None:
        session_id = str(uuid.uuid4())

    print(f"\n{'=' * 60}")
    print("Generating voice with VoxCPM on Modal...")
    print(f"{'=' * 60}\n")

    import httpx as _httpx

    with _using_session(session_id):
        url = _get_vox_url()
        resp = _httpx.post(
            url,
            json={"text": voice_script, "voice_style": voice_style},
            timeout=600,
            follow_redirects=True,
        )
        resp.raise_for_status()
        return resp.content


# ------ Audio Transcription (Cohere Transcribe on Modal) ------


def transcribe_audio(audio_path: str, session_id: str | None = None) -> str:
    import base64
    import httpx as _httpx

    if session_id is None:
        session_id = str(uuid.uuid4())

    url = _get_transcribe_url()

    with _using_session(session_id):
        with open(audio_path, "rb") as f:
            audio_b64 = base64.b64encode(f.read()).decode()

        resp = _httpx.post(url, json={"audio": audio_b64}, timeout=600)
        resp.raise_for_status()
        text = resp.json()["transcription"]
        print(f'Transcribed: "{text}"\n')
        return text


# ------ CLI Entry Point ------

if __name__ == "__main__":
    audio_path = None
    statement = None
    session_id = str(uuid.uuid4())

    if len(sys.argv) > 1:
        if sys.argv[1] == "--audio" and len(sys.argv) > 2:
            audio_path = sys.argv[2]
        else:
            statement = " ".join(sys.argv[1:])

    if audio_path:
        statement = transcribe_audio(audio_path, session_id=session_id)

    if not statement:
        statement = input("Enter a statement to research: ")

    with _using_session(session_id):
        results = run_pipeline(statement, session_id=session_id)
        prompt = results["prompt"]
        corroborate = results["result_corroborate"]
        opposite = results["result_opposite"]

        # Generate image
        if prompt:
            try:
                image_bytes = generate_image(prompt)
                filename = "crew_flux_output.png"
                with open(filename, "wb") as f:
                    f.write(image_bytes)
                print(f"\n✓ Image saved to {filename}")
            except Exception as e:
                print(f"\n✗ Failed to generate image: {e}")
                print(
                    "  Make sure flux_generator.py is deployed: modal deploy flux_generator.py"
                )
        else:
            print("(Skipping image generation — no IMAGE PROMPT found in output)")
            image_bytes = None

        # Generate caption from research + image prompt
        caption_result = None
        if prompt and corroborate and opposite:
            try:
                caption_result = generate_caption(corroborate, opposite, prompt, session_id=session_id)
                if caption_result["caption"]:
                    print(f"\n✓ Caption: {caption_result['caption']}")
                if caption_result["voice_style"]:
                    print(f"\n✓ Voice style: {caption_result['voice_style']}")
            except Exception as e:
                print(f"\n✗ Failed to generate caption: {e}")
        else:
            print("(Skipping caption generation — missing prompt or research)")

        # Generate voice
        if caption_result and caption_result["voice_style"] and caption_result["caption"]:
            try:
                audio_bytes = generate_voice(
                    caption_result["voice_style"], caption_result["caption"], session_id=session_id
                )
                filename = "crew_voice_output.wav"
                with open(filename, "wb") as f:
                    f.write(audio_bytes)
                print(f"\n✓ Voice saved to {filename}")
            except Exception as e:
                print(f"\n✗ Failed to generate voice: {e}")
                print(
                    "  Make sure voxcpm_generator.py is deployed: modal deploy voxcpm_generator.py"
                )
        else:
            print("\n(Skipping voice generation — no caption/style found)")