package me.khodak.claudeusage

import android.util.Log
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import me.khodak.claudeusage.BuildConfig
import me.khodak.claudeusage.data.PreferencesManager
import me.khodak.claudeusage.data.UsageData
import okhttp3.OkHttpClient
import okhttp3.Request
import org.json.JSONArray
import org.json.JSONObject
import java.text.SimpleDateFormat
import java.util.*
import java.util.concurrent.TimeUnit

class UsageRepository(private val prefs: PreferencesManager) {

    private val client = OkHttpClient.Builder()
        .connectTimeout(15, TimeUnit.SECONDS)
        .readTimeout(15, TimeUnit.SECONDS)
        .followRedirects(false)
        .build()

    private val desktopUA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " +
        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"

    private val debugBuf = StringBuilder()
    val lastDebugInfo: String get() = debugBuf.toString()

    suspend fun fetchUsage(): UsageData = withContext(Dispatchers.IO) {
        debugBuf.clear()
        val cookies = prefs.getCookies()
            ?: return@withContext UsageData(errorMessage = "Not logged in")

        prefs.markTodayActive()

        val base = UsageData(
            isLoggedIn = true,
            sessionStartEpoch = prefs.getSessionStart(),
            weeklyActiveDaysMask = prefs.getWeeklyMask(),
            lastUpdated = System.currentTimeMillis()
        )

        // Step 1: get org ID
        var orgId = prefs.getOrgId()
        var orgUsageData: UsageData? = null
        val (fetchedId, fetchedUsage) = fetchOrgInfo(cookies)
        if (orgId == null && fetchedId != null) {
            orgId = fetchedId
            prefs.saveOrgId(orgId)
        }
        orgUsageData = fetchedUsage

        if (orgId == null) return@withContext base.copy(errorMessage = "Can't reach claude.ai")

        // NOTE: message-count data from the org endpoint (orgUsageData) is a FALLBACK only — do
        // NOT return here. The /usage utilization endpoint below is the preferred signal and the
        // one the weekly bar + in-app history chart depend on, so it must be attempted whenever we
        // have an org id. Returning early here was why the history chart stayed empty: any account
        // whose org JSON carried a message limit short-circuited before /usage, so
        // fiveHourUtilization/weeklyUtilization never got populated and recordHistory() had nothing
        // to store. The org message data is used only if /usage yields nothing (fallback below).

        // Step 2: /usage endpoint — returns utilization percentages (preferred signal)
        try {
            val usageUrl = "https://claude.ai/api/organizations/$orgId/usage"
            val resp = client.newCall(buildRequest(usageUrl, cookies)).execute()
            val code = resp.code
            if (BuildConfig.DEBUG) Log.d(TAG, "GET $usageUrl → $code")
            if (code == 401 || code == 403) {
                prefs.clearSession()
                return@withContext UsageData(errorMessage = "Session expired — please sign in again")
            }
            val body = resp.body?.string() ?: ""
            if (BuildConfig.DEBUG) debugBuf.append("$usageUrl\n→ $code: ${body.take(400)}\n\n")
            val utilData = tryParseUtilizationBody(body)
            if (utilData != null) {
                prefs.resetAuthFailCount()
                return@withContext base.copy(
                    fiveHourUtilization = utilData.fiveHourUtilization,
                    weeklyUtilization = utilData.weeklyUtilization,
                    utilizationResetAtEpoch = utilData.utilizationResetAtEpoch,
                    weeklyResetAtEpoch = utilData.weeklyResetAtEpoch
                )
            }
        } catch (e: Exception) {
            if (BuildConfig.DEBUG) Log.w(TAG, "/usage failed: ${e.message}")
        }

        // /usage gave no utilization → fall back to the message-count data the org endpoint
        // already returned (this is the early-return that used to live before Step 2).
        if (orgUsageData?.hasRateLimitData == true) {
            prefs.resetAuthFailCount()
            return@withContext base.copy(
                messagesUsed = orgUsageData.messagesUsed,
                messagesLimit = orgUsageData.messagesLimit,
                messagesRemaining = orgUsageData.messagesRemaining,
                resetAtEpoch = orgUsageData.resetAtEpoch
            )
        }

        // Step 3: fallback endpoints (message-count style)
        val endpoints = listOf(
            "https://claude.ai/api/organizations/$orgId/rate_limit_status",
            "https://claude.ai/api/organizations/$orgId"
        )

        for (url in endpoints) {
            try {
                val req = buildRequest(url, cookies)
                val resp = client.newCall(req).execute()
                val code = resp.code
                if (BuildConfig.DEBUG) Log.d(TAG, "GET $url → $code")

                if (code == 401 || code == 403) {
                    if (prefs.incAuthFailCount() >= AUTH_FAIL_LIMIT) {
                        prefs.clearSession()
                        prefs.resetAuthFailCount()
                        return@withContext UsageData(errorMessage = "Session expired — please sign in again")
                    }
                    // Transient auth failure — keep showing last-good data instead of logging out.
                    return@withContext base
                }

                val rateLimitData = extractRateLimitHeaders(resp.headers)
                val body = resp.body?.string() ?: ""
                if (BuildConfig.DEBUG) {
                    debugBuf.append("$url\n→ $code: ${body.take(300)}\n\n")
                    Log.d(TAG, "Body: ${body.take(300)}")
                }

                val parsed = tryParseUsageBody(body, rateLimitData)
                if (parsed.hasRateLimitData) {
                    return@withContext base.copy(
                        messagesUsed = parsed.messagesUsed,
                        messagesLimit = parsed.messagesLimit,
                        messagesRemaining = parsed.messagesRemaining,
                        resetAtEpoch = parsed.resetAtEpoch,
                        isRateLimited = parsed.isRateLimited
                    )
                }
            } catch (e: Exception) {
                if (BuildConfig.DEBUG) Log.w(TAG, "Endpoint $url failed: ${e.message}")
                if (BuildConfig.DEBUG) debugBuf.append("$url\n→ ERROR: ${e.message}\n\n")
            }
        }

        // Step 3: try /api/me endpoint
        try {
            val req = buildRequest("https://claude.ai/api/me", cookies)
            val resp = client.newCall(req).execute()
            val body = resp.body?.string() ?: ""
            if (BuildConfig.DEBUG) debugBuf.append("https://claude.ai/api/me\n→ ${resp.code}: ${body.take(400)}\n\n")
            if (resp.isSuccessful && body.isNotBlank() && !body.startsWith("<")) {
                val json = JSONObject(body)
                val parsed = tryParseOrgForUsage(json)
                if (parsed?.hasRateLimitData == true) {
                    return@withContext base.copy(
                        messagesUsed = parsed.messagesUsed,
                        messagesLimit = parsed.messagesLimit,
                        messagesRemaining = parsed.messagesRemaining,
                        resetAtEpoch = parsed.resetAtEpoch
                    )
                }
            }
        } catch (e: Exception) {
            if (BuildConfig.DEBUG) debugBuf.append("https://claude.ai/api/me → ERROR: ${e.message}\n\n")
        }

        // Step 4: try page HTML for __NEXT_DATA__
        val htmlData = fetchFromPageHtml(cookies)
        if (htmlData?.hasRateLimitData == true) {
            return@withContext base.copy(
                messagesUsed = htmlData.messagesUsed,
                messagesLimit = htmlData.messagesLimit,
                messagesRemaining = htmlData.messagesRemaining,
                resetAtEpoch = htmlData.resetAtEpoch
            )
        }

        base
    }

    private fun fetchFromPageHtml(cookies: String): UsageData? {
        return try {
            val resp = client.newCall(buildRequest("https://claude.ai/", cookies)).execute()
            val html = resp.body?.string() ?: return null
            val marker = """<script id="__NEXT_DATA__" type="application/json">"""
            val start = html.indexOf(marker)
            if (start < 0) {
                if (BuildConfig.DEBUG) debugBuf.append("HTML: no __NEXT_DATA__ found\n")
                return null
            }
            val jsonStart = start + marker.length
            val jsonEnd = html.indexOf("</script>", jsonStart)
            if (jsonEnd < 0) return null
            val nextData = JSONObject(html.substring(jsonStart, jsonEnd))
            val topKeys = nextData.keys().asSequence().toList()
            if (BuildConfig.DEBUG) debugBuf.append("HTML __NEXT_DATA__ top keys: $topKeys\n")
            tryExtractFromNextData(nextData)
        } catch (e: Exception) {
            if (BuildConfig.DEBUG) debugBuf.append("HTML scrape error: ${e.message}\n")
            null
        }
    }

    private fun tryExtractFromNextData(json: JSONObject): UsageData? {
        // Walk common Next.js page prop paths looking for usage data
        val paths = listOf(
            listOf("props", "pageProps", "rateLimits"),
            listOf("props", "pageProps", "usage"),
            listOf("props", "pageProps", "account"),
            listOf("props", "pageProps", "organization"),
            listOf("props", "pageProps"),
        )
        for (path in paths) {
            var obj: Any? = json
            for (key in path) {
                obj = (obj as? JSONObject)?.opt(key)
            }
            val o = obj as? JSONObject ?: continue
            if (BuildConfig.DEBUG) debugBuf.append("NextData at [${path.joinToString(".")}]: ${o.toString().take(300)}\n")
            val usage = tryParseOrgForUsage(o) ?: tryParseUsageBody(o.toString(), UsageData())
            if (usage.hasRateLimitData) return usage
        }
        return null
    }

    private fun fetchOrgInfo(cookies: String): Pair<String?, UsageData?> {
        return try {
            val resp = client.newCall(buildRequest("https://claude.ai/api/organizations", cookies)).execute()
            if (BuildConfig.DEBUG) Log.d(TAG, "Orgs → ${resp.code}")
            val body = resp.body?.string() ?: return Pair(null, null)
            if (BuildConfig.DEBUG) debugBuf.append("https://claude.ai/api/organizations\n→ ${resp.code}: ${body.take(400)}\n\n")
            if (body.isBlank() || body.startsWith("<")) return Pair(null, null)
            val arr = JSONArray(body)
            val org = arr.optJSONObject(0) ?: return Pair(null, null)
            val orgId = org.optString("uuid").ifBlank { null }
            val usageData = tryParseOrgForUsage(org)
            Pair(orgId, usageData)
        } catch (e: Exception) {
            Log.w(TAG, "fetchOrgInfo: ${e.message}")
            Pair(null, null)
        }
    }

    private fun tryParseOrgForUsage(org: JSONObject): UsageData? {
        val limit = org.optInt("message_limit", -1).takeIf { it > 0 }
            ?: org.optJSONObject("capabilities")?.optInt("message_limit", -1)?.takeIf { it > 0 }
            ?: org.optJSONObject("usage")?.optInt("limit", -1)?.takeIf { it > 0 }
            ?: org.optJSONObject("rate_limit")?.optInt("limit", -1)?.takeIf { it > 0 }
            ?: org.optInt("limit", -1).takeIf { it > 0 }
        val used = org.optInt("messages_sent", -1).takeIf { it >= 0 }
            ?: org.optJSONObject("usage")?.optInt("messages_sent", -1)?.takeIf { it >= 0 }
            ?: org.optJSONObject("rate_limit")?.optInt("used", -1)?.takeIf { it >= 0 }
            ?: org.optInt("used", -1).takeIf { it >= 0 }
        val remaining = org.optInt("messages_remaining", -1).takeIf { it >= 0 }
            ?: org.optJSONObject("usage")?.optInt("remaining", -1)?.takeIf { it >= 0 }
            ?: org.optJSONObject("rate_limit")?.optInt("remaining", -1)?.takeIf { it >= 0 }
            ?: org.optInt("remaining", -1).takeIf { it >= 0 }
        val resetStr = org.optString("usage_reset_at", "")
            .ifBlank { org.optJSONObject("usage")?.optString("reset_at", "") ?: "" }
            .ifBlank { org.optJSONObject("rate_limit")?.optString("reset_at", "") ?: "" }
        if (limit == null && used == null && remaining == null) return null
        return UsageData(
            messagesLimit = limit ?: -1,
            messagesUsed = used ?: -1,
            messagesRemaining = remaining ?: -1,
            resetAtEpoch = if (resetStr.isNotBlank()) parseResetTime(resetStr) else -1L
        )
    }

    // Parses the /usage endpoint response: {"five_hour":{"utilization":75.0,"resets_at":"..."},...}
    private fun tryParseUtilizationBody(body: String): UsageData? {
        if (body.isBlank() || body.startsWith("<")) return null
        return try {
            val json = JSONObject(body)
            val fiveHour = json.optJSONObject("five_hour") ?: return null
            val utilization = fiveHour.optDouble("utilization", -1.0).toFloat()
            if (utilization < 0f) return null
            val resetsAt = fiveHour.optString("resets_at", "")
            val sevenDay = json.optJSONObject("seven_day")
            val weeklyUtil = sevenDay?.optDouble("utilization", -1.0)?.toFloat() ?: -1f
            val weeklyResetsAt = sevenDay?.optString("resets_at", "") ?: ""
            UsageData(
                fiveHourUtilization = utilization,
                weeklyUtilization = weeklyUtil,
                utilizationResetAtEpoch = if (resetsAt.isNotBlank()) parseResetTime(resetsAt) else -1L,
                weeklyResetAtEpoch = if (weeklyResetsAt.isNotBlank()) parseResetTime(weeklyResetsAt) else -1L
            )
        } catch (e: Exception) { null }
    }

    private fun extractRateLimitHeaders(headers: okhttp3.Headers): UsageData {
        val remaining = headers["anthropic-ratelimit-requests-remaining"]?.toIntOrNull()
        val limit = headers["anthropic-ratelimit-requests-limit"]?.toIntOrNull()
        val reset = headers["anthropic-ratelimit-requests-reset"]
        return UsageData(
            messagesRemaining = remaining ?: -1,
            messagesLimit = limit ?: -1,
            resetAtEpoch = parseResetTime(reset)
        )
    }

    private fun tryParseUsageBody(body: String, base: UsageData): UsageData {
        if (body.isBlank() || body.startsWith("<")) return base
        return try {
            val json = JSONObject(body)
            val used = json.optInt("messages_sent", -1).takeIf { it >= 0 }
                ?: json.optInt("message_count", -1).takeIf { it >= 0 }
                ?: json.optInt("used", -1).takeIf { it >= 0 }
                ?: -1
            val limit = json.optInt("messages_limit", -1).takeIf { it >= 0 }
                ?: json.optInt("limit", -1).takeIf { it >= 0 }
                ?: base.messagesLimit
            val remaining = json.optInt("messages_remaining", -1).takeIf { it >= 0 }
                ?: json.optInt("remaining", -1).takeIf { it >= 0 }
                ?: base.messagesRemaining
            val reset = json.optString("reset_at", "").ifBlank { json.optString("resets_at", "") }
            base.copy(
                messagesUsed = used,
                messagesLimit = limit,
                messagesRemaining = remaining,
                resetAtEpoch = if (reset.isNotBlank()) parseResetTime(reset) else base.resetAtEpoch
            )
        } catch (e: Exception) { base }
    }

    private fun parseResetTime(value: String?): Long {
        if (value.isNullOrBlank()) return -1L
        // Normalize: truncate sub-second digits beyond 3 (e.g. microseconds → milliseconds)
        val normalized = value.replace(Regex("(\\.\\d{3})\\d+"), "$1")
        val formats = listOf(
            "yyyy-MM-dd'T'HH:mm:ss.SSSXXX",
            "yyyy-MM-dd'T'HH:mm:ssXXX",
            "yyyy-MM-dd'T'HH:mm:ss.SSS'Z'",
            "yyyy-MM-dd'T'HH:mm:ss'Z'"
        )
        for (fmt in formats) {
            try {
                val sdf = SimpleDateFormat(fmt, Locale.US).also { it.timeZone = TimeZone.getTimeZone("UTC") }
                return sdf.parse(normalized)?.time ?: continue
            } catch (_: Exception) {}
        }
        return value.toLongOrNull()?.let { System.currentTimeMillis() + it * 1000 } ?: -1L
    }

    private fun buildRequest(url: String, cookies: String) = Request.Builder()
        .url(url)
        .header("User-Agent", desktopUA)
        .header("Accept", "application/json, */*")
        .header("Accept-Language", "en-US,en;q=0.9")
        .header("Referer", "https://claude.ai/")
        .header("Cookie", cookies.replace("\r", "").replace("\n", ""))
        .get()
        .build()

    companion object {
        private const val TAG = "UsageRepo"
        // Clear the session only after this many consecutive 401/403s, so one transient
        // auth failure (Cloudflare challenge, brief edge hiccup) doesn't sign the user out.
        private const val AUTH_FAIL_LIMIT = 3
    }
}
