import asyncio
import json
import os
import sys
from pathlib import Path
from typing import TypedDict, Literal
from urllib.parse import urlparse, parse_qs

from dotenv import load_dotenv
from pydantic import BaseModel
from yt_dlp import YoutubeDL

from langchain_groq import ChatGroq

from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import InMemorySaver

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


# =========================================================
# 1. ENVIRONMENT
# =========================================================

load_dotenv()

BASE_DIR = Path(__file__).resolve().parent.parent


# =========================================================
# 2. SETTINGS
# =========================================================

# TEST MODE: 1 minute.
# Production mein 60 kar dena.
BREAK_THRESHOLD_MINUTES = 1

# YouTube open hone ke baad screen lock se pehle delay.
SCREEN_LOCK_DELAY_SECONDS = 15


# =========================================================
# 3. MCP SERVER CONFIG
# =========================================================

server_params = StdioServerParameters(
    command=sys.executable,
    args=["-m", "mcp_server.server"],
    cwd=str(BASE_DIR),
)


# =========================================================
# 4. LANGGRAPH STATE
# =========================================================

class FocusState(TypedDict):
    active_minutes: int
    shame_log: str
    song_catalog: dict

    break_needed: bool

    # AI recommendation
    ai_selected_song: str

    # Final song name/query after HITL override
    selected_song: str

    # Resolved direct YouTube URL
    selected_song_url: str

    # User typed song name; blank means use AI recommendation
    user_song_name: str

    loop_song: bool
    break_message: str
    approved: bool
    action_result: str


# =========================================================
# 5. STRUCTURED LLM OUTPUT
# =========================================================

class SongDecision(BaseModel):
    song_key: Literal[
        "lofi",
        "bollywood",
        "metal",
        "fun",
    ]
    message: str


# =========================================================
# 6. GROQ
# =========================================================

llm = ChatGroq(
    api_key=os.getenv("GROQ_API_KEY"),
    model=os.getenv(
        "GROQ_MODEL",
        "openai/gpt-oss-120b",
    ),
    temperature=0.2,
)

decision_llm = llm.with_structured_output(SongDecision)


# =========================================================
# 7. HELPERS
# =========================================================

def get_mcp_text(result):
    return result.contents[0].text


def extract_youtube_video_id(url: str) -> str | None:
    """Extract a YouTube video ID from common YouTube URL formats."""
    try:
        parsed = urlparse(url)

        if "youtu.be" in parsed.netloc:
            return parsed.path.strip("/") or None

        if "youtube.com" in parsed.netloc:
            if parsed.path == "/watch":
                query = parse_qs(parsed.query)
                values = query.get("v")
                if values:
                    return values[0]

            if parsed.path.startswith("/shorts/"):
                parts = parsed.path.split("/")
                if len(parts) >= 3:
                    return parts[2]

            if parsed.path.startswith("/embed/"):
                parts = parsed.path.split("/")
                if len(parts) >= 3:
                    return parts[2]

        return None

    except Exception:
        return None


def make_autoplay_url(url: str) -> str:
    """
    Convert a resolved YouTube URL into a normal watch URL with autoplay=1.
    Browser autoplay policies can still block audible autoplay.
    """
    video_id = extract_youtube_video_id(url)

    if not video_id:
        return url

    return (
        f"https://www.youtube.com/watch?v={video_id}"
        f"&autoplay=1"
    )


def make_youtube_loop_url(url: str) -> str:
    """
    Best-effort YouTube loop URL.
    Uses embed mode with autoplay + loop + playlist=video_id.
    """
    video_id = extract_youtube_video_id(url)

    if not video_id:
        return url

    return (
        f"https://www.youtube.com/embed/{video_id}"
        f"?autoplay=1"
        f"&loop=1"
        f"&playlist={video_id}"
    )


def resolve_song_to_youtube_url(song_name: str) -> tuple[str, str]:
    """
    Search YouTube for the song name using yt-dlp and return:
        (resolved_title, direct_watch_url)

    No YouTube API key is required.
    """
    clean_name = song_name.strip()

    if not clean_name:
        raise ValueError("Song name cannot be empty.")

    options = {
        "quiet": True,
        "no_warnings": True,
        "skip_download": True,
        "extract_flat": True,
        "socket_timeout": 15,
    }

    with YoutubeDL(options) as ydl:
        result = ydl.extract_info(
            f"ytsearch1:{clean_name}",
            download=False,
        )

    entries = result.get("entries") or []

    if not entries:
        raise RuntimeError(
            f"No YouTube result found for: {clean_name}"
        )

    video = entries[0]

    video_id = video.get("id")
    title = video.get("title") or clean_name

    if not video_id:
        raise RuntimeError(
            f"YouTube result did not contain a video ID for: {clean_name}"
        )

    direct_url = f"https://www.youtube.com/watch?v={video_id}"

    return title, direct_url


# =========================================================
# 8. MCP ASYNC: READ RESOURCES
# =========================================================

async def _read_mcp_context():
    async with stdio_client(server_params) as (
        read_stream,
        write_stream,
    ):
        async with ClientSession(
            read_stream,
            write_stream,
        ) as session:
            await session.initialize()

            timer_result = await session.read_resource(
                "timer://session"
            )
            timer_data = json.loads(
                get_mcp_text(timer_result)
            )

            shame_result = await session.read_resource(
                "shame://log"
            )
            shame_log = get_mcp_text(shame_result)

            song_result = await session.read_resource(
                "songs://catalog"
            )
            songs = json.loads(
                get_mcp_text(song_result)
            )

            return {
                "active_minutes": timer_data["active_minutes"],
                "shame_log": shame_log,
                "song_catalog": songs,
            }


# =========================================================
# 9. NODE 1: READ CONTEXT
# =========================================================

def read_context(state: FocusState):
    print("\n[Node] READ CONTEXT")

    data = asyncio.run(_read_mcp_context())

    print("Active minutes:", data["active_minutes"])
    print("Shame log:", data["shame_log"])

    return data


# =========================================================
# 10. NODE 2: CHECK BREAK
# =========================================================

def check_break(state: FocusState):
    print("\n[Node] CHECK BREAK")

    break_needed = (
        state["active_minutes"] >= BREAK_THRESHOLD_MINUTES
    )

    if break_needed:
        print("BREAK REQUIRED!")
    else:
        print("No break needed yet.")

    return {
        "break_needed": break_needed
    }


# =========================================================
# 11. ROUTER 1
# =========================================================

def route_break(state: FocusState):
    if state["break_needed"]:
        return "choose_song"

    return "end"


# =========================================================
# 12. MCP ASYNC: GET MCP PROMPT
# =========================================================

async def _get_focus_prompt():
    async with stdio_client(server_params) as (
        read_stream,
        write_stream,
    ):
        async with ClientSession(
            read_stream,
            write_stream,
        ) as session:
            await session.initialize()

            result = await session.get_prompt(
                "check-focus"
            )

            return (
                result
                .messages[0]
                .content
                .text
            )


# =========================================================
# 13. NODE 3: AI DEFAULT SONG SELECTION
# =========================================================

def choose_song(state: FocusState):
    print("\n[Node] AI SONG SELECTION")

    mcp_prompt = asyncio.run(
        _get_focus_prompt()
    )

    user_context = f"""
Active minutes:
{state['active_minutes']}

Previous shame history:
{state['shame_log']}

Available song categories:
{json.dumps(state['song_catalog'], indent=2)}

Choose exactly one song_key.
"""

    decision = decision_llm.invoke(
        [
            ("system", mcp_prompt),
            ("user", user_context),
        ]
    )

    song = state["song_catalog"][decision.song_key]

    ai_song_name = song["name"]

    print("AI selected:", ai_song_name)
    print("Message:", decision.message)

    return {
        "ai_selected_song": ai_song_name,
        "selected_song": ai_song_name,
        "break_message": decision.message,
    }


# =========================================================
# 14. NODE 4: HITL #1 — APPROVE BREAK?
# =========================================================

def approve_break(state: FocusState):
    print("\n[Node] HITL — BREAK APPROVAL")

    answer = interrupt(
        {
            "type": "break_approval",
            "question": "Break start karu? (y/n)",
            "message": state["break_message"],
            "active_minutes": state["active_minutes"],
            "ai_song": state["ai_selected_song"],
        }
    )

    approved = str(answer).strip().lower() in {
        "y",
        "yes",
        "true",
        "1",
    }

    return {
        "approved": approved
    }


# =========================================================
# 15. ROUTER 2
# =========================================================

def route_approval(state: FocusState):
    if state["approved"]:
        return "ask_song"

    return "rejected"


# =========================================================
# 16. NODE 5: HITL #2 — ASK SONG NAME
# =========================================================

def ask_song(state: FocusState):
    print("\n[Node] HITL — SONG CHOICE")

    answer = interrupt(
        {
            "type": "song_choice",
            "question": (
                "Kaunsa gaana chalana hai? Song name likho. "
                "AI-selected song use karna ho to blank chhodo."
            ),
            "ai_song": state["ai_selected_song"],
        }
    )

    user_song_name = str(answer).strip()

    if user_song_name:
        final_song_name = user_song_name
    else:
        final_song_name = state["ai_selected_song"]

    return {
        "selected_song": final_song_name,
        "user_song_name": user_song_name,
    }


# =========================================================
# 17. NODE 6: HITL #3 — LOOP?
# =========================================================

def ask_loop(state: FocusState):
    print("\n[Node] HITL — LOOP")

    answer = interrupt(
        {
            "type": "loop_choice",
            "question": "Gaana loop mein chalana hai? (y/n)",
            "song": state["selected_song"],
        }
    )

    loop_song = str(answer).strip().lower() in {
        "y",
        "yes",
        "true",
        "1",
    }

    return {
        "loop_song": loop_song
    }


# =========================================================
# 18. MCP ASYNC: EXECUTE BREAK TOOLS
# =========================================================

async def _execute_break_tools(song_url: str):
    async with stdio_client(server_params) as (
        read_stream,
        write_stream,
    ):
        async with ClientSession(
            read_stream,
            write_stream,
        ) as session:
            await session.initialize()

            # 1. Volume 50%
            await session.call_tool(
                "set_volume",
                {
                    "level": 50
                }
            )

            # 2. Open resolved YouTube video
            play_result = await session.call_tool(
                "play_song",
                {
                    "song_url": song_url
                }
            )

            print("MCP play_song called.")

            # 3. Give browser/video time to open before lock
            print(
                f"\nWaiting {SCREEN_LOCK_DELAY_SECONDS} seconds "
                "before screen lock..."
            )

            await asyncio.sleep(
                SCREEN_LOCK_DELAY_SECONDS
            )

            # 4. Reset focus timer
            await session.call_tool(
                "reset_session",
                {}
            )

            # 5. Lock Windows screen LAST
            await session.call_tool(
                "lock_screen",
                {}
            )

            return play_result


# =========================================================
# 19. NODE 7: RESOLVE SONG + EXECUTE BREAK
# =========================================================

def execute_break(state: FocusState):
    print("\n[Node] EXECUTE BREAK")

    song_query = state["selected_song"].strip()

    print("Searching YouTube for:", song_query)

    try:
        resolved_title, direct_url = (
            resolve_song_to_youtube_url(song_query)
        )

    except Exception as error:
        print("YouTube search failed:", error)

        return {
            "action_result": (
                "Could not find the requested song on YouTube. "
                f"Reason: {error}"
            )
        }

    print("Resolved video:", resolved_title)
    print("Direct URL:", direct_url)

    # User requested loop -> loop embed URL.
    # Otherwise normal watch URL + autoplay=1.
    if state["loop_song"]:
        final_url = make_youtube_loop_url(
            direct_url
        )
        print("Loop mode: ON")
    else:
        final_url = make_autoplay_url(
            direct_url
        )
        print("Loop mode: OFF")

    print("Final YouTube URL:", final_url)
    print("Setting volume to 50%...")
    print("Opening YouTube video...")

    asyncio.run(
        _execute_break_tools(final_url)
    )

    return {
        "selected_song": resolved_title,
        "selected_song_url": final_url,
        "action_result": (
            f"'{resolved_title}' opened on YouTube and "
            "forced break activated."
        ),
    }


# =========================================================
# 20. MCP ASYNC: WRITE SHAME LOG
# =========================================================

async def _write_rejection_log():
    async with stdio_client(server_params) as (
        read_stream,
        write_stream,
    ):
        async with ClientSession(
            read_stream,
            write_stream,
        ) as session:
            await session.initialize()

            await session.call_tool(
                "write_shame_log",
                {
                    "message": (
                        "User rejected the recommended focus break."
                    )
                }
            )


# =========================================================
# 21. NODE 8: REJECTED
# =========================================================

def rejected(state: FocusState):
    print("\n[Node] BREAK REJECTED")

    asyncio.run(
        _write_rejection_log()
    )

    return {
        "action_result": (
            "Break rejected. Shame log updated."
        )
    }


# =========================================================
# 22. BUILD GRAPH
# =========================================================

builder = StateGraph(FocusState)

builder.add_node("read_context", read_context)
builder.add_node("check_break", check_break)
builder.add_node("choose_song", choose_song)
builder.add_node("approve_break", approve_break)
builder.add_node("ask_song", ask_song)
builder.add_node("ask_loop", ask_loop)
builder.add_node("execute_break", execute_break)
builder.add_node("rejected", rejected)


# =========================================================
# 23. EDGES
# =========================================================

builder.add_edge(
    START,
    "read_context"
)

builder.add_edge(
    "read_context",
    "check_break"
)

builder.add_conditional_edges(
    "check_break",
    route_break,
    {
        "choose_song": "choose_song",
        "end": END,
    }
)

builder.add_edge(
    "choose_song",
    "approve_break"
)

builder.add_conditional_edges(
    "approve_break",
    route_approval,
    {
        "ask_song": "ask_song",
        "rejected": "rejected",
    }
)

builder.add_edge(
    "ask_song",
    "ask_loop"
)

builder.add_edge(
    "ask_loop",
    "execute_break"
)

builder.add_edge(
    "execute_break",
    END
)

builder.add_edge(
    "rejected",
    END
)


# =========================================================
# 24. HITL CHECKPOINTER
# =========================================================

memory = InMemorySaver()

graph = builder.compile(
    checkpointer=memory
)


# =========================================================
# 25. MAIN
# =========================================================

def main():
    print("\n==================================")
    print("FOCUS BREAK AI")
    print(
        f"TEST THRESHOLD = {BREAK_THRESHOLD_MINUTES} MINUTE(S)"
    )
    print("==================================")

    initial_state = {
        "active_minutes": 0,
        "shame_log": "",
        "song_catalog": {},
        "break_needed": False,
        "ai_selected_song": "",
        "selected_song": "",
        "selected_song_url": "",
        "user_song_name": "",
        "loop_song": False,
        "break_message": "",
        "approved": False,
        "action_result": "",
    }

    config = {
        "configurable": {
            "thread_id": "focus-session-1"
        }
    }

    result = graph.invoke(
        initial_state,
        config=config,
    )

    # Keep resuming until all HITL interrupts are answered.
    while "__interrupt__" in result:
        data = (
            result["__interrupt__"][0].value
        )

        print("\n==================================")
        print("HUMAN INPUT REQUIRED")
        print("==================================")

        # HITL #1 — break approval
        if data["type"] == "break_approval":
            print("\nAI:", data["message"])
            print(
                "\nActive:",
                data["active_minutes"],
                "minutes"
            )
            print(
                "AI suggested category/song:",
                data["ai_song"]
            )

            answer = input(
                "\nBreak start karu? (y/n): "
            )

        # HITL #2 — song name
        elif data["type"] == "song_choice":
            print(
                "\nAI suggestion:",
                data["ai_song"]
            )

            answer = input(
                "\nSong name likho "
                "(AI suggestion ke liye ENTER): "
            )

        # HITL #3 — loop
        elif data["type"] == "loop_choice":
            print(
                "\nSong:",
                data["song"]
            )

            answer = input(
                "\nLoop mein chalana hai? (y/n): "
            )

        else:
            raise RuntimeError(
                "Unknown HITL interrupt."
            )

        result = graph.invoke(
            Command(resume=answer),
            config=config,
        )

    print("\n==================================")
    print("FINAL RESULT")
    print("==================================")

    print(
        result.get(
            "action_result",
            "Workflow completed."
        )
    )


# =========================================================
# 26. ENTRY
# =========================================================

if __name__ == "__main__":
    main()