# agent.py

import json
from typing import Dict, Any

from re_act import (
    get_plan_from_llm,
    think,
    act,
    store_name_email_mapping,
    extract_sender_info,
    client
)
from schemas import PlanStep
from logger import logger  # from logger.py


def run_agent():
    """
    Main REPL loop for the email agent.
    """
    logger.info("Starting Email Agent REPL...")
    print("🤖 Email Agent ready. Type 'exit' to quit.\n")

    while True:
        try:
            user_query = input("🗨 You: ").strip()
            logger.info("Received user input: %s", user_query)

            if user_query.lower() in ("exit", "quit"):
                logger.info("Exit command detected, shutting down agent.")
                print("👋 Goodbye!")
                break

            # 1) Generate plan
            try:
                plan = get_plan_from_llm(user_query)
                logger.debug("Generated plan: %s", plan)
            except Exception as e:
                logger.error("Failed to generate plan: %s", e)
                print(f"❌ Could not generate a plan: {e}")
                continue

            # print plan for user transparency
            print("\n\nplan:")
            print(plan)
            print("\n\n")

            results: Dict[str, Any] = {}

            # 2) Execute each plan step
            for step in plan.plan:
                logger.info("Processing step: %s", step.action)

                if step.action == "done":
                    logger.info("Encountered 'done' action. Plan complete.")
                    print("✅ Plan complete.")
                    break

                try:
                    should_run, updated_step, user_prompt = think(step, results, user_query)
                    logger.debug(
                        "Think outcome for '%s': should_run=%s, updated_step=%s, user_prompt=%s",
                        step.action, should_run, updated_step, user_prompt
                    )
                except Exception as e:
                    logger.error("Error in think() for step '%s': %s", step.action, e)
                    print(f"❌ Error in planning step '{step.action}': {e}")
                    break

                # Handle user prompt (e.g., missing email mapping)
                if user_prompt:
                    logger.info("User prompt required: %s", user_prompt)
                    print(f"❓ {user_prompt}")
                    user_input = input("📧 Email: ").strip()

                    try:
                        sender_info = extract_sender_info(user_query)
                        sender_intent = sender_info.get("sender_intent", "")
                        store_name_email_mapping(sender_intent, user_input)
                        logger.info("Stored mapping: %s -> %s", sender_intent, user_input)
                        print(f"✅ Stored mapping: {sender_intent} → {user_input}")

                        # Retry current step
                        should_run, updated_step, _ = think(step, results, user_query)
                        logger.debug(
                            "Post-mapping think outcome: should_run=%s, updated_step=%s",
                            should_run, updated_step
                        )
                    except Exception as e:
                        logger.error("Error storing mapping or retrying step '%s': %s", step.action, e)
                        print(f"❌ Error storing mapping or retrying step: {e}")
                        break

                if not should_run:
                    logger.info("Skipping step: %s", step.action)
                    print(f"⏭️  Skipping `{step.action}`")
                    continue

                # Execute the action
                try:
                    output = act(updated_step)
                    results[updated_step.action] = output
                    logger.info("Action '%s' executed successfully.", updated_step.action)
                    print(f"🔧 Ran `{updated_step.action}`")
                except Exception as e:
                    logger.error("Error executing action '%s': %s", updated_step.action, e)
                    print(f"❌ Error running `{updated_step.action}`: {e}")
                    break

            # 3) Summarize results via LLM
            try:
                summary_rsp = client.chat.completions.create(
                    model="gpt-4o-mini",
                    temperature=0.0,
                    messages=[
                        {"role": "system", "content": "Summarize these results for the user in a friendly way."},
                        {"role": "assistant", "content": json.dumps(results)}
                    ],
                )
                summary = summary_rsp.choices[0].message.content
                logger.info("Summary generated successfully.")
                print("\n📋 Summary:\n" + summary)
            except Exception as e:
                logger.error("Failed to generate summary: %s", e)
                print("\n❌ Failed to generate summary.")

            print("\nAnything else I can help you with?\n")

        except KeyboardInterrupt:
            logger.info("KeyboardInterrupt received, shutting down.")
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            logger.exception("Unexpected error in REPL loop: %s", e)
            print(f"❌ Unexpected error: {e}")
            continue


if __name__ == "__main__":
    run_agent()