/*
MIT License

Copyright (c) 2024 SHUGO AMEMIYA

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import openai
import os
import sqlite3
from datetime import datetime

app = FastAPI()

class RequestPayload(BaseModel):
    user_id: str
    user: str
    key: str
    model: str
    system: str = None  # system をオプショナルに設定
    prompt_id: str = None  # prompt_id をオプショナルに設定
    max_rounds: int = 10  # 最大会話回数を指定

class RequestPayloadReset(BaseModel):
    user_id: str

# プロンプトファイルが保存されているディレクトリ
PROMPT_DIR = "./prompts/"

# SQLiteデータベースの初期化
def init_db():
    conn = sqlite3.connect('conversations.db')
    cursor = conn.cursor()
    cursor.execute('''CREATE TABLE IF NOT EXISTS conversations (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        user_id TEXT,
                        role TEXT,
                        content TEXT,
                        timestamp TEXT
                      )''')
    conn.commit()
    conn.close()

# データベースへの会話履歴の追加
def add_message_to_db(user_id: str, role: str, content: str):
    conn = sqlite3.connect('conversations.db')
    cursor = conn.cursor()
    cursor.execute('''INSERT INTO conversations (user_id, role, content, timestamp)
                      VALUES (?, ?, ?, ?)''', (user_id, role, content, datetime.now().isoformat()))
    conn.commit()
    conn.close()

# データベースから履歴を取得
def get_conversation_history(user_id: str, max_messages: int):
    conn = sqlite3.connect('conversations.db')
    cursor = conn.cursor()
    cursor.execute('''SELECT role, content FROM conversations
                      WHERE user_id = ?
                      ORDER BY timestamp ASC
                      LIMIT ?''', (user_id, max_messages))
    history = cursor.fetchall()
    conn.close()
    return [{"role": role, "content": content} for role, content in history]

# 会話回数を超えた場合に古いメッセージを削除
def trim_conversation_history(user_id: str, max_rounds: int):
    conn = sqlite3.connect('conversations.db')
    cursor = conn.cursor()
    cursor.execute('''SELECT COUNT(*) FROM conversations WHERE user_id = ?''', (user_id,))
    total_messages = cursor.fetchone()[0]

    # 指定された回数を超えたメッセージを削除
    max_messages = max_rounds * 2 + 1  # systemメッセージを含める
    if total_messages > max_messages:
        cursor.execute('''DELETE FROM conversations WHERE id IN (
                          SELECT id FROM conversations WHERE user_id = ?
                          ORDER BY timestamp ASC LIMIT ?)''', (user_id, total_messages - max_messages))
    conn.commit()
    conn.close()

# データベースから特定ユーザーの履歴を削除
def reset_conversation_history(user_id: str):
    conn = sqlite3.connect('conversations.db')
    cursor = conn.cursor()
    cursor.execute('''DELETE FROM conversations WHERE user_id = ?''', (user_id,))
    conn.commit()
    conn.close()

# プロンプトを読み込む関数
def load_prompt(prompt_id: str) -> str:
    prompt_path = os.path.join(PROMPT_DIR, f"{prompt_id}.txt")
    if os.path.exists(prompt_path):
        try:
            with open(prompt_path, "r", encoding="utf-8") as file:
                return file.read()
        except UnicodeDecodeError as e:
            raise HTTPException(status_code=400, detail=f"Error decoding the prompt file: {str(e)}")
    else:
        raise HTTPException(status_code=404, detail="Prompt not found.")

@app.on_event("startup")
def startup_event():
    init_db()

@app.post("/chat/")
async def chat(payload: RequestPayload):
    try:
        # APIキーを設定
        openai.api_key = payload.key

        # 直近の会話履歴を取得
        max_messages = payload.max_rounds * 2 + 1  # systemメッセージを含める
        recent_messages = get_conversation_history(payload.user_id, max_messages)

        # プロンプトの読み込み（初回のみ）
        if len(recent_messages) == 0:
            if payload.system:
                add_message_to_db(payload.user_id, "system", payload.system)
            elif payload.prompt_id:
                prompt_text = load_prompt(payload.prompt_id)
                add_message_to_db(payload.user_id, "system", prompt_text)
            recent_messages = get_conversation_history(payload.user_id, max_messages)

        # 履歴にユーザーメッセージを追加
        add_message_to_db(payload.user_id, "user", payload.user)
        recent_messages = get_conversation_history(payload.user_id, max_messages)

        # 送信されるリクエスト内容をターミナルに出力
        print("Sending request to OpenAI API:")
        print({
            "model": payload.model,
            "messages": recent_messages
        })

        # OpenAI APIにチャット完成リクエストを送信
        response = openai.chat.completions.create(
            model=payload.model,
            messages=recent_messages
        )

        # 受信したレスポンスをターミナルに出力
        print("Received response from OpenAI API:")
        print(response)

        # APIの応答を履歴に追加
        if response and response.choices:
            assistant_message = response.choices[0].message.content
            add_message_to_db(payload.user_id, "assistant", assistant_message)
            recent_messages = get_conversation_history(payload.user_id, max_messages)

        # 会話回数を超えた場合に古いメッセージを削除
        trim_conversation_history(payload.user_id, payload.max_rounds)

        return {"response": assistant_message}

    except Exception as e:
        # エラーを捕捉して500エラーを返す
        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")

@app.post("/reset/")
async def reset_history(payload: RequestPayloadReset):
    try:
        reset_conversation_history(payload.user_id)
        return {"status": "success", "message": "Conversation history has been reset."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")