import gradio as gr
import os
import re
import json
from huggingface_hub import InferenceClient

# -------------------------------------------------------------------------
# CONSTANTS & SETUP
# -------------------------------------------------------------------------
HF_TOKEN = os.environ.get("HF_TOKEN")
client = InferenceClient(token=HF_TOKEN)

TRANSCRIBE_MODEL = "CohereLabs/cohere-transcribe-03-2026"
LLM_MODEL = "CohereLabs/tiny-aya-earth"

# -------------------------------------------------------------------------
# CORE LOGIC & POST-PROCESSING MODULES
# -------------------------------------------------------------------------
def clean_and_parse_json(raw_text):
    """
    Small models love adding conversational conversational fluff or markdown wrappers.
    This module cuts out thought tags and code-blocks to isolate raw data.
    """
    try:
        # Strip away potential thoughts block if model uses them
        cleaned = re.sub(r"<think>.*?</think>", "", raw_text, flags=re.DOTALL)
        
        # Extract content between markdown code blocks if present
        json_match = re.search(r"```json\s*(.*?)\s*```", cleaned, re.DOTALL)
        if json_match:
            cleaned = json_match.group(1)
            
        return json.loads(cleaned.strip())
    except Exception:
        # Robust fallback state if model fails to output clean structural JSON
        return {
            "error": "Failed to parse structured layout.",
            "raw_payload": raw_text,
            "tasks": [["Aya parsing anomaly", "High", "Review raw logs"]]
        }

def generate_system_payload(transcript, workflow_type, extra_instructions):
    """
    State machine that routes the transcript into a strict system-prompt schema.
    """
    if not transcript:
        return "No audio track detected.", [], "System Idle."

    # Establish strict structural guardrails for Tiny Aya
    system_prompt = (
        "You are an authoritative backend systems architect. Your job is to analyze incoming speech context "
        "and output a strict JSON structure containing three root keys: 'summary' (string), 'tasks' (a list of lists "
        "where each item is [Task Title, Priority Low/Medium/High, Status]), and 'code' (a clean, unencumbered script or markdown block).\n"
        "Do not write conversational introductions or conclusions. Output ONLY valid JSON."
    )
    
    user_content = f"Workflow: {workflow_type}\nContext: {transcript}\nModifier: {extra_instructions}"
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content}
    ]
    
    try:
        # Call the text generation layer via serverless infrastructure
        response = client.chat_completion(
            model=LLM_MODEL,
            messages=messages,
            max_tokens=1512,
            temperature=0.1 # Lock down low temperature for maximum consistency
        )
        raw_output = response.choices[0].message.content
        
        # Process output through the regex validation wrapper
        parsed_data = clean_and_parse_json(raw_output)
        
        summary = parsed_data.get("summary", "No summary generated.")
        tasks = parsed_data.get("tasks", [["Generic task parsed", "Medium", "Pending"]])
        code_block = parsed_data.get("code", "# No code artifacts extracted.")
        
        # Write file artifact dynamically to the local container disk for download
        output_filename = "fone_system_spec.md"
        with open(output_filename, "w") as f:
            f.write(f"# FONE GENERATED SPEC\n\n## Summary\n{summary}\n\n## Artifact\n```\n{code_block}\n```")
            
        return summary, tasks, code_block, output_filename
        
    except Exception as e:
        return f"Pipeline execution failure: {str(e)}", [], f"```python\n# Raw trace\n{str(e)}\n```", None

def run_pipeline(audio_path, language_code, workflow_type, extra_instructions):
    """
    The orchestrator linking the 2B audio model and 3.35B language model.
    """
    if not audio_path:
        return "Error: Please record an audio segment.", [], "No data.", None
        
    try:
        # Step 1: Fire audio array at Cohere Transcribe 2B
        transcript = client.automatic_speech_recognition(
            audio_path, 
            model=TRANSCRIBE_MODEL
        )
        # Handle empty transcription events
        if not transcript.strip():
            transcript = "[Silence or non-speech artifact detected]"
    except Exception as e:
        transcript = f"[Transcription service offline: {str(e)}]"

    # Step 2: Route text output into structural prompt engine
    summary, tasks, code_block, file_out = generate_system_payload(transcript, workflow_type, extra_instructions)
    
    return summary, tasks, code_block, file_out

# -------------------------------------------------------------------------
# INTERFACE ORCHESTRATION (GRADIO 6 BLOCK COMPLIANT)
# -------------------------------------------------------------------------
with gr.Blocks(title="fone // Sovereign Workspace") as demo:
    gr.Markdown("## 🎛️ fone // Voice Architecture Pipeline")
    gr.Markdown("*Decentralized translation, transcription, and system generation utilizing modular 2B and 3B open weights.*")
    
    with gr.Row():
        # Input Controller Panel
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Input Core")
            audio_feed = gr.Audio(type="filepath", label="Voice Master")
            
            with gr.Row():
                lang_selector = gr.Dropdown(
                    choices=["en", "fr", "es", "de", "ar", "ja", "ko"], 
                    value="en", 
                    label="Target Language"
                )
                workflow_selector = gr.Dropdown(
                    choices=["Feature Engineering Specification", "Database Schema Map", "Automated System Scripts"], 
                    value="Feature Engineering Specification", 
                    label="Routing Class"
                )
                
            instruction_overlay = gr.Textbox(
                label="Execution Modifiers (Optional)", 
                placeholder="e.g., Target an Ubuntu production stack or output valid markdown tables..."
            )
            
            trigger_btn = gr.Button("Execute Pipeline Trace", variant="primary")
            file_download = gr.File(label="Exported System Artifacts")

        # Output Target Panel
        with gr.Column(scale=2):
            gr.Markdown("### 📤 Orchestration Hub")
            
            with gr.Tabs():
                with gr.TabItem("System Summary"):
                    summary_display = gr.Textbox(label="Extracted Scope", lines=8, interactive=False)
                    
                with gr.TabItem("Task Allocation Matrix"):
                    task_matrix = gr.Dataframe(
                        headers=["Objective / Component", "Priority Rank", "Status Context"],
                        datatype=["str", "str", "str"],
                        row_count=5,
                        col_count=(3, "fixed")
                    )
                    
                with gr.TabItem("Code & Schema Artifacts"):
                    code_display = gr.Code(language="markdown", label="Isolated Scripts", lines=15)

    # Wire event listener link
    trigger_btn.click(
        fn=run_pipeline,
        inputs=[audio_feed, lang_selector, workflow_selector, instruction_overlay],
        outputs=[summary_display, task_matrix, code_display, file_download]
    )

if __name__ == "__main__":
    # Force dark high-contrast monochrome runtime configurations
    demo.launch(theme=gr.themes.Monochrome())