from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import os
import re

try:
    from llama_cpp import Llama
    USE_LLAMA_CPP = True
    print("llama-cpp-python is installed. Will use GGUF inference.")
except ImportError:
    USE_LLAMA_CPP = False
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    print("llama-cpp-python is not installed. Falling back to transformers inference.")

app = FastAPI()

# CORS configuration to allow the frontend to communicate with this API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins for local development
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

class TextRequest(BaseModel):
    text: str

class MapResponse(BaseModel):
    markdown: str

# ---------------------------------------------------------
# AI Model Initialization (Phase 5)
# ---------------------------------------------------------

if USE_LLAMA_CPP:
    GGUF_PATH = "./production_mindmap_model.gguf"
    print("Loading AI Model (GGUF) into memory...")
    model = Llama(
        model_path=GGUF_PATH,
        n_ctx=2048,
        n_gpu_layers=-1 # Use all GPU layers if available
    )
    print("AIモデル(GGUF)の起動完了！")
else:
    MERGED_MODEL_PATH = "./production_mindmap_model_merged"
    print("Loading AI Model (Transformers FP16) into memory...")
    tokenizer = AutoTokenizer.from_pretrained(MERGED_MODEL_PATH)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    model = AutoModelForCausalLM.from_pretrained(
        MERGED_MODEL_PATH, 
        torch_dtype=torch.float16, 
        device_map="auto"
    )
    model.eval()
    print("AIモデル(Transformers)の起動完了！")

STRICT_SYSTEM_PROMPT = """あなたは極めて優秀で厳密な情報抽出アシスタントです。入力文章の論理構造を正確に読み取り、Markdown形式の目次（マインドマップ）を出力してください。

【厳格なルール】
1. 企業名、人物名、数値などの固有名詞や事実関係は、入力文章から「正確に」抽出すること。絶対に捏造したり、他の主語と混同したりしてはなりません。
2. ただし、情報を階層化して分かりやすく整理するための「一般的なカテゴリ名（例：背景、特徴、課題、概要など）」を親見出しとして補足することは許可します。
3. あなたの推測や外部知識を一切混ぜず、入力文章内の事実のみを構造化してください。"""

@app.post("/generate", response_model=MapResponse)
def generate_mindmap(request: TextRequest):
    input_text = request.text.strip()
    if not input_text:
        raise HTTPException(status_code=400, detail="Text is empty")

    print(f"APIリクエストを受信しました（文字数: {len(input_text)}文字）")
    print("AIが推論（Markdown構造）を生成中...")

    # Qwen用のチャットプロンプト構築
    USER_PROMPT = f"""以下の文章から論理構造を抽出し、Markdown形式の目次（マインドマップ）を出力してください。

【出力時の厳守ルール（違反厳禁）】
1. 否定表現の厳守：「〜しない」「過度に依存しない」などの否定表現を絶対に見落とさず、意味を逆転させないこと。
2. 創作の禁止：記事に明記されていない具体的な行動や予定（例：「〜への参加」「〜の強化を目指す」など）を勝手に推測して付け足さないこと。
3. 事実の完全一致：抽出した内容が、元の文章の事実と完全に一致していることのみを出力すること。

入力文章:
{input_text}"""
    messages = [
        {"role": "system", "content": STRICT_SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT}
    ]
    
    if USE_LLAMA_CPP:
        response = model.create_chat_completion(
            messages=messages,
            max_tokens=1024,
            temperature=0.0, # Greedy Decoding (推測・創造を完全に排除)
            repeat_penalty=1.1
        )
        generated_markdown = response['choices'][0]['message']['content'].strip()
    else:
        # トークナイザーでプロンプト化
        prompt = tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=1024, # 新しく生成するトークン数の上限
                do_sample=False,     # Greedy Decoding を有効化
                repetition_penalty=1.1,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        # 4. 入力プロンプト部分を切り落として、新しく生成された部分だけを取り出す
        input_length = inputs["input_ids"].shape[1]
        generated_tokens = outputs[0][input_length:]
        
        # 5. 数値（トークン）を人間が読めるMarkdownテキストにデコード
        generated_markdown = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
    
    # 【超重要】T5トークナイザーが改行(\n)を消してしまう問題の対策
    # 見出し記号（#）の前に強制的に改行を挿入して、Markmapが認識できるようにする
    generated_markdown = re.sub(r'\s*(#+ )', r'\n\1', generated_markdown).strip()
    
    # Markmapは必ず「#（大見出し）」から始まる必要があるため、存在しない場合は強制付与
    if not generated_markdown.startswith('#'):
        if '##' in generated_markdown or '###' in generated_markdown:
            generated_markdown = "# マインドマップ\n" + generated_markdown
        else:
            generated_markdown = "# マインドマップ\n## 抽出結果\n- " + generated_markdown.replace('\n', '\n- ')

    print("生成が完了しました！ (コンソール出力時の文字化けエラーを防ぐため内容の表示を省略します)")
    
    return MapResponse(markdown=generated_markdown)

# フロントエンドの静的ファイルをマウント
app.mount("/assets", StaticFiles(directory="frontend"), name="assets")

@app.get("/")
async def root():
    # ルートURLにアクセスされたらフロントエンドのindex.htmlを返す
    return FileResponse("frontend/index.html")

@app.get("/{filename}")
async def get_frontend_file(filename: str):
    # app.js や style.css などのファイルを返す
    file_path = os.path.join("frontend", filename)
    if os.path.exists(file_path):
        return FileResponse(file_path)
    raise HTTPException(status_code=404, detail="File not found")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)