import sqlite3
import operator
from typing import TypedDict, List, Dict, Any, Annotated
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_ollama import ChatOllama 

class SongState(TypedDict):
    script: str
    extracted_themes: str
    interviewer_history: Annotated[List[BaseMessage], operator.add]
    user_preferences: Annotated[List[str], operator.add]
    generated_lyrics: str
    music_style_tags: str
    question_count: int

local_llm = ChatOllama(model="llama3", temperature=0.5) # Lower temperature for more stable adherence

# --- GRAPH NODES ---

def analyze_script_node(state: SongState):
    print("\n[System]: Analyzing script tones...")
    prompt = f"Analyze this script and output raw musical characteristics (Mood, Theme, Tempo):\n{state['script']}"
    response = local_llm.invoke(prompt)
    return {"extracted_themes": response.content, "question_count": 0, "interviewer_history": []}

def interview_user_node(state: SongState):
    history = state.get("interviewer_history", [])
    
    system_instruction = SystemMessage(
        content=(
            f"You are a music producer interviewing an artist based on this script analysis:\n{state['extracted_themes']}\n\n"
            "Ask EXACTLY ONE short, unique question to narrow down style/instruments.\n"
            "Do NOT repeat yourself. Look at the history and move to a new detail (e.g., pace, vocals, era).\n"
            "Output ONLY the question text. No conversational filler."
        )
    )
    
    response = local_llm.invoke([system_instruction] + history)
    print(f"\n[AI Producer]: {response.content}")
    
    return {
        "interviewer_history": [AIMessage(content=response.content)],
        "question_count": state["question_count"] + 1
    }

def human_feedback_node(state: SongState):
    """Pauses internally in the terminal to grab your answer naturally!"""
    user_reply = input("[Your Answer]: ")
    return {
        "user_preferences": [user_reply],
        "interviewer_history": [HumanMessage(content=user_reply)]
    }

def lyrics_and_planner_node(state: SongState):
    print("\n[System]: Finalizing lyrics and music tags...")
    answers = "\n".join(state.get("user_preferences", []))
    
    prompt = f"Write structured song lyrics ([Verse], [Chorus], [Bridge]) based on:\n{state['script']}\nStyle adjustments:\n{answers}"
    lyrics_out = local_llm.invoke(prompt)
    
    tag_prompt = f"Convert these preferences into a brief comma-separated list of musical style tags (e.g., 'lo-fi, 120bpm, synth'):\n{answers}"
    tags_out = local_llm.invoke(tag_prompt)
    
    return {"generated_lyrics": lyrics_out.content, "music_style_tags": tags_out.content}

# --- CONTROL ROUTER ---

def interview_router(state: SongState):
    if int(state.get("question_count", 0)) >= 3:
        return "compile_tracks"
    return "ask_more"

# --- COMPILING THE WORKFLOW ---

conn = sqlite3.connect("song_memory.db", check_same_thread=False)
memory = SqliteSaver(conn)

workflow = StateGraph(SongState)
workflow.add_node("analyzer", analyze_script_node)
workflow.add_node("interviewer", interview_user_node)
workflow.add_node("human_input", human_feedback_node) # New node
workflow.add_node("composer", lyrics_and_planner_node)

workflow.set_entry_point("analyzer")
workflow.add_edge("analyzer", "interviewer")

# Route to human input first, then check the counter loop
workflow.add_edge("interviewer", "human_input")
workflow.add_conditional_edges(
    "human_input",
    interview_router,
    {"ask_more": "interviewer", "compile_tracks": "composer"}
)
workflow.add_edge("composer", END)

langgraph_app = workflow.compile(checkpointer=memory)