"""B1: Task/Problem Classification using LLM (gpt-4o-mini)."""

import json
import os

from dotenv import load_dotenv

from .base import BaseAnalyzer
from .helpers import clean_user_content, is_compaction_message, is_codex_system_message, get_first_user_prompt
from .llm_cache import get_cached, set_cache

# Load OpenAI key from .prod.env
load_dotenv(os.path.join(os.path.dirname(__file__), "..", "..", ".prod.env"), override=True)

VALID_TASK_TYPES = {
    "bug_fix", "build_feature", "refactor", "understand_code",
    "debug", "test", "config_infra", "code_review", "exploration",
}
VALID_COMPLEXITIES = {"simple", "moderate", "complex"}
VALID_DOMAINS = {"backend", "frontend", "mobile", "infra", "data", "mixed"}

SYSTEM_PROMPT = """You are a classifier for AI coding sessions. Given the first user prompt from a coding session, classify:

1. task_type: one of bug_fix, build_feature, refactor, understand_code, debug, test, config_infra, code_review, exploration
2. task_complexity: one of simple, moderate, complex
3. task_domain: one of backend, frontend, mobile, infra, data, mixed

Respond with ONLY valid JSON, no markdown:
{"task_type": "...", "task_complexity": "...", "task_domain": "..."}"""


class TaskClassifier(BaseAnalyzer):
    name = "b01_task_classification"

    def analyze(self, sessions_df, messages_df, tool_calls_df, tool_results_df, token_usage_df) -> dict:
        # Lazy import to avoid failure if openai not installed
        try:
            from openai import OpenAI
            api_key = os.environ.get("OPENAI_API_KEY")
            if not api_key:
                return {"error": "OPENAI_API_KEY not set", "per_session": []}
            client = OpenAI()
        except ImportError:
            return {"error": "openai package not installed", "per_session": []}

        per_session = []
        stats = {"total": 0, "skipped": 0, "cached": 0, "llm_called": 0, "errors": 0}

        for session in self.iter_sessions(sessions_df):
            sid = session["id"]
            source = session.get("source") or ""
            stats["total"] += 1

            msgs = self.get_session_messages(messages_df, sid)
            msg_list = msgs.to_dict("records") if not msgs.empty else []

            first_prompt = get_first_user_prompt(msg_list, source)

            # Heuristic pre-filters
            if len(first_prompt) <= 5:
                stats["skipped"] += 1
                per_session.append({"session_id": sid, "task_type": "trivial",
                                    "task_complexity": "simple", "task_domain": "mixed"})
                continue

            # Check cache
            cached = get_cached(first_prompt)
            if cached:
                stats["cached"] += 1
                per_session.append({"session_id": sid, **cached})
                continue

            # LLM call
            try:
                resp = client.chat.completions.create(
                    model="gpt-4o-mini", temperature=0,
                    messages=[
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": first_prompt[:4000]},
                    ],
                )
                text = resp.choices[0].message.content.strip()
                result = json.loads(text)

                if result.get("task_type") not in VALID_TASK_TYPES:
                    result["task_type"] = "exploration"
                if result.get("task_complexity") not in VALID_COMPLEXITIES:
                    result["task_complexity"] = "moderate"
                if result.get("task_domain") not in VALID_DOMAINS:
                    result["task_domain"] = "mixed"

                set_cache(first_prompt, result)
                stats["llm_called"] += 1
                per_session.append({"session_id": sid, **result})
            except Exception as e:
                stats["errors"] += 1
                print(f"  [b01] LLM error for {sid[:8]}: {e}")
                per_session.append({"session_id": sid, "task_type": "exploration",
                                    "task_complexity": "moderate", "task_domain": "mixed"})

        print(f"  [b01] Stats: {stats}")
        return {"stats": stats, "total_sessions": len(per_session), "per_session": per_session}
