use axum::{
    extract::{Path, State},
    http::StatusCode,
    routing::post,
    Json, Router,
};
use chrono::Utc;
use std::collections::{HashMap, HashSet};
use uuid::Uuid;

use crate::types::{ApiResponse, Edge, Node, NodePosition, NodeType};

use crate::AppState;

type ApiResult<T> = Result<Json<ApiResponse<T>>, (StatusCode, Json<ApiResponse<T>>)>;

const API_KEY_MESSAGE: &str =
    "AI features require API key. Configure via --api-key, ~/.config/theseus/config.toml, or ./theseus.toml";

pub(crate) fn router() -> Router<AppState> {
    Router::new()
        .route("/api/v1/projects/:project_id/tell", post(tell))
        .route("/api/v1/projects/:project_id/ask", post(ask))
}

#[derive(Debug, serde::Deserialize)]
struct ChatRequest {
    message: String,
}

#[derive(Debug, serde::Deserialize)]
struct AiTellResponse {
    #[serde(default)]
    nodes: Vec<AiNodeAction>,
    #[serde(default)]
    edges: Vec<AiEdgeAction>,
    #[serde(default)]
    summary: String,
}

#[derive(Debug, serde::Deserialize)]
struct AiNodeAction {
    label: String,
    #[serde(default)]
    node_type: Option<String>,
    #[serde(default)]
    content: Option<String>,
    #[serde(default)]
    verdict: Option<String>,
    #[serde(default)]
    evidence: Option<String>,
}

#[derive(Debug, serde::Deserialize)]
struct AiEdgeAction {
    source_ref: String,
    target_ref: String,
    #[serde(default)]
    label: Option<String>,
}

#[derive(Debug, serde::Serialize)]
struct TellOutput {
    summary: String,
    nodes: Vec<Node>,
    edges: Vec<Edge>,
    skipped_edges: Vec<SkippedEdge>,
}

#[derive(Debug, serde::Serialize)]
struct SkippedEdge {
    source_ref: String,
    target_ref: String,
    label: Option<String>,
    reason: String,
}

#[derive(Debug, serde::Deserialize)]
struct AiAskResponse {
    answer: String,
    #[serde(default)]
    relevant_node_ids: Vec<String>,
}

#[derive(Debug, serde::Serialize)]
struct AskOutput {
    answer: String,
    relevant_node_ids: Vec<String>,
}

async fn tell(
    State(state): State<AppState>,
    Path(project_id): Path<String>,
    Json(req): Json<ChatRequest>,
) -> ApiResult<TellOutput> {
    let api_key = require_api_key::<TellOutput>(&state)?;
    let pid = crate::resolve_project_id(&state.db, &project_id)
        .await
        .map_err(|s| (s, Json(ApiResponse::err("project not found"))))?;
    let existing_nodes = state.db.list_nodes(pid).map_err(|_| server_err())?;
    let existing_edges = state.db.list_edges(pid).map_err(|_| server_err())?;
    let context = build_tell_context(&req.message, &existing_nodes, &existing_edges);

    let content = crate::deepseek_chat(&state.http_client, api_key, tell_system_prompt(), &context, 0.2, 4000, true)
        .await
        .map_err(|err| {
            tracing::error!("DeepSeek tell action generation failed: {}", err);
            server_err()
        })?;
    let parsed = parse_json_object(&content).ok_or_else(|| {
        tracing::error!("DeepSeek tell response was not a JSON object: {}", content);
        server_err()
    })?;
    let plan: AiTellResponse = serde_json::from_value(parsed).map_err(|err| {
        tracing::error!("DeepSeek tell response schema mismatch: {}", err);
        server_err()
    })?;

    let base_count = existing_nodes.len();
    let mut created_nodes = Vec::new();
    let mut ref_to_id = HashMap::new();
    let existing_ids: HashSet<Uuid> = existing_nodes.iter().map(|node| node.id).collect();

    for (index, action) in plan.nodes.iter().enumerate() {
        let label = action.label.trim();
        if label.is_empty() {
            continue;
        }
        let node = Node {
            id: Uuid::new_v4(),
            project_id: pid,
            label: label.to_string(),
            node_type: parse_node_type(action.node_type.as_deref()),
            position: Some(position_for(base_count + index)),
            properties: serde_json::json!({
                "content": action.content.as_deref().unwrap_or(""),
                "verdict": normalize_verdict(action.verdict.as_deref()),
                "evidence": normalize_evidence(action.evidence.as_deref()),
            }),
            created_at: Utc::now(),
        };
        state.db.insert_node(&node).map_err(|err| {
            tracing::error!("Failed to insert AI-created node: {}", err);
            server_err()
        })?;
        ref_to_id.insert(format!("new:{index}"), node.id);
        created_nodes.push(node);
    }

    let mut created_edges = Vec::new();
    let mut skipped_edges = Vec::new();
    for action in &plan.edges {
        let Some(source_id) = resolve_ai_ref(&action.source_ref, &ref_to_id, &existing_ids) else {
            skipped_edges.push(skipped_edge(action, "source node not found"));
            continue;
        };
        let Some(target_id) = resolve_ai_ref(&action.target_ref, &ref_to_id, &existing_ids) else {
            skipped_edges.push(skipped_edge(action, "target node not found"));
            continue;
        };
        if source_id == target_id {
            skipped_edges.push(skipped_edge(action, "self-loop skipped"));
            continue;
        }
        let edge = Edge {
            id: Uuid::new_v4(),
            project_id: pid,
            source_node_id: source_id,
            target_node_id: target_id,
            label: action.label.clone(),
            created_at: Utc::now(),
        };
        state.db.insert_edge(&edge).map_err(|err| {
            tracing::error!("Failed to insert AI-created edge: {}", err);
            server_err()
        })?;
        created_edges.push(edge);
    }

    if !created_nodes.is_empty() || !created_edges.is_empty() {
        let _ = state.tx.send(format!("graph:{pid}"));
    }

    Ok(Json(ApiResponse::ok(TellOutput {
        summary: plan.summary,
        nodes: created_nodes,
        edges: created_edges,
        skipped_edges,
    })))
}

async fn ask(
    State(state): State<AppState>,
    Path(project_id): Path<String>,
    Json(req): Json<ChatRequest>,
) -> ApiResult<AskOutput> {
    let api_key = require_api_key::<AskOutput>(&state)?;
    let pid = crate::resolve_project_id(&state.db, &project_id)
        .await
        .map_err(|s| (s, Json(ApiResponse::err("project not found"))))?;
    let nodes = state.db.list_nodes(pid).map_err(|_| server_err())?;
    let edges = state.db.list_edges(pid).map_err(|_| server_err())?;
    let context = build_ask_context(&req.message, &nodes, &edges);

    let content = crate::deepseek_chat(&state.http_client, api_key, ask_system_prompt(), &context, 0.2, 4000, true)
        .await
        .map_err(|err| {
            tracing::error!("DeepSeek graph question failed: {}", err);
            server_err()
        })?;
    let parsed = parse_json_object(&content).ok_or_else(|| {
        tracing::error!("DeepSeek ask response was not a JSON object: {}", content);
        server_err()
    })?;
    let answer: AiAskResponse = serde_json::from_value(parsed).map_err(|err| {
        tracing::error!("DeepSeek ask response schema mismatch: {}", err);
        server_err()
    })?;

    Ok(Json(ApiResponse::ok(AskOutput {
        answer: answer.answer,
        relevant_node_ids: answer.relevant_node_ids,
    })))
}

fn require_api_key<T: serde::Serialize>(
    state: &AppState,
) -> Result<&str, (StatusCode, Json<ApiResponse<T>>)> {
    state.config.deepseek_api_key.as_deref().ok_or_else(|| {
        (
            StatusCode::BAD_REQUEST,
            Json(ApiResponse::err(API_KEY_MESSAGE)),
        )
    })
}

fn server_err<T: serde::Serialize>() -> (StatusCode, Json<ApiResponse<T>>) {
    crate::error_response(crate::AppError::internal(
        "ai.operation_failed",
        "chat.ai",
        "AI chat operation failed.",
        "AI 对话操作失败。",
    ))
}

fn tell_system_prompt() -> &'static str {
    r#"你是 Theseus 攻击路径图谱的自然语言录入助手。你的任务是把渗透测试人员的中文自然语言记录转换为图谱变更动作。

必须只返回一个 JSON 对象，不要输出 Markdown、解释或代码块。

## 判断流程

### 第一步：判断是否值得创建节点
分析用户输入是否包含以下任一实质性安全发现：
- 具体的漏洞/弱点（有 URL、端口、服务名、漏洞类型、CVE 编号等）
- 具体的攻击操作（有利用了什么、做了什么、结果如何）
- 具体的情报（有版本号、配置、协议、指纹等可验证的事实）
- 具体的凭据/权限（有用户名、密码、Token、密钥等）

以下输入类型 **不值得创建节点**，必须返回空动作：
- 模糊输入："发现了一个东西"、"嗯"、"有个漏洞"（无具体内容）
- 意图/计划："我打算..."、"明天试试..."、"换个方向"
- 元认知："今天先这样"、"卡住了"、"好难打"
- 纯感叹/情绪："这个系统好难"、"有点意思"

**不值得时返回：**
```json
{"nodes": [], "edges": [], "summary": "已记录，但未发现实质性安全发现，不创建图谱节点。"}
```

### 第二步：创建节点和边
如果输入包含实质性发现，返回：
```json
{
  "nodes": [{"label": "...", "node_type": "...", "content": "...", "verdict": "...", "evidence": "..."}],
  "edges": [{"source_ref": "new:0|existing:<uuid>", "target_ref": "new:0|existing:<uuid>", "label": "..."}],
  "summary": "一句中文总结"
}
```

## 节点类型选择标准（必须精确匹配语义）

| node_type | 何时使用 | 示例 |
|-----------|---------|------|
| host | 目标主机/服务器 | "192.168.1.1 是一台 Ubuntu 服务器" |
| service | 运行中的服务/端口 | "42101 端口运行 Nacos 2.2.3" |
| vulnerability | 通用漏洞（注入、XSS、SSRF 等） | "存在 SQL 注入"、"发现 SSRF" |
| privilege_escalation | 提权漏洞（内核漏洞、SUID、配置错误提权） | "CVE-2023-0386 OverlayFS 可提权"、"SUID vim 可提权到 root" |
| credential | 凭据/密钥 | "发现默认密码 root:root"、"找到 AWS AK/SK" |
| finding | 通用安全发现（不属于上述类型） | "WAF 拦截了 union select" |
| entry_point | 攻击入口点 | "外网暴露的管理后台" |
| attack_step | 攻击利用步骤/动作 | "利用 SSRF 读取内网服务"、"通过 SQL 注入获取数据库" |
| impact | 攻击影响/结果 | "获取了 root 权限"、"拿到了数据库全部数据" |
| objective | 渗透目标 | "拿下域控"、"获取数据库中的客户数据" |
| note | 仅当以上类型都不匹配时 | |

**关键：提权相关的漏洞必须用 privilege_escalation，不要用 vulnerability。**

## Evidence 判断标准

- **confirmed**：用户以陈述句描述已观察到的事实。包括：观察到的指纹、测试过的行为、确认存在的配置、获取到的凭据。只要用户在描述"看到了什么、测试了什么"，就是 confirmed。
- **unverified**：提到了某个发现但尚未验证或测试。例如："可能有漏洞"、"看起来像是"。
- **hypothesis**：推测、假设、计划中的攻击思路。例如："可能可用于绕过"、"猜测可以提权"。

**一段话中可能混合多种 evidence：** 事实部分标 confirmed，推测部分单独标 hypothesis。不要把确认的事实降级。

## 边规则
- new:N 引用本次 nodes 数组中的第 N 个节点，existing:<node_id> 引用上下文中已有节点的真实 UUID。
- 尽量连接到最相关的已有节点；没有明确关系时不连线。
- label 应简短描述关系或动作，例如"发现服务"、"验证漏洞"、"利用得到"。

## Content 规则
使用信息丰富的 Markdown，保留用户消息中的目标、端点、命令、响应、证据和推理细节，不要编造事实。"#
}

fn ask_system_prompt() -> &'static str {
    r#"你是 Theseus 攻击路径图谱的中文问答助手，服务于渗透测试复盘和下一步决策。

必须只返回一个 JSON 对象，不要输出 Markdown 代码块。JSON 格式：
{
  "answer": "中文自然语言回答",
  "relevant_node_ids": ["uuid1", "uuid2"]
}

规则：
- 只能基于用户提供的图谱节点、边和节点完整内容回答，不要编造图谱中没有的事实。
- 回答要引用具体节点标签和 UUID，说明依据来自哪些节点或边。
- 如果图谱信息不足，明确说明缺少什么信息，并给出可执行的下一步检查方向。
- relevant_node_ids 只包含和回答直接相关、且确实存在于图谱上下文中的节点 UUID。"#
}

fn build_tell_context(message: &str, nodes: &[Node], edges: &[Edge]) -> String {
    let graph = serde_json::json!({
        "user_message": message,
        "existing_nodes": nodes.iter().map(|node| {
            serde_json::json!({
                "id": node.id.to_string(),
                "label": node.label,
                "node_type": node.node_type.to_string(),
                "verdict": node.properties.get("verdict").and_then(|value| value.as_str()).unwrap_or("pending"),
                "evidence": node.properties.get("evidence").and_then(|value| value.as_str()).unwrap_or("unverified"),
                "content": truncate(node.properties.get("content").and_then(|value| value.as_str()).unwrap_or(""), 200),
            })
        }).collect::<Vec<_>>(),
        "existing_edges": edges.iter().map(|edge| {
            serde_json::json!({
                "id": edge.id.to_string(),
                "source_node_id": edge.source_node_id.to_string(),
                "target_node_id": edge.target_node_id.to_string(),
                "label": edge.label.as_deref().unwrap_or(""),
            })
        }).collect::<Vec<_>>(),
    });
    serde_json::to_string_pretty(&graph).unwrap_or_else(|_| "{}".to_string())
}

fn build_ask_context(message: &str, nodes: &[Node], edges: &[Edge]) -> String {
    let graph = serde_json::json!({
        "question": message,
        "nodes": nodes.iter().map(|node| {
            serde_json::json!({
                "id": node.id.to_string(),
                "label": node.label,
                "node_type": node.node_type.to_string(),
                "verdict": node.properties.get("verdict").and_then(|value| value.as_str()).unwrap_or("pending"),
                "evidence": node.properties.get("evidence").and_then(|value| value.as_str()).unwrap_or("unverified"),
                "content": node.properties.get("content").and_then(|value| value.as_str()).unwrap_or(""),
                "properties": node.properties,
            })
        }).collect::<Vec<_>>(),
        "edges": edges.iter().map(|edge| {
            serde_json::json!({
                "id": edge.id.to_string(),
                "source_node_id": edge.source_node_id.to_string(),
                "target_node_id": edge.target_node_id.to_string(),
                "label": edge.label.as_deref().unwrap_or(""),
            })
        }).collect::<Vec<_>>(),
    });
    serde_json::to_string_pretty(&graph).unwrap_or_else(|_| "{}".to_string())
}

fn parse_node_type(raw: Option<&str>) -> NodeType {
    match raw.unwrap_or("note").trim().to_ascii_lowercase().as_str() {
        "host" => NodeType::Host,
        "service" => NodeType::Service,
        "vulnerability" => NodeType::Vulnerability,
        "privilege_escalation" | "privesc" => NodeType::PrivilegeEscalation,
        "credential" => NodeType::Credential,
        "finding" => NodeType::Finding,
        "note" => NodeType::Note,
        "entry_point" => NodeType::EntryPoint,
        "attack_step" => NodeType::AttackStep,
        "impact" => NodeType::Impact,
        "objective" => NodeType::Objective,
        _ => NodeType::Note,
    }
}

fn normalize_verdict(raw: Option<&str>) -> &'static str {
    match raw.unwrap_or("pending").trim() {
        "success" => "success",
        "failed" => "failed",
        "in_progress" => "in_progress",
        _ => "pending",
    }
}

fn normalize_evidence(raw: Option<&str>) -> &'static str {
    match raw.unwrap_or("unverified").trim() {
        "confirmed" => "confirmed",
        "hypothesis" => "hypothesis",
        _ => "unverified",
    }
}

fn position_for(index: usize) -> NodePosition {
    let column = (index % 8) as f64;
    let row = (index / 8) as f64;
    NodePosition {
        x: 100.0 + column * 90.0,
        y: 100.0 + row * 90.0,
    }
}

fn resolve_ai_ref(
    raw: &str,
    new_refs: &HashMap<String, Uuid>,
    existing_ids: &HashSet<Uuid>,
) -> Option<Uuid> {
    if let Some(id) = new_refs.get(raw.trim()) {
        return Some(*id);
    }
    let existing_raw = raw.trim().strip_prefix("existing:")?;
    let id = Uuid::parse_str(existing_raw).ok()?;
    existing_ids.contains(&id).then_some(id)
}

fn skipped_edge(action: &AiEdgeAction, reason: &str) -> SkippedEdge {
    SkippedEdge {
        source_ref: action.source_ref.clone(),
        target_ref: action.target_ref.clone(),
        label: action.label.clone(),
        reason: reason.to_string(),
    }
}

fn truncate(value: &str, max_chars: usize) -> String {
    let mut out = String::new();
    for ch in value.chars().take(max_chars) {
        out.push(ch);
    }
    if value.chars().count() > max_chars {
        out.push_str("...");
    }
    out
}

fn parse_json_object(content: &str) -> Option<serde_json::Value> {
    if let Ok(value) = serde_json::from_str::<serde_json::Value>(content.trim()) {
        if value.is_object() {
            return Some(value);
        }
    }
    let without_fence = content
        .trim()
        .trim_start_matches("```json")
        .trim_start_matches("```")
        .trim_end_matches("```")
        .trim();
    if let Ok(value) = serde_json::from_str::<serde_json::Value>(without_fence) {
        if value.is_object() {
            return Some(value);
        }
    }
    let start = content.find('{')?;
    let end = content.rfind('}')?;
    serde_json::from_str::<serde_json::Value>(&content[start..=end]).ok()
}
