"""AI-powered tree generation service.

Implements the 4-stage wizard flow:
  Stage 2 (scaffold): AI suggests 4-7 top-level branches
  Stage 3 (branch_detail): AI generates detailed nodes per branch
  Stage 4 (assemble): Pure assembly logic — zero AI calls

Uses the provider abstraction from ai_provider.py (supports Gemini + Anthropic).
"""
import json
import logging
import re
import uuid
from typing import Any

from app.core.ai_provider import get_ai_provider
from app.core.config import settings
from app.core.ai_tree_validator import validate_generated_tree, count_tree_stats

logger = logging.getLogger(__name__)

# ── Cost estimation (Haiku 4.5 pricing) ──
COST_PER_INPUT_TOKEN = 1.0 / 1_000_000    # $1.00 per 1M input tokens
COST_PER_OUTPUT_TOKEN = 5.0 / 1_000_000   # $5.00 per 1M output tokens


# ── System Prompts ──

SCAFFOLD_SYSTEM_PROMPT = """You are ResolutionFlow AI, assisting MSP engineers to build troubleshooting and procedural flows for IT service management.

Context: Your audience is technical MSP staff experienced with Windows Server, Active Directory, networking, and common MSP tooling (ConnectWise, Datto, SonicWall, etc.).

Task: Given a flow type, category, name, description, and environment tags, suggest 4-7 top-level branches for the flow.

For TROUBLESHOOTING flows:
- Branches should describe the symptom the user observes — written as what the user sees or reports
- The branch name becomes a selectable option on the first screen, so it must be self-identifying from a user's perspective
- Good: "Drive letter missing after login", "Mapped drive shows as disconnected (red X)", "Access denied when opening files"
- Bad: "Authentication Failures", "GPO Issues", "Connectivity Problems" — too vague for users to self-identify
- Order from most common to least common

For PROCEDURE flows:
- Branches should be phase-based stages (e.g., "Prerequisites", "Configuration", "Verification", "Documentation")
- Each branch represents a major step in the process
- Order in logical execution sequence

Rules:
- Suggest 4-7 branches
- Be specific to the technology/service described — avoid generic internal category names
- Branch names should be concise (3-7 words) and describe the observable symptom or phase
- Each branch needs a brief description (1 sentence) explaining what scenarios it covers
- Return ONLY valid JSON, no markdown, no explanation

Output format:
{"branches": [{"name": "Branch Name", "description": "Brief description of what this covers"}]}"""


BRANCH_DETAIL_SYSTEM_PROMPT = """You are ResolutionFlow AI generating step-by-step detail for one branch of a troubleshooting or procedural flow for MSP engineers.

Context: Your audience is technical MSP staff experienced with Windows Server, Active Directory, networking, and common MSP tooling.

You must return ONLY valid JSON — no markdown, no code fences, no explanation.

Required node schema:

Decision nodes (branching diagnostic questions — choose the right path):
{"id": "unique-slug", "type": "decision", "question": "The diagnostic question", "help_text": "Optional context or command hint", "options": [{"id": "opt-id", "label": "Specific observable answer", "next_node_id": "child-node-id"}], "children": [<all child nodes listed here>]}

Action nodes (a single investigation or remediation step — MUST have next_node_id pointing to the next node):
{"id": "unique-slug", "type": "action", "title": "Short title", "description": "Detailed instructions", "commands": ["PowerShell or CMD commands"], "expected_outcome": "What success looks like", "next_node_id": "id-of-next-sibling-node"}

Solution nodes (leaf nodes — the final resolution, no children):
{"id": "unique-slug", "type": "solution", "title": "Resolution title", "description": "Full resolution description", "resolution_steps": ["Step 1", "Step 2"]}

CRITICAL NAVIGATION RULES:
- Decision node: each option's next_node_id MUST exactly match the "id" of a direct child in that decision node's "children" array
- Action node: next_node_id MUST exactly match the "id" of a sibling node (another child of the same parent decision node)
- Every action node MUST have a next_node_id — action nodes with no next step are broken dead ends
- Solution nodes have no children and no next_node_id — they are the terminus
- Every path through the tree MUST end at a solution node

Additional rules:
1. Generate 4-10 nodes total for this branch
2. Start with a decision node if troubleshooting, action node if procedure
3. Decision nodes must have at least 2 options with specific, observable answer choices
4. Include realistic MSP commands (PowerShell preferred for Windows)
5. Use unique node IDs prefixed with the branch context (e.g., "gpo-check-link")
6. Build the tree bottom-up in your head: create solution/leaf nodes first, then build parent nodes referencing their IDs

Few-shot example showing correct action node next_node_id usage:
{"id": "dns-root", "type": "decision", "question": "Can the client resolve any DNS names?", "help_text": "Run: nslookup google.com", "options": [{"id": "dns-opt-none", "label": "No — nslookup times out or returns 'server failed'", "next_node_id": "dns-check-service"}, {"id": "dns-opt-partial", "label": "Some names resolve but others fail", "next_node_id": "dns-check-specific"}], "children": [{"id": "dns-check-service", "type": "action", "title": "Check DNS Client Service", "description": "Verify the DNS Client service is running on the affected machine", "commands": ["Get-Service -Name Dnscache | Select-Object Status,StartType"], "expected_outcome": "Status should be Running", "next_node_id": "dns-service-solution"}, {"id": "dns-service-solution", "type": "solution", "title": "DNS Service Was Stopped", "description": "The DNS Client service was stopped, preventing all name resolution", "resolution_steps": ["Run: Start-Service Dnscache", "Set startup type: Set-Service Dnscache -StartupType Automatic", "Flush cache: ipconfig /flushdns", "Test: nslookup google.com"]}, {"id": "dns-check-specific", "type": "solution", "title": "Selective DNS Failure — Stale or Missing Records", "description": "Some records resolve correctly, indicating DNS is functional but specific records are stale or missing", "resolution_steps": ["Check DNS server for missing A/CNAME records", "Clear DNS cache on the DNS server: Clear-DnsServerCache", "Flush client cache: ipconfig /flushdns", "Verify with: nslookup <failing-hostname>"]}]}"""


CORRECTIVE_PROMPT_TEMPLATE = """Your previous JSON was invalid for ResolutionFlow's tree schema.

Validation errors:
{error_list}

CRITICAL RULES TO FIX THESE ERRORS:

1. Decision node options → next_node_id MUST match the "id" of a direct child in that SAME decision node's "children" array.
   Example: if decision node has children [A, B, C], then option next_node_id must be "A", "B", or "C".

2. Action node → next_node_id MUST match the "id" of a SIBLING node — another child of the SAME parent decision node.
   Example: if parent decision has children [action-1, solution-1, solution-2], then action-1's next_node_id must be "solution-1" or "solution-2".
   The next node must ALREADY EXIST in the parent's children array — do NOT nest the next node inside the action node.

3. Every referenced ID must physically exist somewhere in the tree as a node with that exact "id" value.

Return a corrected full JSON object only. No markdown, no prose, no code fences.
Fix ALL listed errors while maintaining the same troubleshooting/procedural logic."""


def _strip_markdown_fences(text: str) -> str:
    """Strip markdown code fences if the model wrapped its JSON response."""
    text = text.strip()
    match = re.match(r"^```(?:json)?\s*([\s\S]*?)```$", text)
    if match:
        return match.group(1).strip()
    return text



def _estimate_cost(input_tokens: int, output_tokens: int) -> float:
    """Estimate USD cost from token counts."""
    return (input_tokens * COST_PER_INPUT_TOKEN) + (
        output_tokens * COST_PER_OUTPUT_TOKEN
    )


async def scaffold_branches(
    wizard_state: dict[str, Any],
) -> tuple[list[dict[str, str]], int, int, float]:
    """Stage 2: AI suggests top-level branches.

    Returns (branches, input_tokens, output_tokens, estimated_cost).
    Raises ValueError on invalid response.
    """
    provider = get_ai_provider()

    flow_type = wizard_state.get("flow_type", "troubleshooting")
    name = wizard_state.get("name", "")
    description = wizard_state.get("description", "")
    tags = wizard_state.get("environment_tags", [])

    user_message = (
        f"Flow type: {flow_type}\n"
        f"Name: {name}\n"
        f"Description: {description}\n"
    )
    if tags:
        user_message += f"Environment: {', '.join(tags)}\n"

    raw_text, input_tokens, output_tokens = await provider.generate_json(
        system_prompt=SCAFFOLD_SYSTEM_PROMPT,
        messages=[{"role": "user", "content": user_message}],
        max_tokens=2048,
    )

    logger.info(
        "scaffold raw response (tokens in=%d out=%d, len=%d): %s",
        input_tokens,
        output_tokens,
        len(raw_text),
        raw_text[:500],
    )
    raw_text = _strip_markdown_fences(raw_text)
    cost = _estimate_cost(input_tokens, output_tokens)

    try:
        data = json.loads(raw_text)
    except json.JSONDecodeError as e:
        logger.error("scaffold JSON parse failed. Full text (%d chars): %s", len(raw_text), raw_text)
        raise ValueError(f"AI returned invalid JSON: {e}")

    branches = data.get("branches", [])
    if not isinstance(branches, list) or len(branches) < 2:
        raise ValueError("AI returned fewer than 2 branches")

    return branches, input_tokens, output_tokens, cost


async def generate_branch_detail(
    wizard_state: dict[str, Any],
    branch_name: str,
    existing_branches: list[str],
) -> tuple[dict[str, Any], int, int, float]:
    """Stage 3: AI generates detailed nodes for one branch.

    Returns (branch_tree, input_tokens, output_tokens, estimated_cost).
    On validation failure, retries once with corrective prompt.
    Raises ValueError if both attempts fail.
    """
    provider = get_ai_provider()

    flow_type = wizard_state.get("flow_type", "troubleshooting")
    name = wizard_state.get("name", "")
    description = wizard_state.get("description", "")

    user_message = (
        f"Flow: {name} ({flow_type})\n"
        f"Description: {description}\n"
        f"Branch to detail: {branch_name}\n"
    )
    if existing_branches:
        other = [b for b in existing_branches if b != branch_name]
        if other:
            user_message += f"Other branches (avoid overlap): {', '.join(other)}\n"

    messages = [{"role": "user", "content": user_message}]
    total_input = 0
    total_output = 0

    for attempt in range(3):
        raw_text, input_tokens, output_tokens = await provider.generate_json(
            system_prompt=BRANCH_DETAIL_SYSTEM_PROMPT,
            messages=messages,
            max_tokens=8192,
        )

        total_input += input_tokens
        total_output += output_tokens
        logger.debug(
            "branch_detail attempt=%d output_tokens=%d",
            attempt,
            output_tokens,
        )
        raw_text = _strip_markdown_fences(raw_text) if raw_text else ""
        if not raw_text:
            logger.warning("branch_detail attempt=%d returned empty text", attempt)

        try:
            branch_tree = json.loads(raw_text)
        except json.JSONDecodeError as e:
            logger.error(
                "branch_detail attempt=%d JSON parse failed (%d chars): %s",
                attempt, len(raw_text), raw_text[:500],
            )
            if attempt < 2:
                messages.append({"role": "assistant", "content": raw_text})
                messages.append({
                    "role": "user",
                    "content": CORRECTIVE_PROMPT_TEMPLATE.format(
                        error_list=f"JSON parse error: {e}"
                    ),
                })
                continue
            raise ValueError(f"AI returned invalid JSON after retry: {e}")

        errors = validate_generated_tree(branch_tree)
        if not errors:
            cost = _estimate_cost(total_input, total_output)
            return branch_tree, total_input, total_output, cost

        if attempt < 2:
            messages.append({"role": "assistant", "content": raw_text})
            messages.append({
                "role": "user",
                "content": CORRECTIVE_PROMPT_TEMPLATE.format(
                    error_list="\n".join(f"- {e}" for e in errors)
                ),
            })
            continue

        raise ValueError(
            f"AI tree validation failed after retry: {'; '.join(errors)}"
        )

    # Should not reach here
    raise ValueError("Branch detail generation failed")


def assemble_tree(
    wizard_state: dict[str, Any],
    branches: list[dict[str, Any]],
) -> tuple[dict[str, Any], str, str, dict[str, int]]:
    """Stage 4: Assemble branches into a complete tree.

    Zero AI calls — pure assembly logic.
    Returns (tree_structure, suggested_name, suggested_description, summary_stats).
    """
    flow_type = wizard_state.get("flow_type", "troubleshooting")
    name = wizard_state.get("name", "Untitled Flow")
    description = wizard_state.get("description", "")

    # Build root decision node pointing to each branch
    options = []
    children = []
    for i, branch in enumerate(branches):
        branch_name = branch.get("name", f"Branch {i + 1}")
        branch_tree = branch.get("steps")

        if not branch_tree or not isinstance(branch_tree, dict):
            # Skip branches without detail
            continue

        branch_id = branch_tree.get("id", f"branch_{i}")
        options.append({
            "id": f"opt_{i + 1}",
            "label": branch_name,
            "next_node_id": branch_id,
        })
        children.append(branch_tree)

    if len(options) < 2:
        raise ValueError("Need at least 2 branches with detail to assemble a tree")

    # Determine root question based on flow type
    if flow_type == "troubleshooting":
        root_question = f"What is the user experiencing? Select the symptom that best matches their report."
        root_help = "Choose the option that most closely describes the user's reported problem."
    else:
        root_question = f"Which phase of {name} are you working on?"
        root_help = None

    tree_structure = {
        "id": "root",
        "type": "decision",
        "question": root_question,
        **({"help_text": root_help} if root_help else {}),
        "options": options,
        "children": children,
    }

    stats = count_tree_stats(tree_structure)

    return tree_structure, name, description, stats
