server.ts

import express, { Request, Response } from 'express'
import cors from 'cors'
import rateLimit from 'express-rate-limit'
import dotenv from 'dotenv'
import fs from 'fs/promises'
import path from 'path'
import OpenAI from 'openai'
import { getMCPClient } from './mcp-client.js'
import { generatePythonCodeFromJSON } from './utils/pythonConvert.js'
import { generateDrawIoXml } from './src/utils/drawioGenerator.js'

dotenv.config()

const app = express()
const PORT = process.env.PORT || 3001
const NVIDIA_MODEL = process.env.NVIDIA_MODEL || 'meta/llama-3.3-70b-instruct'
const CORS_ORIGIN = process.env.CORS_ORIGIN || 'http://localhost:5173'

app.use(cors({ origin: CORS_ORIGIN }))
app.use(express.json())

const apiLimiter = rateLimit({
  windowMs: 60 * 1000,
  max: 20,
  standardHeaders: true,
  legacyHeaders: false,
  message: { error: 'Too many requests, please try again later.' },
})
app.use('/api/', apiLimiter)

app.use('/generated-diagrams', express.static('generated-diagrams'))
app.use('/assets', express.static('generated-diagrams/assets'))

if (process.env.NODE_ENV === 'production') {
  const distPath = path.join(process.cwd(), 'dist')
  app.use(express.static(distPath))
}

interface DiagramRequest {
  prompt: string
}

interface DiagramResponse {
  diagram: string
  format: string
  dot_source?: string
  diagram_json?: unknown
  edit_url?: string
}

const nvidia = new OpenAI({
  apiKey: process.env.NVIDIA_API_KEY || '',
  baseURL: 'https://integrate.api.nvidia.com/v1',
})

let mcpClientReady = false
let mcpClientError: string | null = null

async function initializeMCPClient() {
  try {
    console.log('Initializing MCP client...')
    await getMCPClient()
    mcpClientReady = true
    console.log('MCP client initialized successfully')
  } catch (error) {
    mcpClientError = error instanceof Error ? error.message : 'Unknown error'
    console.error('Failed to initialize MCP client:', mcpClientError)
    console.error('Make sure python3 and the diagrams library are installed.')
  }
}

async function getMCPDiagramTool() {
  const mcpClient = await getMCPClient()
  const tools = await mcpClient.listTools()
  const tool = tools.find(
    t => t.name.includes('diagram') || t.name.includes('create') || t.name.includes('generate')
  )
  if (!tool) {
    throw new Error('Diagram generation tool not available')
  }
  return { mcpClient, tool }
}

function extractDiagramPath(result: { content?: unknown }): string | null {
  if (!result.content || !Array.isArray(result.content)) return null
  for (const item of result.content) {
    if (item.type === 'text' && item.text) {
      const match = (item.text as string).match(/(generated-diagrams\/[\w/\-.]+\.png)/)
      if (match) return match[1].trim()
    }
  }
  return null
}

function extractMCPError(result: { content?: unknown }): string {
  if (!result.content || !Array.isArray(result.content)) return 'Unknown MCP error'
  for (const item of result.content) {
    if (item.type === 'text' && item.text) {
      try {
        const parsed = JSON.parse(item.text as string)
        if (parsed.message) return parsed.message as string
        if (parsed.error) return parsed.error as string
      } catch {
        if ((item.text as string).length < 200) return item.text as string
      }
    }
  }
  return 'Failed to extract diagram path from MCP response'
}

async function resolveSafePath(diagramPath: string): Promise<string> {
  const cwd = process.cwd()
  const fullPath = path.resolve(cwd, diagramPath)
  const allowedPrefix = path.join(cwd, 'generated-diagrams') + path.sep
  if (!fullPath.startsWith(allowedPrefix)) {
    throw new Error('Invalid diagram path')
  }
  return fullPath
}

async function generateDiagramCode(prompt: string): Promise<string> {
  console.log(`Generating diagram JSON with NVIDIA AI (Model: ${NVIDIA_MODEL})...`)

  const systemPrompt = `You are a Universal Cloud Architect.
Convert a natural language description into a STRUCTURED JSON representation of the architecture.

OUTPUT FORMAT:
Return a single JSON object with this schema:
{
  "title": "Diagram Title",
  "nodes": [
    { "id": "web", "label": "Web Server", "type": "aws.ec2", "cluster": "vpc1" },
    { "id": "db", "label": "Database", "type": "aws.rds", "cluster": "vpc1" }
  ],
  "edges": [
    { "from": "web", "to": "db" }
  ],
  "clusters": [
    { "id": "vpc1", "label": "VPC", "type": "aws.vpc" }
  ]
}

NODE TYPES (Use these prefixes):
- AWS: aws.ec2, aws.lambda, aws.s3, aws.rds, aws.dynamodb, aws.sns, aws.sqs, aws.alb, aws.eks, aws.api_gateway, aws.dms, aws.sct, aws.migration, aws.sagemaker
- Azure: azure.vm, azure.functions, azure.blob, azure.sql, azure.aks, azure.ai, azure.migration
- GCP: gcp.compute, gcp.storage, gcp.sql, gcp.gke, gcp.run, gcp.ml, gcp.ai
- On-Prem: onprem.oracle, onprem.sql, onprem.server, onprem.database, onprem.client
- DevOps: devops.docker, devops.kubernetes, devops.github, devops.terraform

RULES:
1. "id" must be unique and alphanumeric (lowercase, no spaces).
2. "cluster" is optional ID of a parent cluster.
3. Return ONLY JSON. No markdown blocks, no code fences.
4. Use specific service types where possible (e.g. 'aws.dms' not 'aws.migration').`

  const completion = await nvidia.chat.completions.create({
    model: NVIDIA_MODEL,
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: prompt },
    ],
    temperature: 0,
    max_tokens: 3000,
  })

  let text = completion.choices[0]?.message?.content || ''
  text = text.replace(/```json/g, '').replace(/```/g, '').trim()
  console.log('AI Generated JSON length:', text.length)
  return text
}

async function generateIaCCode(diagramData: unknown): Promise<string> {
  const systemPrompt = `You are a Senior DevOps Engineer and Terraform Expert.
Generate production-ready Terraform (HCL) code based on the provided Architecture JSON.

RULES:
1. Use AWS Provider (hashicorp/aws) unless Azure/GCP is detected in the node types.
2. Create complete, working logic with provider blocks and matching resources.
   - aws.ec2 -> aws_instance
   - aws.lambda -> aws_lambda_function
   - aws.s3 -> aws_s3_bucket
   - aws.vpc -> aws_vpc (and associated subnets/gateways if implied)
   - aws.rds -> aws_db_instance
3. Use the "edges" to create security group rules or IAM policies (A -> B means A can communicate with B).
4. Return ONLY Terraform HCL. No markdown formatting, no explanations.
5. Use best practices: specific versions, tagging, variable definitions.`

  const completion = await nvidia.chat.completions.create({
    model: NVIDIA_MODEL,
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: `Architecture JSON:\n${JSON.stringify(diagramData, null, 2)}` },
    ],
    temperature: 0,
    max_tokens: 4000,
  })

  let text = completion.choices[0]?.message?.content || ''
  text = text.replace(/```hcl/g, '').replace(/```terraform/g, '').replace(/```/g, '').trim()
  return text
}

app.post('/api/generate-iac', async (req: Request, res: Response) => {
  try {
    const { diagramData } = req.body
    if (!diagramData) {
      return res.status(400).json({ error: 'Diagram data is required' })
    }
    console.log('Generating IaC for diagram...')
    const iacCode = await generateIaCCode(diagramData)
    res.json({ iac: iacCode })
  } catch (error) {
    console.error('IaC generation failed:', error)
    res.status(500).json({ error: 'Failed to generate IaC' })
  }
})

app.get('/api/health', (_req: Request, res: Response) => {
  res.json({
    status: mcpClientReady ? 'ready' : 'initializing',
    mcpConnected: mcpClientReady,
    error: mcpClientError,
  })
})

app.get('/api/docs', (_req: Request, res: Response) => {
  res.json({
    name: 'SkyArch AI API',
    version: '1.2.0',
    description: 'Universal Cloud Architecture Diagram Generator API',
    endpoints: [
      {
        path: '/api/generate-diagram',
        method: 'POST',
        description: 'Generate a diagram from natural language prompt',
        body: { prompt: 'string (max 2000 chars)' },
        response: { diagram: 'string (base64 encoded png)', format: 'string (png)' },
      },
      { path: '/api/health', method: 'GET', description: 'Check server and MCP status' },
    ],
    supported_providers: ['AWS', 'Azure', 'GCP', 'On-Premise', 'DevOps', 'K8s'],
    features: ['Fuzzy Match Icon Resolution', 'Context-Aware Provider Selection', 'Auto-Import Injection', 'Multi-Cloud Support'],
  })
})

app.post('/api/generate-diagram', async (req: Request, res: Response) => {
  try {
    const { prompt } = req.body as DiagramRequest

    if (!prompt || typeof prompt !== 'string') {
      return res.status(400).json({ error: 'Prompt is required' })
    }
    if (prompt.length > 2000) {
      return res.status(400).json({ error: 'Prompt is too long (max 2000 characters)' })
    }
    if (!mcpClientReady) {
      return res.status(503).json({
        error: 'MCP server not ready. Please try again in a moment.',
        details: mcpClientError,
      })
    }

    const aiJsonString = await generateDiagramCode(prompt)
    let diagramJSON
    try {
      diagramJSON = JSON.parse(aiJsonString)
    } catch {
      throw new Error('AI returned invalid JSON')
    }
    console.log('AI Generated JSON:\n', JSON.stringify(diagramJSON, null, 2))

    const pythonCode = generatePythonCodeFromJSON(diagramJSON)
    console.log('Generated Python code for Preview:\n', pythonCode)
    const drawIoXml = generateDrawIoXml(diagramJSON)

    const { mcpClient, tool } = await getMCPDiagramTool()
    const result = await mcpClient.callTool(tool.name, {
      code: pythonCode,
      filename: `aws_diagram_${Date.now()}`,
      workspace_dir: process.cwd(),
    })

    const diagramPath = extractDiagramPath(result)
    if (!diagramPath) {
      console.error('No diagram path found in MCP response:', JSON.stringify(result, null, 2))
      return res.status(500).json({ error: extractMCPError(result), result })
    }

    let fullPath: string
    try {
      fullPath = await resolveSafePath(diagramPath)
    } catch {
      return res.status(500).json({ error: 'Invalid diagram path' })
    }

    let diagram: string
    try {
      const imageBuffer = await fs.readFile(fullPath)
      diagram = `data:image/png;base64,${imageBuffer.toString('base64')}`
    } catch (error) {
      console.error('Failed to read diagram file:', error)
      return res.status(500).json({ error: 'Failed to read generated diagram file', path: diagramPath })
    }

    res.json({
      diagram,
      format: 'png',
      dot_source: drawIoXml,
      diagram_json: diagramJSON,
      edit_url: diagramPath.replace('.png', '.json'),
    } as DiagramResponse)

  } catch (error) {
    console.error('Error generating diagram:', error)
    res.status(500).json({
      error: error instanceof Error ? error.message : 'Failed to generate diagram',
    })
  }
})

app.post('/api/update-diagram', async (req: Request, res: Response) => {
  try {
    const { diagram_json } = req.body
    if (!diagram_json) {
      return res.status(400).json({ error: 'Diagram JSON is required' })
    }
    if (!mcpClientReady) {
      return res.status(503).json({ error: 'MCP server not ready' })
    }

    const pythonCode = generatePythonCodeFromJSON(diagram_json)
    console.log('Regenerated Python Code:\n', pythonCode)

    const { mcpClient, tool } = await getMCPDiagramTool()
    const result = await mcpClient.callTool(tool.name, {
      code: pythonCode,
      filename: `aws_diagram_update_${Date.now()}`,
      workspace_dir: process.cwd(),
    })

    const diagramPath = extractDiagramPath(result)
    if (!diagramPath) {
      return res.status(500).json({ error: 'Failed to generate image', details: result })
    }

    let fullPath: string
    try {
      fullPath = await resolveSafePath(diagramPath)
    } catch {
      return res.status(500).json({ error: 'Invalid diagram path' })
    }

    const imageBuffer = await fs.readFile(fullPath)
    const diagram = `data:image/png;base64,${imageBuffer.toString('base64')}`

    res.json({ diagram, diagram_path: diagramPath })

  } catch (error) {
    console.error('Update failed:', error)
    res.status(500).json({ error: 'Failed to update diagram' })
  }
})

if (process.env.NODE_ENV === 'production') {
  app.get('*', (_req: Request, res: Response) => {
    res.sendFile(path.join(process.cwd(), 'dist', 'index.html'))
  })
}

initializeMCPClient().then(() => {
  app.listen(PORT, () => {
    console.log(`Server running on http://localhost:${PORT}`)
    console.log(`NVIDIA AI Model: ${NVIDIA_MODEL}`)
  })
})