use std::collections::HashSet;

use regex::Regex;
use serde_json::{json, Value};
use url::Url;

use crate::{
    auth::codex_oauth::{chatgpt_plan_from_id_token, ChatGptPlan},
    credentials::{
        load_codex_tokens, load_provider_api_key, CodexTokens, OsCredentialStore,
        WebSearchCredential,
    },
    model::openai::auth::{refresh_codex_token, CodexAuthSource},
    tool::ToolError,
};

use super::{SearchBackendConfig, SearchItem};
use crate::tools::web::util::http_client;

const OPENAI_RESPONSES_URL: &str = "https://api.openai.com/v1/responses";
const CODEX_RESPONSES_URL: &str = "https://chatgpt.com/backend-api/codex/responses";
const OPENAI_SEARCH_MODEL: &str = "gpt-5.6-luna";
const CODEX_SEARCH_MODEL: &str = "gpt-5.6-terra";

#[derive(Clone, Debug)]
enum OpenAiSearchAuth {
    Codex {
        tokens: CodexTokens,
        source: CodexAuthSource,
    },
    ApiKey(String),
}

impl OpenAiSearchAuth {
    fn endpoint(&self) -> &'static str {
        match self {
            Self::Codex { .. } => CODEX_RESPONSES_URL,
            Self::ApiKey(_) => OPENAI_RESPONSES_URL,
        }
    }

    fn model(&self) -> &'static str {
        match self {
            Self::Codex { tokens, .. } => match tokens
                .id_token
                .as_deref()
                .map(chatgpt_plan_from_id_token)
                .unwrap_or(ChatGptPlan::Unknown)
            {
                ChatGptPlan::Free | ChatGptPlan::Go | ChatGptPlan::Unknown => CODEX_SEARCH_MODEL,
                ChatGptPlan::Plus
                | ChatGptPlan::Pro
                | ChatGptPlan::ProLite
                | ChatGptPlan::Team
                | ChatGptPlan::SelfServeBusinessUsageBased
                | ChatGptPlan::Business
                | ChatGptPlan::EnterpriseCbpUsageBased
                | ChatGptPlan::Enterprise
                | ChatGptPlan::Edu => OPENAI_SEARCH_MODEL,
            },
            Self::ApiKey(_) => OPENAI_SEARCH_MODEL,
        }
    }

    fn bearer_token(&self) -> &str {
        match self {
            Self::Codex { tokens, .. } => &tokens.access_token,
            Self::ApiKey(key) => key,
        }
    }
}

pub(super) fn is_available(config: &SearchBackendConfig) -> bool {
    resolve_auth(config).is_ok()
}

fn resolve_auth(config: &SearchBackendConfig) -> Result<OpenAiSearchAuth, ToolError> {
    if let Ok(access_token) = std::env::var("CODEX_ACCESS_TOKEN") {
        return Ok(OpenAiSearchAuth::Codex {
            tokens: CodexTokens {
                access_token,
                refresh_token: None,
                id_token: None,
                account_id: std::env::var("CODEX_ACCOUNT_ID").ok(),
            },
            source: CodexAuthSource::Env,
        });
    }
    if let Ok(Some(tokens)) = load_codex_tokens(&OsCredentialStore) {
        return Ok(OpenAiSearchAuth::Codex {
            tokens,
            source: CodexAuthSource::Store,
        });
    }
    if let Ok(key) = std::env::var("OPENAI_API_KEY") {
        return Ok(OpenAiSearchAuth::ApiKey(key));
    }
    if let Some(key) = config.credential(WebSearchCredential::OpenAi) {
        return Ok(OpenAiSearchAuth::ApiKey(key));
    }
    if let Ok(Some(key)) = load_provider_api_key(&OsCredentialStore, "openai") {
        return Ok(OpenAiSearchAuth::ApiKey(key));
    }
    Err(ToolError::Message(
        "OpenAI web search unavailable: sign in with /login openai-codex, /login openai, or set OPENAI_API_KEY"
            .into(),
    ))
}

pub(super) async fn search(
    query: &str,
    num_results: usize,
    recency_filter: Option<&str>,
    domain_filter: Option<&[String]>,
    config: &SearchBackendConfig,
) -> Result<Vec<SearchItem>, ToolError> {
    let mut auth = resolve_auth(config)?;
    let body = openai_search_body(&auth, query, num_results, recency_filter, domain_filter);
    let (mut status, mut text) = send_openai_search_request(&auth, &body).await?;
    if status == reqwest::StatusCode::UNAUTHORIZED {
        if let OpenAiSearchAuth::Codex {
            tokens,
            source: CodexAuthSource::Store,
        } = &auth
        {
            if let Some(refresh_token) = tokens.refresh_token.as_deref() {
                let client = http_client();
                let store = OsCredentialStore;
                let refreshed = refresh_codex_token(
                    &client,
                    &store,
                    refresh_token,
                    CodexAuthSource::Store,
                    tokens,
                )
                .await
                .map_err(|err| {
                    ToolError::Message(format!("OpenAI web search token refresh failed: {err}"))
                })?;
                auth = OpenAiSearchAuth::Codex {
                    tokens: refreshed,
                    source: CodexAuthSource::Store,
                };
                let body =
                    openai_search_body(&auth, query, num_results, recency_filter, domain_filter);
                let retried = send_openai_search_request(&auth, &body).await?;
                status = retried.0;
                text = retried.1;
            }
        }
    }
    if !status.is_success() {
        return Err(ToolError::Message(format!(
            "OpenAI web search failed: HTTP {status}: {}",
            text.chars().take(300).collect::<String>()
        )));
    }

    let output = parse_openai_search_output(&text)?;
    let results = extract_openai_search_results(&output, num_results);
    if results.is_empty() {
        Err(ToolError::Message(
            "OpenAI web_search returned no sources".into(),
        ))
    } else {
        Ok(results)
    }
}

fn openai_search_body(
    auth: &OpenAiSearchAuth,
    query: &str,
    num_results: usize,
    recency_filter: Option<&str>,
    domain_filter: Option<&[String]>,
) -> Value {
    json!({
        "model": auth.model(),
        "instructions": openai_search_instructions(num_results, recency_filter, domain_filter),
        "input": [{"role": "user", "content": [{"type": "input_text", "text": query}]}],
        "tools": [openai_web_search_tool(domain_filter)],
        "include": ["web_search_call.action.sources"],
        "store": false,
        "stream": true,
        "tool_choice": "required",
        "parallel_tool_calls": true,
    })
}

async fn send_openai_search_request(
    auth: &OpenAiSearchAuth,
    body: &Value,
) -> Result<(reqwest::StatusCode, String), ToolError> {
    let mut request = http_client()
        .post(auth.endpoint())
        .bearer_auth(auth.bearer_token())
        .header("Content-Type", "application/json");
    if let OpenAiSearchAuth::Codex { tokens, .. } = auth {
        request = request
            .header("User-Agent", "codex-cli")
            .header("originator", "codex_cli_rs");
        if let Some(account_id) = tokens.account_id.as_deref() {
            request = request.header("ChatGPT-Account-ID", account_id);
        }
    }

    let response =
        request.json(body).send().await.map_err(|err| {
            ToolError::Message(format!("OpenAI web search request failed: {err}"))
        })?;
    let status = response.status();
    let text = response
        .text()
        .await
        .map_err(|err| ToolError::Message(format!("OpenAI web search response failed: {err}")))?;
    Ok((status, text))
}

fn openai_search_instructions(
    num_results: usize,
    recency_filter: Option<&str>,
    domain_filter: Option<&[String]>,
) -> String {
    let mut lines = vec![
        "Search the web and return concise source-backed results.".to_string(),
        "Prefer clickable source citations when possible.".to_string(),
        format!("Prefer around {} distinct sources.", num_results.min(20)),
    ];
    if let Some(recency) = recency_filter.and_then(openai_recency_label) {
        lines.push(format!("Prefer sources from the {recency}."));
    }
    let filters = normalize_domain_filters(domain_filter);
    if !filters.allowed.is_empty() {
        lines.push(format!(
            "Only use sources from: {}.",
            filters.allowed.join(", ")
        ));
    }
    if !filters.blocked.is_empty() {
        lines.push(format!(
            "Do not use sources from: {}.",
            filters.blocked.join(", ")
        ));
    }
    lines.join(" ")
}

pub(super) fn openai_recency_label(recency_filter: &str) -> Option<&'static str> {
    match recency_filter {
        "day" => Some("past 24 hours"),
        "week" => Some("past week"),
        "month" => Some("past month"),
        "year" => Some("past year"),
        _ => None,
    }
}

fn openai_web_search_tool(domain_filter: Option<&[String]>) -> Value {
    let filters = normalize_domain_filters(domain_filter);
    let mut tool = serde_json::Map::from_iter([("type".into(), json!("web_search"))]);
    if !filters.allowed.is_empty() || !filters.blocked.is_empty() {
        tool.insert(
            "filters".into(),
            json!({
                "allowed_domains": filters.allowed,
                "blocked_domains": filters.blocked,
            }),
        );
    }
    Value::Object(tool)
}

#[derive(Default)]
pub(super) struct DomainFilters {
    pub(super) allowed: Vec<String>,
    pub(super) blocked: Vec<String>,
}

pub(super) fn normalize_domain_filters(domain_filter: Option<&[String]>) -> DomainFilters {
    let mut filters = DomainFilters::default();
    for raw in domain_filter.into_iter().flatten() {
        let Some(domain) = normalize_domain(raw) else {
            continue;
        };
        let target = if raw.trim().starts_with('-') {
            &mut filters.blocked
        } else {
            &mut filters.allowed
        };
        if !target.contains(&domain) {
            target.push(domain);
        }
    }
    filters.allowed.truncate(100);
    filters.blocked.truncate(100);
    filters
}

fn normalize_domain(raw: &str) -> Option<String> {
    let mut input = raw
        .trim()
        .trim_start_matches('-')
        .trim()
        .to_ascii_lowercase();
    if input.is_empty() {
        return None;
    }
    if let Ok(url) = Url::parse(&input).or_else(|_| Url::parse(&format!("https://{input}"))) {
        input = url.host_str()?.to_string();
    } else {
        input = input.split('/').next()?.split(':').next()?.to_string();
    }
    let input = input.trim_matches('.').to_string();
    Regex::new(r"^[a-z0-9][a-z0-9.-]*\.[a-z]{2,}$")
        .ok()?
        .is_match(&input)
        .then_some(input)
}

fn parse_openai_search_output(text: &str) -> Result<Vec<Value>, ToolError> {
    let trimmed = text.trim();
    if trimmed.starts_with('{') || trimmed.starts_with('[') {
        let value: Value = serde_json::from_str(trimmed).map_err(|err| {
            ToolError::Message(format!("OpenAI web search returned invalid JSON: {err}"))
        })?;
        return Ok(value
            .get("output")
            .and_then(Value::as_array)
            .cloned()
            .unwrap_or_default());
    }

    let mut output = Vec::new();
    let mut completed_output = None;
    for line in text.lines() {
        let Some(data) = line.strip_prefix("data: ") else {
            continue;
        };
        let data = data.trim();
        if data.is_empty() || data == "[DONE]" {
            continue;
        }
        let Ok(value) = serde_json::from_str::<Value>(data) else {
            continue;
        };
        if value.get("type").and_then(Value::as_str) == Some("response.output_item.done") {
            if let Some(item) = value.get("item") {
                output.push(item.clone());
            }
        }
        if matches!(
            value.get("type").and_then(Value::as_str),
            Some("response.done" | "response.completed")
        ) {
            completed_output = value
                .get("response")
                .and_then(|response| response.get("output"))
                .and_then(Value::as_array)
                .cloned();
        }
    }
    Ok(completed_output.unwrap_or(output))
}

fn extract_openai_search_results(output: &[Value], num_results: usize) -> Vec<SearchItem> {
    let mut results = Vec::new();
    let mut seen = HashSet::new();

    for item in output {
        if item.get("type").and_then(Value::as_str) != Some("message") {
            continue;
        }
        for part in item
            .get("content")
            .and_then(Value::as_array)
            .into_iter()
            .flatten()
        {
            let text = part.get("text").and_then(Value::as_str).unwrap_or_default();
            for annotation in part
                .get("annotations")
                .and_then(Value::as_array)
                .into_iter()
                .flatten()
            {
                if annotation.get("type").and_then(Value::as_str) != Some("url_citation") {
                    continue;
                }
                add_openai_search_result(
                    &mut results,
                    &mut seen,
                    annotation.get("url").and_then(Value::as_str),
                    annotation.get("title").and_then(Value::as_str),
                    citation_snippet(
                        text,
                        annotation.get("start_index").and_then(Value::as_u64),
                        annotation.get("end_index").and_then(Value::as_u64),
                    ),
                );
            }
        }
    }

    for item in output {
        if item.get("type").and_then(Value::as_str) != Some("web_search_call") {
            continue;
        }
        for group in [
            item.get("action")
                .and_then(|action| action.get("sources"))
                .and_then(Value::as_array),
            item.get("sources").and_then(Value::as_array),
            item.get("results").and_then(Value::as_array),
        ]
        .into_iter()
        .flatten()
        {
            for source in group {
                add_openai_search_result(
                    &mut results,
                    &mut seen,
                    source
                        .get("url")
                        .or_else(|| source.get("source_website_url"))
                        .and_then(Value::as_str),
                    source
                        .get("title")
                        .or_else(|| source.get("caption"))
                        .and_then(Value::as_str),
                    String::new(),
                );
            }
        }
    }

    results.truncate(num_results.min(20));
    results
}

fn add_openai_search_result(
    results: &mut Vec<SearchItem>,
    seen: &mut HashSet<String>,
    url: Option<&str>,
    title: Option<&str>,
    snippet: String,
) {
    let Some(url) = url.filter(|url| !url.trim().is_empty()) else {
        return;
    };
    let url = clean_openai_source_url(url);
    if !seen.insert(url.clone()) {
        return;
    }
    results.push(SearchItem {
        title: title
            .filter(|title| !title.trim().is_empty())
            .map(str::to_string)
            .or_else(|| Some(url.clone())),
        url: Some(url),
        snippet,
    });
}

fn clean_openai_source_url(raw_url: &str) -> String {
    Url::parse(raw_url)
        .map(|mut url| {
            let query = url
                .query_pairs()
                .filter(|(key, value)| !(key == "utm_source" && value == "openai"))
                .map(|(key, value)| (key.into_owned(), value.into_owned()))
                .collect::<Vec<_>>();
            url.set_query(None);
            if !query.is_empty() {
                url.query_pairs_mut().extend_pairs(query);
            }
            url.to_string()
        })
        .unwrap_or_else(|_| raw_url.replace("?utm_source=openai", ""))
}

fn citation_snippet(text: &str, start: Option<u64>, end: Option<u64>) -> String {
    let (Some(start), Some(end)) = (start, end) else {
        return String::new();
    };
    let start = start as usize;
    let end = end.max(start as u64) as usize;
    let before = start.saturating_sub(100);
    let after = end.saturating_add(100);
    text.chars()
        .skip(before)
        .take(after.saturating_sub(before))
        .collect::<String>()
        .replace(['[', ']', '(', ')'], "")
        .trim()
        .chars()
        .take(300)
        .collect()
}
