# backend/main.py
from fastapi import FastAPI, BackgroundTasks
from git import Repo
from langchain import OpenAI, VectorStore
from terraform import validate, format_code

app = FastAPI()

@app.post("/api/generate")
async def generate_iac(request: IaCRequest, background_tasks: BackgroundTasks):
    # 1. 创建任务
    task = Task.create(
        project=request.project,
        prompt=request.prompt,
        environment=request.environment,
        status="pending"
    )
    
    # 2. 后台处理
    background_tasks.add_task(
        process_request,
        task_id=task.id,
        prompt=request.prompt,
        project=request.project,
        is_new_project=request.is_new_project
    )
    
    return {"task_id": task.id, "status": "pending"}

async def process_request(task_id: str, prompt: str, project: str, is_new_project: bool):
    try:
        # 更新状态
        Task.update(task_id, status="processing")
        
        if is_new_project:
            # 新项目：创建仓库
            repo = create_new_repo(project)
        else:
            # 已有项目：克隆仓库
            repo = clone_repo(project)
        
        # 分析仓库结构
        context = analyze_repo(repo)
        
        # RAG：检索相关代码
        relevant_code = vector_store.search(prompt, filter={"repo": project})
        
        # 生成代码
        generated_code = llm.generate(
            prompt=f"""
            基于以下上下文生成 Terraform 代码：
            
            需求：{prompt}
            
            项目结构：
            {context.structure}
            
            相关代码示例：
            {relevant_code}
            
            标签规范：
            {context.tags}
            
            变量定义：
            {context.variables}
            """
        )
        
        # 验证
        format_code(generated_code)
        validate(generated_code)
        
        # 提交到 Git
        branch = f"feature/{generate_branch_name(prompt)}"
        commit_to_repo(repo, generated_code, branch)
        pr = create_pull_request(repo, branch, prompt)
        
        # 更新状态
        Task.update(task_id, status="completed", pr_url=pr.url)
        
        # 通知
        notify_user(task_id, "代码已生成", pr_url=pr.url)
        
    except Exception as e:
        Task.update(task_id, status="failed", error=str(e))
        notify_user(task_id, "生成失败", error=str(e))

def analyze_repo(repo_path):
    """提取仓库上下文"""
    return {
        "structure": get_directory_tree(repo_path),
        "modules": find_terraform_modules(repo_path),
        "variables": extract_variables(repo_path),
        "tags": extract_tag_patterns(repo_path),
        "providers": extract_provider_config(repo_path),
    }

// frontend/GenerateForm.jsx
function GenerateForm() {
  const [project, setProject] = useState("");
  const [prompt, setPrompt] = useState("");
  const [environment, setEnvironment] = useState("dev");
  const [status, setStatus] = useState(null);
  
  const handleSubmit = async () => {
    const response = await fetch("/api/generate", {
      method: "POST",
      body: JSON.stringify({
        project,
        prompt,
        environment,
        is_new_project: project === "new"
      })
    });
    
    const { task_id } = await response.json();
    
    // 轮询状态
    pollStatus(task_id);
  };
  
  const pollStatus = async (task_id) => {
    const interval = setInterval(async () => {
      const response = await fetch(`/api/task/${task_id}`);
      const data = await response.json();
      
      setStatus(data.status);
      
      if (data.status === "completed") {
        clearInterval(interval);
        showCodePreview(data.pr_url);
      }
    }, 2000);
  };
  
  return (
    <div>
      <select value={project} onChange={e => setProject(e.target.value)}>
        <option value="new">新建项目</option>
        <option value="project-a">Project A</option>
        <option value="project-b">Project B</option>
      </select>
      
      <textarea
        value={prompt}
        onChange={e => setPrompt(e.target.value)}
        placeholder="描述你的基础设施需求..."
      />
      
      <select value={environment} onChange={e => setEnvironment(e.target.value)}>
        <option value="dev">开发环境</option>
        <option value="staging">预发布环境</option>
        <option value="prod">生产环境</option>
      </select>
      
      <button onClick={handleSubmit}>生成代码</button>
      
      {status && <TaskStatus status={status} />}
    </div>
  );
}

# streamlit_app.py
import streamlit as st
import subprocess
from git import Repo

st.title("IaC Code Generator")

# 项目选择
project = st.selectbox(
    "选择项目",
    ["新建项目", "project-a", "project-b", "project-c"]
)

# 环境选择
environment = st.selectbox(
    "选择环境",
    ["dev", "staging", "prod"]
)

# 需求输入
prompt = st.text_area(
    "描述你的基础设施需求",
    placeholder="创建一个 VPC + 3 层架构 + RDS 数据库"
)

# 提交按钮
if st.button("生成代码"):
    with st.spinner("正在生成..."):
        
        if project == "新建项目":
            # 创建新仓库
            repo_path = create_new_repo("new-project-infra")
        else:
            # 克隆现有仓库
            repo_path = clone_repo(f"{project}-infra")
        
        # 使用 aiac 生成代码
        result = subprocess.run(
            ["aiac", "get", "terraform", prompt],
            capture_output=True,
            text=True
        )
        
        generated_code = result.stdout
        
        # 验证
        subprocess.run(["terraform", "fmt"], cwd=repo_path)
        
        # 提交
        repo = Repo(repo_path)
        repo.index.add(["generated.tf"])
        repo.index.commit(f"feat: {prompt}")
        
        st.success("✅ 代码已生成并提交！")
        
        # 显示代码
        st.code(generated_code, language="hcl")
        
        # 显示 Git 信息
        st.info(f"仓库：{repo_path}")
        st.info(f"分支：{repo.active_branch}")

# docker-compose.yml
version: "3.8"

services:
  # UI 前端
  frontend:
    image: iac-generator-frontend:latest
    ports:
      - "3000:3000"
  
  # 后端 API
  backend:
    image: iac-generator-backend:latest
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - GITHUB_TOKEN=${GITHUB_TOKEN}
      - DATABASE_URL=postgresql://user:pass@db:5432/iac
    depends_on:
      - db
      - vector_db
  
  # PostgreSQL
  db:
    image: postgres:15
    volumes:
      - postgres_data:/var/lib/postgresql/data
  
  # 向量数据库
  vector_db:
    image: chromadb/chroma:latest
    volumes:
      - vector_data:/chroma
  
  # Terraform MCP Server
  mcp_server:
    image: hashicorp/terraform-mcp-server:latest
    environment:
      - PROVIDERS=aws,alicloud,azurerm

volumes:
  postgres_data:
  vector_data:

# 向量化代码仓库
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma

embeddings = OpenAIEmbeddings()
vector_store = Chroma(embedding_function=embeddings)

# 索引所有仓库代码
for repo in ["project-a", "project-b", "project-c"]:
    repo_path = f"/repos/{repo}-infra"
    
    for tf_file in find_terraform_files(repo_path):
        code = read_file(tf_file)
        
        # 生成嵌入向量
        embedding = embeddings.embed_documents([code])
        
        # 存储到向量数据库
        vector_store.add_texts(
            texts=[code],
            metadatas=[{
                "repo": repo,
                "file": tf_file,
                "type": classify_file(tf_file)  # module/variable/resource
            }]
        )

# 检索相关代码
def retrieve_context(prompt: str, project: str):
    results = vector_store.similarity_search(
        query=prompt,
        k=5,
        filter={"repo": project}
    )
    return [doc.page_content for doc in results]

def determine_file_path(prompt: str, repo_structure: dict) -> str:
    """根据项目结构智能决定文件路径"""
    
    resource_type = extract_resource_type(prompt)  # redis, vpc, eks, etc.
    
    # 模式 1：模块化结构
    if "modules/" in repo_structure:
        return f"modules/{resource_type}/main.tf"
    
    # 模式 2：环境分层
    elif "environments/" in repo_structure:
        return f"environments/prod/{resource_type}.tf"
    
    # 模式 3：平铺结构
    else:
        return f"{resource_type}.tf"

def ensure_style_consistency(generated_code: str, existing_code: str):
    """确保生成代码符合项目风格"""
    
    # 1. 提取标签模式
    tag_pattern = extract_tag_pattern(existing_code)
    generated_code = apply_tag_pattern(generated_code, tag_pattern)
    
    # 2. 提取命名约定
    naming_convention = extract_naming_convention(existing_code)
    generated_code = apply_naming_convention(generated_code, naming_convention)
    
    # 3. 提取变量引用
    var_pattern = extract_variable_pattern(existing_code)
    generated_code = apply_variable_pattern(generated_code, var_pattern)
    
    return generated_code

def extract_tag_pattern(code: str):
    """提取标签模式"""
    # 例如：tags = { env = "prod", team = "devops", project = "..." }
    return re.search(r'tags\s*=\s*\{[^}]+\}', code)

from git import Repo
import github

def commit_and_create_pr(
    repo_path: str,
    generated_code: str,
    file_path: str,
    branch_name: str,
    prompt: str
):
    """提交代码并创建 PR"""
    
    # 1. 克隆仓库
    repo = Repo(repo_path)
    
    # 2. 创建分支
    repo.git.checkout("-b", branch_name)
    
    # 3. 写入文件
    full_path = os.path.join(repo_path, file_path)
    os.makedirs(os.path.dirname(full_path), exist_ok=True)
    with open(full_path, "w") as f:
        f.write(generated_code)
    
    # 4. 添加 + 提交
    repo.index.add([file_path])
    repo.index.commit(f"feat: {prompt}")
    
    # 5. 推送
    repo.git.push("origin", branch_name)
    
    # 6. 创建 PR（GitHub API）
    g = github.Github(os.environ["GITHUB_TOKEN"])
    gh_repo = g.get_repo(repo.remote().url.split("github.com/")[1].replace(".git", ""))
    pr = gh_repo.create_pull(
        title=f"Add {extract_resource_type(prompt)}",
        body=f"Generated from natural language:\n\n{prompt}",
        head=branch_name,
        base="main"
    )
    
    return pr.html_url