inference.py

from __future__ import annotations

import json
import os
from typing import Any, Dict

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # dotenv is optional

from openai import OpenAI

from src.env import EmailTriageEnv
from src.models import Action
from src.score_utils import SAFE_SCORE, safe_ratio_score


def _display_score(value: float) -> float:
    # Keep reported scores strictly inside (0, 1) even after 2-decimal rounding.
    return max(0.01, min(0.99, float(value)))


def _fmt_bool(value: bool) -> str:
    return "1" if value else "0"


def _extract_json_object(content: str) -> Dict[str, Any] | None:
    text = content.strip()
    if not text:
        return None
    try:
        parsed = json.loads(text)
        return parsed if isinstance(parsed, dict) else None
    except Exception:
        start = text.find("{")
        end = text.rfind("}")
        if start == -1 or end == -1 or end <= start:
            return None
        try:
            parsed = json.loads(text[start : end + 1])
            return parsed if isinstance(parsed, dict) else None
        except Exception:
            return None


def _format_email_for_log(email_text: str, max_len: int = 80) -> str:
    # Keep STEP lines readable by collapsing whitespace and truncating long emails.
    compact = " ".join(email_text.split())
    if len(compact) <= max_len:
        return compact
    return compact[: max_len - 3] + "..."


def _emit_start(task_id: str, model_name: str, api_enabled: bool, total_steps: int) -> None:
    print(
        "[START]"
        f" task_id={task_id}"
        f" model_name={model_name}"
        f" api_enabled={_fmt_bool(api_enabled)}"
        f" total_steps={total_steps}"
    )


def _emit_step(
    task_id: str,
    step: int,
    email_id: str,
    email: str,
    reward: float,
    cumulative_reward: float,
    category: str,
    priority: str | None,
    action: str | None,
    reply_template: str | None,
) -> None:
    reward_display = _display_score(reward)
    cumulative_reward_display = _display_score(cumulative_reward)
    line = (
        "[STEP]"
        f" task_id={task_id}"
        f" step={step:02d}"
        f" email_id={email_id}"
        f" email=\"{email}\""
        f" reward={reward_display:.2f}"
        f" cumulative_reward={cumulative_reward_display:.2f}"
        f" category={category}"
    )

    if task_id in {"task_medium", "task_hard"} and priority is not None:
        line += f" priority={priority}"
    if task_id == "task_hard" and action is not None and reply_template is not None:
        line += f" action={action}"
        line += f" reply_template={reply_template}"

    print(line)


def _emit_end(
    task_id: str,
    steps: int,
    final_score: float,
    avg_reward: float,
    category_accuracy: float | None,
    priority_accuracy: float | None,
    action_accuracy: float | None,
    reply_accuracy: float | None,
) -> None:
    final_score = _display_score(final_score)
    avg_reward = _display_score(avg_reward)

    line = (
        "[END]"
        f" task_id={task_id}"
        f" steps={steps}"
        f" final_score={final_score:.2f}"
        f" avg_reward={avg_reward:.2f}"
    )

    if category_accuracy is not None:
        category_accuracy = _display_score(category_accuracy)
        line += f" category_accuracy={category_accuracy:.2f}"
    if task_id in {"task_medium", "task_hard"} and priority_accuracy is not None:
        priority_accuracy = _display_score(priority_accuracy)
        line += f" priority_accuracy={priority_accuracy:.2f}"
    if task_id == "task_hard" and action_accuracy is not None:
        action_accuracy = _display_score(action_accuracy)
        line += f" action_accuracy={action_accuracy:.2f}"
    if task_id == "task_hard" and reply_accuracy is not None:
        reply_accuracy = _display_score(reply_accuracy)
        line += f" reply_accuracy={reply_accuracy:.2f}"

    print(line)


SYSTEM_PROMPT = (
    "You are an email triage assistant. Return strict JSON with keys: "
    "category, priority, action, reply_template. "
    "Use one of categories: billing, technical, sales, account, complaint, shipping, other. "
    "Use one of priorities: low, medium, high, urgent. "
    "Use one of actions: reply, escalate, archive. "
    "Use one of reply_templates: billing_refund, billing_invoice, account_unlock, "
    "escalate_specialist, tech_troubleshoot, sales_pricing, shipping_update, "
    "complaint_apology, archive_no_reply."
)

# Confidence thresholds for hybrid decision making
CONFIDENCE_THRESHOLD_HIGH = 0.8
CONFIDENCE_THRESHOLD_MEDIUM = 0.5


def calculate_heuristic_confidence(email_text: str, category: str) -> float:
    """
    Calculate confidence score for heuristic classification.
    Returns value between 0.0 and 1.0.
    """
    text = email_text.lower()
    confidence = 0.5  # Base confidence
    
    # Strong category indicators increase confidence
    strong_indicators = {
        "billing": ["invoice", "charged", "refund", "billing", "payment"],
        "technical": ["crash", "500", "bug", "api", "error"],
        "account": ["password", "login", "account locked"],
        "sales": ["pricing", "quote", "discount"],
        "shipping": ["package", "shipment", "tracking"],
        "complaint": ["unacceptable", "rude", "complaint"],
        "other": ["thanks", "amazing", "no further action"],
    }
    
    if category in strong_indicators:
        matches = sum(1 for indicator in strong_indicators[category] if indicator in text)
        confidence += min(matches * 0.15, 0.4)  # Up to +0.4 for multiple matches
    
    # Clear urgency words increase confidence
    urgency_words = ["urgent", "immediately", "asap", "emergency", "production down"]
    if any(word in text for word in urgency_words):
        confidence += 0.1
    
    return SAFE_SCORE(confidence)


def heuristic_policy_with_confidence(email_text: str) -> tuple[Dict[str, str], float]:
    """
    Returns action dict and confidence score.
    This enables hybrid decision making.
    """
    result = heuristic_policy(email_text)
    confidence = calculate_heuristic_confidence(email_text, result["category"])
    return result, confidence


def rule_category(text: str) -> str:
    """Determine category based on keyword matching with prioritized checks."""
    text = text.lower()

    # Priority 1: COMPLAINT (strong emotional indicators)
    complaint_keywords = [
        "complaint", "angry", "bad service", "not happy", "unhappy",
        "disappointed", "terrible", "worst", "poor service", "unacceptable",
        "rude", "ignored", "closed my case", "without resolution",
        "refund because", "missing feature"
    ]
    if any(word in text for word in complaint_keywords):
        return "complaint"

    # Priority 2: BILLING (financial indicators)
    billing_keywords = [
        "refund", "invoice", "charge", "charged", "payment", "billing",
        "money", "price charged", "overcharged", "receipt", "card", "due", "mismatch",
        "confirm if my plan renews", "canceled invoice"
    ]
    if any(word in text for word in billing_keywords):
        return "billing"

    # Priority 3: SHIPPING (delivery indicators)
    shipping_keywords = [
        "delivery", "shipping", "courier", "late",
        "package", "shipment", "tracking", "delivered", "order status",
        "received someone", "wrong order", "compensation"
    ]
    if any(word in text for word in shipping_keywords):
        return "shipping"

    # Priority 4: ACCOUNT (access/login indicators - check before technical)
    account_keywords = [
        "login", "password", "account", "reset", "signin",
        "signup", "locked", "access", "verify", "unlock", "forgot",
        "close my account", "delete all personal"
    ]
    if any(word in text for word in account_keywords):
        return "account"

    # Priority 5: SALES (pricing/business indicators)
    sales_keywords = [
        "price", "plan", "demo", "quote", "purchase",
        "buy", "subscription", "upgrade", "cost", "pricing", "discount",
        "seats", "hipaa", "soc2", "enterprise", "annual", "student",
        "50 seats", "pricing details"
    ]
    if any(word in text for word in sales_keywords):
        return "sales"

    # Priority 6: TECHNICAL (technical problem indicators)
    technical_keywords = [
        "error", "bug", "not working", "crash", "failed",
        "unable", "fix", "exception", "doesn't work", "api", "sso", "logout",
        "configure", "endpoint", "deleted", "workspace", "restore", "backup",
        "documentation", "rate limit", "examples", "500", "reset"
    ]
    if any(word in text for word in technical_keywords):
        return "technical"

    # Priority 7: OTHER (default)
    # Check for positive/closure indicators
    other_keywords = [
        "thanks", "thank you", "issue solved", "no further action",
        "just reporting", "amazing", "solved", "docs link", "broken link"
    ]
    if any(word in text for word in other_keywords):
        return "other"

    return "other"


def rule_priority(text: str) -> str:
    """Determine priority based on urgency keywords - dataset-aware."""
    text = text.lower()

    # Urgent: production down, all admins locked, payroll, immediate callback, 500 errors
    urgent_keywords = [
        "urgent", "asap", "immediately", "right now",
        "critical", "emergency", "production", "all admins", 
        "payroll", "immediate callback", "500", "all requests"
    ]

    # High: crashes, locked, not arriving, ignored, refund issues, before friday, unacceptable
    high_keywords = [
        "high", "crash", "crashes", "locked", "not arriving", 
        "ignored", "refund", "before friday", "unacceptable",
        "rude", "without resolution", "restore", "deleted",
        "wrong order", "keeps spinning", "missing"
    ]

    # Low: student, discount, hipaa, soc2, docs, examples, just reporting, thanks, amazing
    low_keywords = [
        "low", "student", "discount", "hipaa", "soc2", 
        "documentation", "examples", "just reporting", 
        "thanks", "amazing", "solved", "no further action",
        "automatically", "know if"
    ]

    if any(word in text for word in urgent_keywords):
        return "urgent"
    elif any(word in text for word in high_keywords):
        return "high"
    elif any(word in text for word in low_keywords):
        return "low"
    else:
        return "medium"


def rule_action(category: str, priority: str, text: str) -> str:
    """Determine action based on category and priority."""
    text = text.lower()
    
    if priority == "urgent":
        return "escalate"
    elif category in ["complaint"]:
        # Check if complaint needs escalation
        if any(k in text for k in ["closed my case", "without resolution", "missing feature", "refund because", "rude"]):
            return "escalate"
        return "reply"
    elif category in ["technical"]:
        # Check if technical issue needs escalation
        if any(k in text for k in ["production", "500", "restore", "deleted workspace", "crash", "all requests"]):
            return "escalate"
        return "reply"
    elif category in ["billing", "sales", "account", "shipping"]:
        return "reply"
    else:
        return "archive"


def rule_reply_template(category: str, action: str, text: str) -> str:
    """
    Map to exact reply templates used in dataset:
    billing_refund, billing_invoice, account_unlock, escalate_specialist,
    tech_troubleshoot, sales_pricing, shipping_update, complaint_apology, archive_no_reply
    """
    text = text.lower()
    
    # Archive action
    if action == "archive":
        return "archive_no_reply"
    
    # Escalate action
    if action == "escalate":
        return "escalate_specialist"
    
    # Category-specific templates for "reply" action
    if category == "billing":
        # Check for refund-related billing
        if any(k in text for k in ["refund", "charged twice", "duplicate", "overcharged", "money back"]):
            return "billing_refund"
        return "billing_invoice"
    
    elif category == "account":
        return "account_unlock"
    
    elif category == "technical":
        return "tech_troubleshoot"
    
    elif category == "sales":
        return "sales_pricing"
    
    elif category == "shipping":
        return "shipping_update"
    
    elif category == "complaint":
        return "complaint_apology"
    
    else:
        return "archive_no_reply"


def load_local_env(env_path: str = ".env") -> None:
    # Load variables from a local .env file when present.
    # Existing environment variables keep precedence.
    if not os.path.exists(env_path):
        return

    with open(env_path, "r", encoding="utf-8") as env_file:
        for raw_line in env_file:
            line = raw_line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue

            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            existing = os.environ.get(key)
            if key and (existing is None or not existing.strip()):
                os.environ[key] = value


load_local_env()


def heuristic_policy(email_text: str) -> Dict[str, str]:
    """Optimized heuristic policy using modular rule functions."""
    category = rule_category(email_text)
    priority = rule_priority(email_text)
    action = rule_action(category, priority, email_text)
    reply_template = rule_reply_template(category, action, email_text)

    return {
        "category": category,
        "priority": priority,
        "action": action,
        "reply_template": reply_template,
    }


def llm_policy(client: OpenAI, model_name: str, email_text: str) -> Dict[str, str]:
    llm_temperature = float(os.getenv("LLM_TEMPERATURE", "0.2"))
    response = client.chat.completions.create(
        model=model_name,
        temperature=llm_temperature,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    "Classify this email and return JSON only with keys "
                    "category, priority, action, reply_template:\n\n"
                    f"{email_text}"
                ),
            },
        ],
    )
    content = (response.choices[0].message.content or "").strip()
    parsed = _extract_json_object(content)
    if parsed is None:
        raise ValueError("LLM response did not contain valid JSON object")
    return {k: str(v).strip().lower() for k, v in parsed.items()}


class HybridEmailAgent:
    """
    Deterministic rule-based agent.
    
    Strategy:
    - Category: rule-based keyword matching
    - Priority: rule-based keyword matching  
    - Action: rule-based logic
    - Reply Template: mapped directly from category
    
    No LLM used for reply template - fully deterministic.
    """
    
    def __init__(self, client: OpenAI | None, model_name: str):
        self.client = client
        self.model_name = model_name
        self.rule_calls = 0
        self.llm_calls = 0

    def _should_use_llm(self, task_id: str) -> bool:
        if self.client is None:
            return False
        return task_id == "task_hard"

    def _normalize_llm_payload(self, raw: Dict[str, str], fallback: Dict[str, str]) -> Dict[str, str]:
        allowed_categories = {"billing", "technical", "sales", "account", "complaint", "shipping", "other"}
        allowed_priorities = {"low", "medium", "high", "urgent"}
        allowed_actions = {"reply", "escalate", "archive"}
        allowed_templates = {
            "billing_refund",
            "billing_invoice",
            "account_unlock",
            "escalate_specialist",
            "tech_troubleshoot",
            "sales_pricing",
            "shipping_update",
            "complaint_apology",
            "archive_no_reply",
            "general_reply",
        }

        category = raw.get("category", "").lower()
        priority = raw.get("priority", "").lower()
        action = raw.get("action", "").lower()
        reply_template = raw.get("reply_template", "").lower()

        if category not in allowed_categories:
            category = fallback["category"]
        if priority not in allowed_priorities:
            priority = fallback["priority"]
        if action not in allowed_actions:
            action = fallback["action"]
        if reply_template not in allowed_templates:
            reply_template = rule_reply_template(category, action, "")

        return {
            "category": category,
            "priority": priority,
            "action": action,
            "reply_template": reply_template,
        }
    
    def decide_action(self, email_text: str, task_id: str = "task_hard") -> Dict[str, str]:
        """
        Main decision method - fully rule-based with task-specific logic.
        """
        self.rule_calls += 1
        
        # Get category, priority, action from rules
        category = rule_category(email_text)
        priority = rule_priority(email_text)
        action = rule_action(category, priority, email_text)
        reply_template = rule_reply_template(category, action, email_text)

        fallback_hard = {
            "category": category,
            "priority": priority,
            "action": action,
            "reply_template": reply_template,
        }

        if self._should_use_llm(task_id):
            try:
                llm_raw = llm_policy(client=self.client, model_name=self.model_name, email_text=email_text)
                llm_payload = self._normalize_llm_payload(raw=llm_raw, fallback=fallback_hard)
                self.llm_calls += 1

                if task_id == "task_easy":
                    return {
                        "category": llm_payload["category"],
                        "priority": "medium",
                        "action": "reply",
                        "reply_template": "general_reply",
                    }
                if task_id == "task_medium":
                    return {
                        "category": llm_payload["category"],
                        "priority": llm_payload["priority"],
                        "action": llm_payload["action"],
                        "reply_template": "general_reply",
                    }
                return llm_payload
            except Exception:
                pass
        
        # Task-specific handling with safe defaults
        if task_id == "task_easy":
            # Only category matters - use safe defaults for other fields
            return {
                "category": category,
                "priority": "medium",
                "action": "reply",
                "reply_template": "general_reply",
            }
        elif task_id == "task_medium":
            # Category + Priority + Action - use generic reply template
            return {
                "category": category,
                "priority": priority,
                "action": action,
                "reply_template": "general_reply",
            }
        else:  # task_hard
            # Category + Priority + Action + Reply Template
            return {
                "category": category,
                "priority": priority,
                "action": action,
                "reply_template": reply_template,
            }
    
    def get_stats(self) -> Dict[str, int]:
        """Return usage statistics."""
        return {
            "rule_calls": self.rule_calls,
            "llm_calls": self.llm_calls,
        }


def choose_action(client: OpenAI | None, model_name: str, email_text: str) -> Dict[str, str]:
    """Legacy function - creates temporary agent."""
    agent = HybridEmailAgent(client, model_name)
    return agent.decide_action(email_text)


def make_client() -> OpenAI | None:
    """Create OpenAI client from environment variables."""
    api_base_url = os.getenv("API_BASE_URL")
    hf_token = os.getenv("HF_TOKEN")
    
    if not api_base_url:
        return None
    
    return OpenAI(
        base_url=api_base_url,
        api_key=hf_token or ""
    )


def _new_component_metric() -> Dict[str, float]:
    return {"correct": 0, "total": 0, "accuracy": 0.05}


def _safe_accuracy(correct: int, total: int) -> float:
    return SAFE_SCORE(safe_ratio_score(correct=correct, total=total))


def run_task(task_id: str, client: OpenAI | None, model_name: str) -> Dict[str, object]:
    env = EmailTriageEnv(task_id=task_id)
    required_components = {
        "category": env.task.require_category,
        "priority": env.task.require_priority,
        "action": env.task.require_action,
        "reply": env.task.require_reply_template,
    }
    obs = env.reset()
    _emit_start(
        task_id=task_id,
        model_name=model_name,
        api_enabled=bool(os.getenv("API_BASE_URL", "").strip()),
        total_steps=len(env.dataset),
    )
    
    # Create hybrid agent for this task
    agent = HybridEmailAgent(client, model_name)

    done = False
    step_count = 0
    step_rewards: list[float] = []
    cumulative_rewards: list[float] = []

    component_accuracy = {
        "category": _new_component_metric(),
        "priority": _new_component_metric(),
        "action": _new_component_metric(),
        "reply": _new_component_metric(),
    }

    while not done:
        step_count += 1
        current_email_text = obs.email_text
        action_payload = agent.decide_action(email_text=obs.email_text, task_id=task_id)
        action = Action.model_validate(action_payload)
        obs, reward, done, info = env.step(action)
        step_rewards.append(reward)
        cumulative_rewards.append(_display_score(env.state().cumulative_reward))

        truth = info["truth"]
        category_correct = action.category is not None and action.category.value == truth["category"]
        priority_correct = action.priority is not None and action.priority.value == truth["priority"]
        action_correct = action.action is not None and action.action.value == truth["action"]
        reply_correct = action.reply_template is not None and action.reply_template == truth["reply_template"]

        if required_components["category"]:
            component_accuracy["category"]["total"] += 1
            component_accuracy["category"]["correct"] += int(category_correct)
        if required_components["priority"]:
            component_accuracy["priority"]["total"] += 1
            component_accuracy["priority"]["correct"] += int(priority_correct)
        if required_components["action"]:
            component_accuracy["action"]["total"] += 1
            component_accuracy["action"]["correct"] += int(action_correct)
        if required_components["reply"]:
            component_accuracy["reply"]["total"] += 1
            component_accuracy["reply"]["correct"] += int(reply_correct)

        _emit_step(
            task_id=task_id,
            step=step_count,
            email_id=info["email_id"],
            email=_format_email_for_log(current_email_text),
            reward=reward,
            cumulative_reward=env.state().cumulative_reward,
            category=action_payload["category"],
            priority=action_payload.get("priority"),
            action=action_payload.get("action"),
            reply_template=action_payload.get("reply_template"),
        )

    metric_report: Dict[str, float | None] = {
        "category": None,
        "priority": None,
        "action": None,
        "reply": None,
    }
    for key in ["category", "priority", "action", "reply"]:
        if required_components[key]:
            correct = int(component_accuracy[key]["correct"])
            total = int(component_accuracy[key]["total"])
            accuracy = _safe_accuracy(correct=correct, total=total)
            component_accuracy[key]["accuracy"] = SAFE_SCORE(accuracy)
            metric_report[key] = float(SAFE_SCORE(accuracy))

    category_accuracy = metric_report["category"]
    priority_accuracy = metric_report["priority"]
    action_accuracy = metric_report["action"]
    reply_accuracy = metric_report["reply"]
    if category_accuracy is not None:
        category_accuracy = SAFE_SCORE(category_accuracy)
    if priority_accuracy is not None:
        priority_accuracy = SAFE_SCORE(priority_accuracy)
    if action_accuracy is not None:
        action_accuracy = SAFE_SCORE(action_accuracy)
    if reply_accuracy is not None:
        reply_accuracy = SAFE_SCORE(reply_accuracy)

    final_score = _display_score(SAFE_SCORE(env.final_score()))
    cumulative_reward = env.state().cumulative_reward
    avg_reward = _display_score(cumulative_reward / max(step_count, 1))
    
    # Get agent stats
    agent_stats = agent.get_stats()

    _emit_end(
        task_id=task_id,
        steps=step_count,
        final_score=final_score,
        avg_reward=avg_reward,
        category_accuracy=category_accuracy,
        priority_accuracy=priority_accuracy,
        action_accuracy=action_accuracy,
        reply_accuracy=reply_accuracy,
    )

    return {
        "task_id": task_id,
        "steps": step_count,
        "step_rewards": step_rewards,
        "cumulative_rewards": cumulative_rewards,
        "final_score": final_score,
        "avg_reward": avg_reward,
        "component_accuracy": component_accuracy,
        "agent_stats": agent_stats,
    }


def main() -> None:
    model_name = os.getenv("MODEL_NAME", "gpt-4o-mini").strip()
    client = make_client()

    task_ids = ["task_easy", "task_medium", "task_hard"]
    task_results: list[Dict[str, object]] = []

    for index, task_id in enumerate(task_ids):
        if index > 0:
            print()
        print("=" * 64)
        print(f"[RUNNING] task_id={task_id}")
        result = run_task(task_id=task_id, client=client, model_name=model_name)
        task_results.append(result)

    total_steps = sum(int(result["steps"]) for result in task_results)
    if task_results:
        mean_final_score = sum(float(result["final_score"]) for result in task_results) / len(task_results)
    else:
        mean_final_score = 0.05

    mean_final_score = _display_score(SAFE_SCORE(mean_final_score))
    score_parts = " ".join(
        f"{str(result['task_id'])}={_display_score(float(result['final_score'])):.2f}"
        for result in task_results
    )
    print()
    print(
        "[SUMMARY]"
        f" tasks_run={len(task_results)}"
        f" total_steps={total_steps}"
        f" mean_final_score={mean_final_score:.2f}"
        f" {score_parts}"
    )


if __name__ == "__main__":
    main()