import gradio as gr
from agent_state import AgentState
from workflow import build_graph
from tools.transcriber import transcribe_fast # We need to import the tool directly
from datetime import datetime

# --- Main Application Logic ---

# Build the graph once when the app starts
app = build_graph()

# This function is now ONLY for transcription
def run_transcription(audio_file):
    """
    Step 1: Runs ONLY the transcription tool and prepares the state.
    """
    if not audio_file:
        return None, None, gr.update(visible=False)

    print("--- Step 1: Transcribing Audio ---")
    # Create a new agent state for this session
    initial_state = AgentState(audio_path=audio_file)
    
    # Call the transcription tool directly
    state_after_transcription = transcribe_fast(initial_state)

    # Return the text for the UI, the updated state for the session,
    # and make the next button visible.
    return (
        state_after_transcription,
        state_after_transcription.transcribed_text,
        gr.update(visible=True)
    )

# This new function handles the rest of the agent's workflow
def generate_report(current_state):
    """
    Step 2: Takes the transcribed state and runs it through the main graph
    to generate the report and enable the review process.
    """
    if not current_state or not current_state.transcribed_text:
        return current_state, "Transcription not found. Please complete Step 1.", gr.update(visible=False)

    print("--- Step 2: Generating News Report ---")
    
    final_state = None
    # Run the stream. The graph will start from the news_reporter node
    # because the transcribed_text already exists.
    # Note: Your graph needs to be robust enough to handle this.
    # A simple way is to have the first node check if transcription exists.
    # If not, run it. If yes, skip to the next node.
    # For now, we assume the graph continues from where the state is.
    
    # A simplified invocation for this flow would be to call the reporter tool directly
    # and then handle the loop, but let's stick to the graph.
    # To make this work, we'll manually call the next step for clarity.
    
    from tools.news_reporter import create_news_report # this is our tool
    state_after_report = create_news_report(current_state)

    return state_after_report, state_after_report.news_report, gr.update(visible=True)


def handle_revision(feedback, current_state):
    """
    Handles the revision loop.
    """
    if not feedback:
        return current_state, current_state.news_report, "Please provide feedback to revise."

    print("Revising with feedback...")
    current_state.feedback = feedback
    current_state.approved = False
    
    # Re-run the summarization/report tool with the new feedback
    from tools.news_reporter import create_news_report
    state_after_revision = create_news_report(current_state)

    return state_after_revision, state_after_revision.news_report, "✅ Report revised. Please review again."


def handle_save(current_state):
    """
    Handles the final save action.
    """
    print("Saving final report...")
    from tools.saver import save_summary
    save_summary(current_state) # Call the save tool directly
    
    return "✅ Final report has been saved successfully!"




#####################################
# --- Gradio UI ---
#####################################
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎙️ Audio to News Agent")
    gr.Markdown(f"### Reporting live, {datetime.now().strftime('%B %d, %Y')}.")
    
    agent_state_gr = gr.State(value=None)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Step 1: Transcribe Audio")
            audio_input = gr.Audio(type="filepath", sources=["upload", "microphone"], label="Upload or Record")
            transcribe_btn = gr.Button("1️⃣ Transcribe Audio", variant="secondary")
            
            # This button will appear after step 1 is complete
            generate_btn = gr.Button("2️⃣ Generate News Report", variant="primary", visible=False)
            
        with gr.Column(scale=2):
            transcribed_output = gr.Textbox(label="📝 Transcription Result", lines=5, interactive=False)
            report_output = gr.Textbox(label="📰 Generated News Report", lines=10, interactive=False)
            
            with gr.Group(visible=False) as review_group:
                feedback_input = gr.Textbox(label="❌ Provide Feedback for Revision", lines=2)
                with gr.Row():
                    revise_btn = gr.Button("🔁 Revise Report")
                    save_btn = gr.Button("✅ Approve & Save Report", variant="primary")
                status_output = gr.Textbox(label="Status", interactive=False)

    # --- Button Click Logic ---

    transcribe_btn.click(
        fn=run_transcription,
        inputs=[audio_input],
        outputs=[agent_state_gr, transcribed_output, generate_btn]
    )
    
    generate_btn.click(
        fn=generate_report,
        inputs=[agent_state_gr],
        outputs=[agent_state_gr, report_output, review_group]
    )

    revise_btn.click(
        fn=handle_revision,
        inputs=[feedback_input, agent_state_gr],
        outputs=[agent_state_gr, report_output, status_output]
    )

    save_btn.click(
        fn=handle_save,
        inputs=[agent_state_gr],
        outputs=[status_output]
    )

demo.launch(share=True)