import concurrent.futures
import contextvars
import sys
import uuid
from contextlib import contextmanager, nullcontext
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.litellm import LiteLLMInstrumentor

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

    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(),
    )
    LiteLLMInstrumentor().instrument(tracer_provider=tracer_provider)
    CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
    tracer = trace.get_tracer("crewai")
except ImportError:
    tracer = None


###### 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
_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=1)


def _run_with_timeout(fn, *, timeout=CREW_TIMEOUT):
    """Run a callable with a hard timeout. Returns result or fallback str.

    Propagates OTel context (session ID, etc.) to the worker thread.
    """
    ctx = contextvars.copy_context()
    fut = _POOL.submit(ctx.run, fn)
    try:
        return fut.result(timeout=timeout)
    except concurrent.futures.TimeoutError:
        print(f"\n⚠ Crew timed out after {timeout}s — using fallback\n")
        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 and merges 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

    with (
        tracer.start_as_current_span("flux.generate_image") if tracer else nullcontext()
    ):
        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):
        _span_cm = (
            tracer.start_as_current_span("llm.generate_caption")
            if tracer
            else nullcontext()
        )
        with _span_cm as _span:
            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"]

            if tracer and "usage" in body:
                usage = body["usage"]
                _span.set_attribute(
                    "llm.token_count.prompt", usage.get("prompt_tokens", 0)
                )
                _span.set_attribute(
                    "llm.token_count.completion", usage.get("completion_tokens", 0)
                )
                _span.set_attribute(
                    "llm.token_count.total", usage.get("total_tokens", 0)
                )
                _span.set_attribute("llm.model_name", _VLLM_SERVED_MODEL)

    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):
        with (
            tracer.start_as_current_span("voxcpm.generate_voice")
            if tracer
            else nullcontext()
        ):
            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 (
            tracer.start_as_current_span("cohere.transcribe_audio")
            if tracer
            else nullcontext()
        ):
            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)")