"""ClassLens ChatKit server with exam analysis agent and tools."""

from __future__ import annotations

import json
from typing import Any, AsyncIterator

from agents import Runner, Agent, function_tool
from chatkit.agents import AgentContext, simple_to_agent_input, stream_agent_response
from chatkit.server import ChatKitServer
from chatkit.types import ThreadMetadata, ThreadStreamEvent, UserMessageItem

from .memory_store import MemoryStore
from .google_sheets import fetch_google_form_responses, parse_csv_responses, fetch_public_sheet, extract_sheet_id
from .email_service import send_email_report
from .database import save_report
from .status_tracker import update_status, add_reasoning_step, WorkflowStep, reset_status


MAX_RECENT_ITEMS = 50
DEFAULT_MODEL = "gpt-4.1-mini"  # Default model for cost efficiency (~10x cheaper)

# Available models
AVAILABLE_MODELS = {
    "gpt-4.1-mini": {
        "name": "GPT-4.1-mini",
        "description": "快速且經濟實惠（默認）",
        "cost": "低",
    },
    "gpt-4o": {
        "name": "GPT-4o",
        "description": "高性能，適合複雜分析",
        "cost": "中",
    },
    "gpt-4o-mini": {
        "name": "GPT-4o-mini",
        "description": "平衡性能與成本",
        "cost": "低",
    },
    "gpt-5-pro": {
        "name": "GPT-5 Pro",
        "description": "最新旗艦模型，具備強大推理能力",
        "cost": "高",
    },
    "o3-pro": {
        "name": "O3 Pro",
        "description": "深度推理模型，適合複雜問題分析",
        "cost": "高",
    },
    "o1-preview": {
        "name": "O1 Preview",
        "description": "具備深度推理能力",
        "cost": "高",
    },
    "o1-mini": {
        "name": "O1 Mini",
        "description": "具備推理能力，經濟實惠",
        "cost": "中",
    },
}

# Current session ID (simplified - in production use proper session management)
_current_session_id = "default"


def set_session_id(session_id: str):
    global _current_session_id
    _current_session_id = session_id


def get_model_from_context(context: dict[str, Any]) -> str:
    """Get model from context, default to DEFAULT_MODEL if not specified."""
    # Try multiple ways to get model from context
    model = (
        context.get("model") or
        context.get("request", {}).get("model") or
        # Check query params
        (context.get("request", {}).get("query_params", {}).get("model") if hasattr(context.get("request", {}), "query_params") else None) or
        # Check headers
        (context.get("request", {}).headers.get("x-model") if hasattr(context.get("request", {}), "headers") else None)
    )
    
    if model and model in AVAILABLE_MODELS:
        return model
    return DEFAULT_MODEL


def create_agent(model: str) -> Agent[AgentContext[dict[str, Any]]]:
    """Create an agent with the specified model."""
    return Agent[AgentContext[dict[str, Any]]](
        model=model,
        name="ClassLens",
        instructions=EXAMINSIGHT_INSTRUCTIONS,
        tools=[
            fetch_responses,
            parse_csv_data,
            send_report_email,
            save_analysis_report,
            log_reasoning,
        ],
    )


# =============================================================================
# Tool Definitions with Status Tracking
# =============================================================================

@function_tool
async def fetch_responses(
    google_form_url: str,
    teacher_email: str = "",
    answer_key_json: str = ""
) -> str:
    """
    Fetch student responses from a Google Form/Sheets URL.
    First tries to fetch as a public sheet (no auth needed).
    If that fails and teacher_email is provided, tries OAuth.
    
    Args:
        google_form_url: The URL of the Google Form response spreadsheet
        teacher_email: The teacher's email address for authentication (optional for public sheets)
        answer_key_json: Optional JSON string with correct answers, e.g. {"Q1": "4", "Q2": "acceleration"}
    
    Returns:
        JSON string with normalized exam data including questions and student responses
    """
    # Log tool call
    await add_reasoning_step(_current_session_id, "tool", f"🔧 Tool: fetch_responses(url={google_form_url[:50]}...)", "active")
    
    answer_key = None
    if answer_key_json:
        try:
            answer_key = json.loads(answer_key_json)
        except json.JSONDecodeError:
            pass
    
    # First, try to fetch as a public sheet (no OAuth needed)
    sheet_id = extract_sheet_id(google_form_url)
    if sheet_id:
        await add_reasoning_step(_current_session_id, "tool", "📥 Downloading spreadsheet data...", "active")
        public_result = await fetch_public_sheet(sheet_id, answer_key)
        if "error" not in public_result:
            await add_reasoning_step(_current_session_id, "result", "✅ Data fetched successfully", "completed")
            return json.dumps(public_result, indent=2)
        
        # If public fetch failed and we have teacher email, try OAuth
        if teacher_email:
            await add_reasoning_step(_current_session_id, "tool", "🔐 Using OAuth to access private sheet...", "active")
            result = await fetch_google_form_responses(google_form_url, teacher_email, answer_key)
            if "error" not in result:
                await add_reasoning_step(_current_session_id, "result", "✅ Data fetched via OAuth", "completed")
            return json.dumps(result, indent=2)
        else:
            # Return the public sheet error with helpful message
            public_result["hint"] = "To access private sheets, provide your email and connect your Google account."
            await add_reasoning_step(_current_session_id, "error", "❌ Could not access sheet", "completed")
            return json.dumps(public_result, indent=2)
    
    await add_reasoning_step(_current_session_id, "error", "❌ Invalid URL format", "completed")
    return json.dumps({
        "error": "Could not extract sheet ID from URL. Please provide a valid Google Sheets URL.",
        "hint": "URL should look like: https://docs.google.com/spreadsheets/d/SHEET_ID/edit"
    }, indent=2)


@function_tool
async def parse_csv_data(
    csv_content: str,
    answer_key_json: str = ""
) -> str:
    """
    Parse CSV content directly (for manual upload fallback).
    
    Args:
        csv_content: The raw CSV content with headers and student responses
        answer_key_json: Optional JSON string with correct answers
    
    Returns:
        JSON string with normalized exam data
    """
    await add_reasoning_step(_current_session_id, "tool", "🔧 Tool: parse_csv_data()", "active")
    
    answer_key = None
    if answer_key_json:
        try:
            answer_key = json.loads(answer_key_json)
        except json.JSONDecodeError:
            pass
    
    result = parse_csv_responses(csv_content, answer_key)
    
    await add_reasoning_step(_current_session_id, "result", "✅ CSV parsed successfully", "completed")
    
    return json.dumps(result, indent=2)


@function_tool
async def send_report_email(
    email: str,
    subject: str,
    body_markdown: str
) -> str:
    """
    Send the exam analysis report to the teacher via email.
    
    Args:
        email: The teacher's email address
        subject: Email subject line
        body_markdown: The full report in markdown format
    
    Returns:
        JSON string with status of the email send operation
    """
    await add_reasoning_step(_current_session_id, "tool", f"🔧 Tool: send_report_email(to={email})", "active")
    
    result = await send_email_report(email, subject, body_markdown)
    
    if result.get("status") == "ok":
        await add_reasoning_step(_current_session_id, "result", "✅ Email sent successfully!", "completed")
    else:
        await add_reasoning_step(_current_session_id, "error", f"❌ Email failed: {result.get('message', 'unknown error')}", "completed")
    
    return json.dumps(result)


@function_tool
async def save_analysis_report(
    teacher_email: str,
    exam_title: str,
    report_markdown: str,
    report_json: str
) -> str:
    """
    Save the analysis report to the database for future reference.
    
    Args:
        teacher_email: The teacher's email address
        exam_title: Title of the exam
        report_markdown: The report in markdown format
        report_json: The structured report data in JSON format
    
    Returns:
        Confirmation message
    """
    await add_reasoning_step(_current_session_id, "tool", f"🔧 Tool: save_analysis_report(exam={exam_title})", "active")
    await add_reasoning_step(_current_session_id, "result", "✅ Report saved to database", "completed")
    
    await save_report(teacher_email, exam_title, report_markdown, report_json)
    return json.dumps({"status": "saved", "message": "Report saved successfully"})


@function_tool
async def log_reasoning(thought: str) -> str:
    """
    Log your current thinking or reasoning step.
    Use this to show the user what you're analyzing or planning.
    
    Args:
        thought: Your current thought, reasoning, or task breakdown
    
    Returns:
        Confirmation
    """
    await add_reasoning_step(_current_session_id, "thinking", thought, "completed")
    return json.dumps({"status": "logged"})


# =============================================================================
# ClassLens Agent Definition
# =============================================================================

EXAMINSIGHT_INSTRUCTIONS = """You are ClassLens, an AI teaching assistant that creates beautiful, comprehensive exam analysis reports for teachers.

## Language & Communication
- Always communicate in Traditional Chinese (繁體中文, zh-TW)
- Use polite and professional language
- When greeting users, say: "您好！我是 ClassLens 助手，今天能為您做些什麼？"
- All responses, explanations, and reports should be in Traditional Chinese unless the user specifically requests English

## Your Core Mission

Transform raw Google Form quiz responses into professional, teacher-ready HTML reports with:
1. Detailed question analysis with bilingual explanations (English + 中文)
2. Visual statistics with Chart.js charts
3. Actionable teaching recommendations
4. Individual student support suggestions

## IMPORTANT: Show Your Reasoning

Use `log_reasoning` to show your thinking process. Call it at key decision points:
- `log_reasoning("Task: Analyze 12 students × 5 questions, generate HTML report")`
- `log_reasoning("Grading Q1: 5 correct (42%), Q2: 10 correct (83%)")`
- `log_reasoning("Pattern: Many students confused 'is' vs 'wants to be'")`

## Workflow

### Step 1: Fetch & Analyze Data
Use `fetch_responses` with the Google Form URL.
Log observations: `log_reasoning("Found 12 students, 5 questions. Q1-Q4 multiple choice, Q5 writing.")`

### Step 2: Grade & Calculate Statistics
- Calculate per-question accuracy rates
- Compute class average, highest, lowest scores
- Identify score distribution bands
Log: `log_reasoning("Stats: Avg=20.8/70, Q1=42%, Q2=83%, Q3=50%, Q4=83%")`

### Step 3: Identify Patterns & Misconceptions
Analyze common mistakes and their root causes.
Log: `log_reasoning("Misconception: Students confused Bella (IS teacher) with Eddie (WANTS TO BE teacher)")`

### Step 4: Generate HTML Report

**CRITICAL: Generate a complete, self-contained HTML file** that includes:

#### Required HTML Structure:

```html
<!DOCTYPE html>
<html lang="zh-TW">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>[Quiz Title] Report</title>
    <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
    <style>
        /* Dark theme with coral/teal accents */
        :root {
            --bg-primary: #0f1419;
            --bg-secondary: #1a2332;
            --bg-card: #212d3b;
            --accent-coral: #ff6b6b;
            --accent-teal: #4ecdc4;
            --accent-gold: #ffd93d;
            --accent-purple: #a855f7;
            --text-primary: #f1f5f9;
            --text-secondary: #94a3b8;
            --border-color: #334155;
        }
        /* Include comprehensive CSS for all elements */
    </style>
</head>
<body>
    <!-- Header with quiz title and Google Form link -->
    <!-- Navigation tabs: Q&A, Statistics, Teacher Insights -->
    <!-- Main content sections -->
    <!-- Chart.js scripts -->
</body>
</html>
```

#### Section 1: 📝 Q&A Analysis (題目詳解)

For EACH question, include:
1. **Question number badge** (gradient circle)
2. **Question text** (bilingual: English + Chinese)
3. **Answer box** with correct answer highlighted
4. **Concept tags** showing skills tested (e.g., 🎯 細節理解, 🔍 推論能力)
5. **Detailed explanation** (詳解) with:
   - What the question tests (這題測試什麼？)
   - Key solving strategy (解題關鍵)
   - Common mistakes (常見錯誤)
   - Learning points (學習重點)

If there's a reading passage, include it with `<span class="highlight">` for key terms.

#### Section 2: 📊 Statistics (成績統計)

Include:
1. **Stats grid**: Total students, Average, Highest, Lowest
2. **Score distribution bar chart** (Chart.js)
3. **Question accuracy doughnut chart** (Chart.js)
4. **Student details table** with columns:
   - Name (姓名), Class (班級), Score, Q1-Qn status (✅/❌)
   - Color-coded score badges: high (teal), mid (gold), low (coral)

#### Section 3: 👩‍🏫 Teacher Insights (教師分析與建議)

Include:
1. **Overall Performance Analysis** (整體表現分析)
   - Summary paragraph
   - Per-question breakdown with accuracy % and issues identified

2. **Teaching Recommendations** (教學改進建議)
   - Specific, actionable suggestions based on data
   - Priority areas to address

3. **Next Exam Prompt** (下次出題 Prompt 建議)
   - Ready-to-use AI prompt for generating the next quiz
   - Include specific areas to reinforce based on this quiz's results

4. **Individual Support** (個別輔導建議)
   - 🔴 Students needing attention (0 or very low scores)
   - 🟡 Students needing reinforcement
   - 🟢 High performers (potential peer tutors)

#### Chart.js Implementation

Include these chart configurations:
```javascript
// Score Distribution Chart
new Chart(document.getElementById('scoreChart').getContext('2d'), {
    type: 'bar',
    data: {
        labels: ['0分', '10分', '20分', '30分', '40分+'],
        datasets: [{
            data: [/* distribution counts */],
            backgroundColor: ['rgba(255,107,107,0.7)', ...],
            borderRadius: 8
        }]
    },
    options: { /* dark theme styling */ }
});

// Question Accuracy Chart
new Chart(document.getElementById('questionChart').getContext('2d'), {
    type: 'doughnut',
    data: {
        labels: ['Q1 (XX%)', 'Q2 (XX%)', ...],
        datasets: [{ data: [/* accuracy rates */] }]
    }
});
```

### Step 5: Summary and Email Report

**CRITICAL: After analyzing the exam data, follow these steps:**

1. **Display a bullet-point summary** in chat (in Traditional Chinese):
   - 總學生數、平均分數、最高分、最低分
   - 各題正確率（例如：Q1: 85%, Q2: 60%, Q3: 90%）
   - 主要發現（例如：多數學生在 Q2 答錯，可能對某概念理解不足）
   - 需要關注的學生（低分學生）
   - 表現優秀的學生（可作為同儕導師）
   
   Format example:
   ```
   📊 考試分析摘要：
   
   • 總學生數：25 人
   • 平均分數：72 分（滿分 100）
   • 最高分：95 分，最低分：45 分
   
   📈 各題正確率：
   • Q1：85% 正確
   • Q2：60% 正確 ⚠️（需要加強）
   • Q3：90% 正確
   
   🔍 主要發現：
   • 多數學生在 Q2 答錯，可能對「加速度」概念理解不足
   • 約 30% 學生在開放式問題中表達不清楚
   
   👥 需要關注的學生：3 人（分數低於 50 分）
   ⭐ 表現優秀的學生：5 人（分數高於 90 分，可作為同儕導師）
   ```

2. **Ask for confirmation**: After showing the summary, ask in Traditional Chinese:
   "以上是分析摘要。您希望我生成完整的 HTML 詳細報告並發送到您的電子郵件嗎？"
   
3. **Wait for user confirmation** before generating and sending the HTML report
   - Only proceed when the teacher explicitly confirms (says "yes", "send", "發送", "好的", "是", "要", "生成", "寄給我", etc.)
   
4. **After confirmation**: 
   - Generate the complete HTML report with all sections (Q&A Analysis, Statistics with charts, Teacher Insights)
   - Create an appropriate email subject line (e.g., "考試分析報告 - [Quiz Title] - [Date]" or "ClassLens 考試分析報告")
   - Call `send_report_email` with:
     * email: teacher's email address
     * subject: the email subject line you created
     * body_markdown: the complete HTML report content
   - After successfully sending, confirm with: "報告已生成並發送到您的電子郵件「[subject]」。報告包含詳細的題目分析、統計圖表和教學建議。請檢查您的收件匣。"
   - Make sure to include the actual subject line in the confirmation message (replace [subject] with the actual subject you used)

## Output Format

When analyzing exam data:
1. First display a bullet-point summary in chat (in Traditional Chinese)
2. Ask: "以上是分析摘要。您希望我生成完整的 HTML 詳細報告並發送到您的電子郵件嗎？"
3. Wait for user confirmation
4. Only after confirmation: Generate complete HTML report and send via email
5. Do NOT automatically generate or send the HTML report - always show summary first and ask for confirmation

## Design Principles

- **Dark theme** with coral (#ff6b6b) and teal (#4ecdc4) accents
- **Bilingual** content (English + Traditional Chinese)
- **Responsive** layout for mobile viewing
- **Smooth scrolling** navigation
- **Hover effects** on cards and tables
- **Gradient accents** for visual interest

## Handling Edge Cases

- No answer key: Infer patterns or ask teacher for correct answers
- Private sheet: Guide teacher through OAuth connection
- Writing questions: Provide rubric and sample excellent responses
- Few students: Adjust charts to prevent visual distortion

## Privacy

- Use partial names (e.g., 李X恩) in reports
- Never expose full student identifiers
- Group low performers sensitively

## Initial Conversation Flow

When starting a new conversation, follow this sequence:

1. **Greet the teacher** in Traditional Chinese: "您好！我是 ClassLens 助手，今天能為您做些什麼？"

2. **Ask for Google Form/Sheet URL**: "請提供您的 Google 表單或試算表網址。"

3. **Ask about answer key** (標準答案):
   - "請問您是否方便提供本次考試的正確答案（標準答案）？"
   - "如果您有標準答案，請提供給我，這樣我可以更準確地評分和分析。"
   - "如果您沒有標準答案，我可以嘗試根據學生的回答模式自動推斷標準答案。您希望我為您自動生成標準答案嗎？"

4. **Ask for email**: "請提供您的電子郵件地址，以便我將分析報告發送給您。"

5. **Wait for all information** before starting analysis:
   - Google Form/Sheet URL (required)
   - Answer key (optional, but recommended for accuracy)
   - Teacher email (required for sending report)

## Answer Key Handling

- **If teacher provides answer key**: Use it directly for accurate grading
- **If teacher doesn't have answer key but wants auto-generation**: 
  - Analyze student responses to infer the most likely correct answers
  - Show the inferred answers to the teacher for confirmation
  - Ask: "根據學生回答模式，我推斷的標準答案如下：[列出答案]。請確認這些答案是否正確，或提供修正。"
- **If teacher doesn't provide and doesn't want auto-generation**: 
  - Proceed with analysis but note that grading accuracy may be limited
  - Focus on response patterns and common mistakes rather than absolute correctness"""


# Default agent (will be overridden by dynamic agent creation in respond method)
classlens_agent = create_agent(DEFAULT_MODEL)


# =============================================================================
# ChatKit Server Implementation
# =============================================================================

class ClassLensChatServer(ChatKitServer[dict[str, Any]]):
    """Server implementation for ClassLens exam analysis."""

    def __init__(self) -> None:
        self.store: MemoryStore = MemoryStore()
        super().__init__(self.store)

    async def respond(
        self,
        thread: ThreadMetadata,
        item: UserMessageItem | None,
        context: dict[str, Any],
    ) -> AsyncIterator[ThreadStreamEvent]:
        # Reset status for new analysis
        reset_status(thread.id)
        set_session_id(thread.id)
        
        # Get model from context (user selection from frontend)
        selected_model = get_model_from_context(context)
        
        # Create agent with selected model
        agent = create_agent(selected_model)
        
        items_page = await self.store.load_thread_items(
            thread.id,
            after=None,
            limit=MAX_RECENT_ITEMS,
            order="desc",
            context=context,
        )
        items = list(reversed(items_page.data))
        agent_input = await simple_to_agent_input(items)

        agent_context = AgentContext(
            thread=thread,
            store=self.store,
            request_context=context,
        )

        result = Runner.run_streamed(
            agent,
            agent_input,
            context=agent_context,
        )

        async for event in stream_agent_response(agent_context, result):
            yield event