import yaml
import logging
from agent.database import get_connection, release_connection

logger = logging.getLogger("agent.system_skills.jelaskan_workflow")

def run(inputs: dict) -> dict:
    """
    Retrieves the list of workflows or detailed YAML content for a specific workflow,
    and formats it with a Mermaid flowchart if a specific workflow is requested.
    """
    workflow_name = inputs.get("workflow_name")
    
    # Fallback to query if workflow_name was not extracted properly
    if not workflow_name and "_raw_query" in inputs:
        query = inputs["_raw_query"].lower()
        if "customer service" in query:
            workflow_name = "customer service"
        elif "order_fulfillment" in query or "order fulfillment" in query:
            workflow_name = "order_fulfillment"
    
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            if workflow_name:
                cur.execute("SELECT yaml_content FROM workflow_definitions WHERE workflow_name = %s", (workflow_name,))
                row = cur.fetchone()
                if row:
                    yaml_str = row[0]
                    parsed = yaml.safe_load(yaml_str)
                    
                    # Generate Mermaid Diagram
                    mermaid = ["```mermaid", "stateDiagram-v2"]
                    states = parsed.get("states", [])
                    
                    for state in states:
                        name = state.get("name")
                        desc = state.get("description", "No description")
                        # Clean description for mermaid
                        clean_desc = desc.replace('"', "'").replace("\n", " ")
                        mermaid.append(f"    {name} : {clean_desc}")
                        
                        transitions = state.get("transitions", [])
                        for tr in transitions:
                            to_state = tr.get("to")
                            mermaid.append(f"    {name} --> {to_state}")
                            
                        # Mark start and end
                        if state.get("type") == "start":
                            mermaid.append(f"    [*] --> {name}")
                        if state.get("type") == "end":
                            mermaid.append(f"    {name} --> [*]")
                            
                    mermaid.append("```")
                    
                    # Format output nicely
                    lines = [
                        f"### Workflow: **{parsed.get('workflow', workflow_name)}**",
                        f"_{parsed.get('description', '')}_",
                        "",
                        "#### Diagram Alur Kerja",
                        "",
                        "\n".join(mermaid),
                        "",
                        "#### Detail State",
                        "",
                    ]
                    
                    for state in states:
                        lines.append(f"- **{state.get('name')}** ({state.get('type', 'intermediate')}): {state.get('description')}")
                        skills = state.get("skills", [])
                        if skills:
                            lines.append(f"  - *Skills*: {', '.join(skills)}")
                            
                    output = "\n".join(lines)
                else:
                    output = f"Workflow '{workflow_name}' not found."
            else:
                cur.execute("SELECT workflow_name, description, type, state_count FROM workflow_definitions ORDER BY workflow_name")
                rows = cur.fetchall()
                if rows:
                    lines = ["Available Workflows:"]
                    for r in rows:
                        name, desc, wtype, states = r
                        desc_str = f" - {desc}" if desc else ""
                        lines.append(f"- **{name}** (Type: {wtype}, States: {states}){desc_str}")
                    output = "\n".join(lines)
                else:
                    output = "No workflows found in the registry."
                    
        return {"output": output}
        
    except Exception as e:
        logger.error(f"Error fetching workflows: {e}")
        return {"output": f"Error fetching workflows: {str(e)}"}
    finally:
        release_connection(conn)