use std::sync::Arc;

use anyhow::{bail, Result};
use async_trait::async_trait;
use serde::de::DeserializeOwned;

use super::prompt_builder::{
    build_round_n_prompt, build_round_one_analysis_prompt, build_round_one_bodies_prompt,
    build_round_one_config_prompt, system_prompt_from_request, user_prompt_from_request,
};
use super::response_parser::{
    build_parse_repair_prompt, extract_json_payload, parse_generation_response,
};
use super::stages::{AnalysisStage, BodiesStage, ConfigStage};
use crate::generator::domain::generation_response::{GenerationResponse, GenerationResult};
use crate::generator::ports::outbound::{GenerationPort, GenerationRequest, LlmClientPort};
use crate::shared::models::{GenerationUsage, PromptMode};

const MAX_ATTEMPTS: usize = 3;

pub struct LiteLlmGenerationAdapter {
    model: String,
    api_key: String,
    pub client: Arc<dyn LlmClientPort>,
    prompt_mode: PromptMode,
}

impl LiteLlmGenerationAdapter {
    pub fn new(
        model: impl Into<String>,
        api_key: impl Into<String>,
        client: Arc<dyn LlmClientPort>,
        prompt_mode: PromptMode,
    ) -> Self {
        Self {
            model: model.into(),
            api_key: api_key.into(),
            client,
            prompt_mode,
        }
    }

    fn set_api_key(&self) {
        if let Some(prefix) = self.model.split('/').next() {
            let env_var = format!("{}_API_KEY", prefix.to_uppercase());
            std::env::set_var(env_var, &self.api_key);
        }
    }

    fn merge_usage(total: &mut GenerationUsage, usage: Option<GenerationUsage>) {
        if let Some(usage) = usage {
            total.calls = total.calls.saturating_add(usage.calls);
            total.prompt_tokens = total.prompt_tokens.saturating_add(usage.prompt_tokens);
            total.completion_tokens = total
                .completion_tokens
                .saturating_add(usage.completion_tokens);
            total.total_tokens = total.total_tokens.saturating_add(usage.total_tokens);
            total.cached_prompt_tokens = total
                .cached_prompt_tokens
                .saturating_add(usage.cached_prompt_tokens);
            total.reasoning_tokens = total
                .reasoning_tokens
                .saturating_add(usage.reasoning_tokens);
            total.thinking_tokens = total.thinking_tokens.saturating_add(usage.thinking_tokens);
        }
    }

    async fn request_json<T>(
        &self,
        system_prompt: &str,
        initial_prompt: String,
        stage_name: &str,
        schema_hint: &str,
        usage_total: &mut GenerationUsage,
    ) -> Result<T>
    where
        T: DeserializeOwned,
    {
        let mut user_prompt = initial_prompt;
        let mut last_error = String::new();

        for attempt in 1..=MAX_ATTEMPTS {
            let (content, usage) = self.client.complete(system_prompt, &user_prompt).await?;
            Self::merge_usage(usage_total, usage);
            let payload = extract_json_payload(&content)?;

            match serde_json::from_str::<T>(&payload) {
                Ok(parsed) => return Ok(parsed),
                Err(err) => {
                    last_error = format!("JSON parse error: {err}");
                    if attempt == MAX_ATTEMPTS {
                        break;
                    }
                    user_prompt =
                        build_parse_repair_prompt(stage_name, schema_hint, &payload, &last_error);
                }
            }
        }

        bail!("{stage_name} failed after {MAX_ATTEMPTS} attempts: {last_error}")
    }

    async fn generate_round_one(&self, request: &GenerationRequest) -> Result<GenerationResult> {
        let system_prompt = system_prompt_from_request(request);
        let mut usage = GenerationUsage::default();

        let analysis: AnalysisStage = self
            .request_json(
                &system_prompt,
                build_round_one_analysis_prompt(),
                "analysis",
                "vulnerability_analysis, handler_logic_pseudocode, invariant_mathematical_proofs",
                &mut usage,
            )
            .await?;

        let bodies_stage: BodiesStage = self
            .request_json(
                &system_prompt,
                build_round_one_bodies_prompt(
                    &analysis,
                    &request.contract_name,
                    &request.contract_path,
                    &request.source_code,
                    &self.prompt_mode,
                )?,
                "bodies",
                "bodies object with valid Solidity syntax",
                &mut usage,
            )
            .await?;

        let config_stage: ConfigStage = self
            .request_json(
                &system_prompt,
                build_round_one_config_prompt(&analysis, &bodies_stage.bodies)?,
                "config",
                "foundry_config mapping to handler functions",
                &mut usage,
            )
            .await?;

        Ok(GenerationResult {
            response: GenerationResponse::Full {
                bodies: Box::new(bodies_stage.bodies),
                foundry_config: Box::new(config_stage.foundry_config),
            },
            usage,
        })
    }

    async fn generate_round_n(&self, request: &GenerationRequest) -> Result<GenerationResult> {
        let system_prompt = system_prompt_from_request(request);
        let domain_context = user_prompt_from_request(request);
        let patch_prompt = build_round_n_prompt(request, &self.prompt_mode)?;
        let mut user_prompt = format!("{domain_context}\n\n{patch_prompt}");
        let mut last_error = String::new();
        let mut usage = GenerationUsage::default();

        for attempt in 1..=MAX_ATTEMPTS {
            let (content, call_usage) = self.client.complete(&system_prompt, &user_prompt).await?;
            Self::merge_usage(&mut usage, call_usage);
            let payload = extract_json_payload(&content)?;

            match parse_generation_response(&payload) {
                Ok(parsed) => {
                    return Ok(GenerationResult {
                        response: parsed,
                        usage,
                    })
                }
                Err(err) => {
                    last_error = err.to_string();
                    if attempt == MAX_ATTEMPTS {
                        break;
                    }
                    user_prompt = build_parse_repair_prompt(
                        "round generation",
                        "{\"mode\":\"patch\",\"bodies_updates\":[...],\"foundry_config_updates\":[...]} — mode MUST be \"patch\", NO other top-level keys",
                        &payload,
                        &last_error,
                    );
                }
            }
        }

        bail!(
            "model returned invalid structured output after {MAX_ATTEMPTS} attempts: {last_error}"
        )
    }
}

#[async_trait]
impl GenerationPort for LiteLlmGenerationAdapter {
    async fn generate(&self, request: GenerationRequest) -> Result<GenerationResult> {
        self.set_api_key();

        if request.round == 1 {
            self.generate_round_one(&request).await
        } else {
            self.generate_round_n(&request).await
        }
    }
}
