from dotenv import load_dotenv
load_dotenv()

import json
from datetime import datetime
import pandas as pd
import yaml
from crewai import Crew, Process, Agent, Task

# ✅ Import your glossary RAG tool
from inclusive_writer.tools.glossary_rag import glossary_rag

# -------------------------------
# Utility: Load YAML files with encoding fallback
# -------------------------------
def load_yaml_file(path):
    for enc in ("utf-8", "utf-8-sig", "cp1252", "latin-1"):
        try:
            with open(path, "r", encoding=enc) as f:
                return yaml.safe_load(f)
        except UnicodeDecodeError:
            continue
    raise UnicodeDecodeError(f"Unable to read {path} with common encodings. Re-save the file as UTF-8.")

# -------------------------------
# Utility: Replace placeholders in task descriptions
# -------------------------------
def substitute_placeholders(template: str, **kwargs) -> str:
    for k, v in kwargs.items():
        template = template.replace("{" + k + "}", v)
    return template

# -------------------------------
# Utility: Clean terminal output
# -------------------------------
def print_section(title, content):
    print("\n" + "="*60)
    print(f"🧠 {title.upper()}")
    print("="*60)
    print(content)

# === CONFIG ===
agents_config = load_yaml_file("inclusive_writer/config/agents.yaml")
tasks_config = load_yaml_file("inclusive_writer/config/tasks.yaml")

# === AGENTS ===
analyzer = Agent(**agents_config["analyzer"])

strategist = Agent(**agents_config["strategist"])

# ✅ Attach glossary_rag to generator
generator = Agent(
    role=agents_config["generator"]["role"],
    goal=agents_config["generator"]["goal"],
    backstory=agents_config["generator"].get("backstory", ""),
    tools=[glossary_rag],  # 🔌 TOOL INTEGRATION
    verbose=True
)

judge = Agent(**agents_config["judge"])

# === QUERY INPUT ===
# query = "Write a bedtime story about a child who grows up to be a policeman."
print(f"Input Query: ")
query = input("Enter your writing query (or press Enter to use default): ") 

# === TASKS ===
analyze_query_task = Task(
    description=substitute_placeholders(tasks_config["analyze_query_task"]["description"], query=query),
    expected_output=tasks_config["analyze_query_task"]["expected_output"],
    agent=analyzer
)

generate_guidance_task = Task(
    description=substitute_placeholders(tasks_config["generate_guidance_task"]["description"], query=query),
    expected_output=tasks_config["generate_guidance_task"]["expected_output"],
    agent=strategist
)

generate_inclusive_task = Task(
    description=substitute_placeholders(tasks_config["generate_inclusive_task"]["description"], query=query),
    expected_output=tasks_config["generate_inclusive_task"]["expected_output"],
    agent=generator
)

evaluate_output_task = Task(
    description=substitute_placeholders(tasks_config["evaluate_output_task"]["description"], query=query),
    expected_output=tasks_config["evaluate_output_task"]["expected_output"],
    agent=judge
)

# === STEP 1: ANALYSIS ===
crew_analysis = Crew(agents=[analyzer], tasks=[analyze_query_task], process=Process.sequential)
analysis_result = crew_analysis.kickoff(inputs={"query": query})
analysis_output = str(analysis_result)
print_section("Analysis Output", analysis_output)

# === STEP 2: GUIDANCE ===
crew_guidance = Crew(agents=[strategist], tasks=[generate_guidance_task], process=Process.sequential)
guidance_result = crew_guidance.kickoff(inputs={"query": query, "analysis": analysis_output})
guidance_output = str(guidance_result)
print_section("Guidance Output", guidance_output)

# === STEP 3: INCLUSIVE GENERATION ===
crew_generation = Crew(agents=[generator], tasks=[generate_inclusive_task], process=Process.sequential)
inclusive_result = crew_generation.kickoff(inputs={
    "query": query,
    "analysis": analysis_output,
    "guidance": guidance_output
})
inclusive_output = str(inclusive_result)
print_section("Inclusive Generation Output", inclusive_output)

# === STEP 4: EVALUATION ===
crew_evaluation = Crew(agents=[judge], tasks=[evaluate_output_task], process=Process.sequential)
evaluation_result = crew_evaluation.kickoff(inputs={"query": query, "output": inclusive_output})
evaluation_output = str(evaluation_result)
print_section("Evaluation Output", evaluation_output)

# === PARSE EVALUATION JSON ===
try:
    judge_data = json.loads(evaluation_output)
    scores = judge_data.get("scores", {})
    notes = judge_data.get("notes", {})
    fairness_index = sum(scores.values()) / 12 * 100 if scores else 0.0
except Exception:
    scores, notes = {}, {}
    fairness_index = 0.0

# === SAVE RESULTS TO EXCEL ===
row = {
    "Timestamp": datetime.now().isoformat(timespec="seconds"),
    "Query": query,
    "Analysis": analysis_output,
    "Guidance": guidance_output,
    "Inclusive Output": inclusive_output,
    "Judge Scores": json.dumps(scores),
    "Judge Notes": json.dumps(notes),
    "Fairness Index": fairness_index
}

output_file = "output_log.xlsx"
try:
    df = pd.read_excel(output_file)
    df = pd.concat([df, pd.DataFrame([row])], ignore_index=True)
except FileNotFoundError:
    df = pd.DataFrame([row])

df.to_excel(output_file, index=False)
print("\n✅ Results saved to", output_file)