use std::io::{BufRead, Write};

use orfail::OrFail;

use crate::{
    command::Command,
    message::{Message, MessageLog, Role},
};

const API_END_POINT: &str = "https://api.anthropic.com/v1/messages";
const ANTHROPIC_VERSION: &str = "2023-06-01";
const MAX_TOKENS: u32 = 10_000;

#[derive(Debug)]
pub struct Claude {
    api_key: String,
    model: String,
}

impl Claude {
    pub fn new(command: &Command, model: String) -> orfail::Result<Self> {
        Ok(Self {
            api_key: command.anthropic_api_key.clone().or_fail()?,
            model,
        })
    }

    pub fn run(&self, log: &MessageLog) -> orfail::Result<Message> {
        let (log, system_message) = log.strip_system_message();
        let request = nojson::json(|f| {
            f.object(|f| {
                f.member("model", &self.model)?;
                f.member("stream", true)?;
                f.member("max_tokens", MAX_TOKENS)?;
                f.member("messages", &log.messages)?;
                if let Some(system_message) = &system_message {
                    f.member("system", system_message)?;
                }
                Ok(())
            })
        });

        let request_builder = crate::curl::CurlRequest::new(API_END_POINT)
            .header("Content-Type", "application/json")
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", ANTHROPIC_VERSION);
        let response = request_builder.post(request)?;
        let reader = response.check_success()?;
        let reply = self.handle_stream_response(reader).or_fail()?;
        Ok(reply)
    }

    fn handle_stream_response<R: BufRead>(&self, reader: R) -> orfail::Result<Message> {
        let mut content = String::new();

        for line in reader.lines() {
            let line = line.or_fail()?;
            if line.is_empty() {
                continue;
            }
            if !line.starts_with("data: ") {
                continue;
            }
            if line == "data: [DONE]" {
                break;
            }

            let nojson::Json(data) = line["data: ".len()..]
                .parse::<nojson::Json<Data>>()
                .or_fail_with(|e| format!("failed to parse line: {line} ({e})"))?;
            match data {
                Data::MessageStart { stop_reason } => {
                    if let Some(reason) = stop_reason {
                        (reason == "end_turn").or_fail_with(|()| format!("API error: {reason}"))?;
                    }
                }
                Data::MessageDelta { stop_reason } => {
                    if let Some(reason) = stop_reason {
                        (reason == "end_turn").or_fail_with(|()| format!("API error: {reason}"))?;
                    }
                }
                Data::MessageStop => {}
                Data::Ping => {}
                Data::ContentBlockStart { content_block } => match content_block {
                    ContentBlock::Text(text) => {
                        content.push_str(&text);
                        print!("{text}");
                        std::io::stdout().flush().or_fail()?;
                    }
                    ContentBlock::ServerToolUse { id, name, input } => {
                        eprintln!("Server tool use: id={id}, name={name}, input={input}");
                    }
                    ContentBlock::ToolResult {
                        tool_use_id,
                        content: tool_content,
                    } => {
                        eprintln!("Tool result: tool_use_id={tool_use_id}, content={tool_content}");
                    }
                },
                Data::ContentBlockDelta { delta } => {
                    content.push_str(&delta.text);
                    print!("{}", delta.text);
                    std::io::stdout().flush().or_fail()?;
                }
                Data::ContentBlockStop => {}
                Data::Error { error } => {
                    return Err(orfail::Failure::new(format!(
                        "Claude API error: reason={error}"
                    )));
                }
            }
        }
        println!();

        Ok(Message {
            role: Role::Assistant,
            content,
            model: Some(self.model.clone()),
        })
    }
}

#[derive(Debug)]
enum Data {
    MessageStart { stop_reason: Option<String> },
    MessageDelta { stop_reason: Option<String> },
    MessageStop,
    ContentBlockStart { content_block: ContentBlock },
    ContentBlockDelta { delta: Delta },
    ContentBlockStop,
    Ping,
    Error { error: String },
}

impl<'text, 'raw> TryFrom<nojson::RawJsonValue<'text, 'raw>> for Data {
    type Error = nojson::JsonParseError;

    fn try_from(value: nojson::RawJsonValue<'text, 'raw>) -> Result<Self, nojson::JsonParseError> {
        let ty = value.to_member("type")?.required()?;
        match ty.to_unquoted_string_str()?.as_ref() {
            "message_start" => Ok(Self::MessageStart {
                stop_reason: value.to_member("stop_reason")?.try_into()?,
            }),
            "message_delta" => Ok(Self::MessageDelta {
                stop_reason: value.to_member("stop_reason")?.try_into()?,
            }),
            "message_stop" => Ok(Self::MessageStop),
            "content_block_start" => {
                let content_block = value.to_member("content_block")?.required()?;
                let block_type = content_block.to_member("type")?.required()?;

                match block_type.to_unquoted_string_str()?.as_ref() {
                    "text" => {
                        let text = content_block.to_member("text")?.required()?;
                        Ok(Self::ContentBlockStart {
                            content_block: ContentBlock::Text(text.try_into()?),
                        })
                    }
                    "server_tool_use" => {
                        let id = content_block.to_member("id")?.required()?;
                        let name = content_block.to_member("name")?.required()?;
                        let input = content_block.to_member("input")?.required()?;

                        Ok(Self::ContentBlockStart {
                            content_block: ContentBlock::ServerToolUse {
                                id: id.try_into()?,
                                name: name.try_into()?,
                                input: input.extract().into_owned(),
                            },
                        })
                    }
                    ty => Err(content_block.invalid(format!("unknown content block type: {ty}"))),
                }
            }
            "content_block_delta" => {
                let delta = value.to_member("delta")?.required()?;
                let text = delta.to_member("text")?.required()?;
                Ok(Self::ContentBlockDelta {
                    delta: Delta {
                        text: text.try_into()?,
                    },
                })
            }
            "content_block_stop" => Ok(Self::ContentBlockStop),
            "ping" => Ok(Self::Ping),
            "error" => {
                let error = value.to_member("error")?.required()?;
                Ok(Self::Error {
                    error: error.to_string(),
                })
            }
            ty => Err(nojson::JsonParseError::invalid_value(
                value,
                format!("unknown message type: {ty}"),
            )),
        }
    }
}

#[derive(Debug)]
enum ContentBlock {
    Text(String),
    ServerToolUse {
        id: String,
        name: String,
        input: nojson::RawJsonOwned,
    },
    ToolResult {
        tool_use_id: String,
        content: nojson::RawJsonOwned,
    },
}

impl<'text, 'raw> TryFrom<nojson::RawJsonValue<'text, 'raw>> for ContentBlock {
    type Error = nojson::JsonParseError;

    fn try_from(value: nojson::RawJsonValue<'text, 'raw>) -> Result<Self, nojson::JsonParseError> {
        let ty = value.to_member("type")?.required()?;
        let type_str = ty.to_unquoted_string_str()?;

        match type_str.as_ref() {
            "text" => {
                let text = value.to_member("text")?.required()?;
                Ok(Self::Text(text.try_into()?))
            }
            "server_tool_use" => {
                let id = value.to_member("id")?.required()?;
                let name = value.to_member("name")?.required()?;
                let input = value.to_member("input")?.required()?;
                Ok(Self::ServerToolUse {
                    id: id.try_into()?,
                    name: name.try_into()?,
                    input: input.extract().into_owned(),
                })
            }
            // Handle all tool result types dynamically (ends with _tool_result)
            ty if ty.ends_with("_tool_result") => {
                let tool_use_id = value.to_member("tool_use_id")?.required()?;
                let content = value.to_member("content")?.required()?;
                Ok(Self::ToolResult {
                    tool_use_id: tool_use_id.try_into()?,
                    content: content.extract().into_owned(),
                })
            }
            ty => Err(value.invalid(format!("unknown content block type: {ty}"))),
        }
    }
}

#[derive(Debug)]
struct Delta {
    text: String,
}
